Zillion: fix logic cache (#3719)

worlds/zillion/__init__.py

@@ -4,7 +4,7 @@ import functools
 import settings
 import threading
 import typing
-from typing import Any, Dict, List, Set, Tuple, Optional
+from typing import Any, Dict, List, Set, Tuple, Optional, Union
 import os
 import logging
 
@@ -12,7 +12,7 @@ from BaseClasses import ItemClassification, LocationProgressType, \
     MultiWorld, Item, CollectionState, Entrance, Tutorial
 
 from .gen_data import GenData
-from .logic import cs_to_zz_locs
+from .logic import ZillionLogicCache
 from .region import ZillionLocation, ZillionRegion
 from .options import ZillionOptions, validate, z_option_groups
 from .id_maps import ZillionSlotInfo, get_slot_info, item_name_to_id as _item_name_to_id, \
@@ -21,7 +21,6 @@ from .id_maps import ZillionSlotInfo, get_slot_info, item_name_to_id as _item_na
 from .item import ZillionItem
 from .patch import ZillionPatch
 
-from zilliandomizer.randomizer import Randomizer as ZzRandomizer
 from zilliandomizer.system import System
 from zilliandomizer.logic_components.items import RESCUE, items as zz_items, Item as ZzItem
 from zilliandomizer.logic_components.locations import Location as ZzLocation, Req
@@ -121,6 +120,7 @@ class ZillionWorld(World):
     """ This is kind of a cache to avoid iterating through all the multiworld locations in logic. """
     slot_data_ready: threading.Event
     """ This event is set in `generate_output` when the data is ready for `fill_slot_data` """
+    logic_cache: Union[ZillionLogicCache, None] = None
 
     def __init__(self, world: MultiWorld, player: int):
         super().__init__(world, player)
@@ -134,9 +134,6 @@ class ZillionWorld(World):
         self.id_to_zz_item = id_to_zz_item
 
     def generate_early(self) -> None:
-        if not hasattr(self.multiworld, "zillion_logic_cache"):
-            setattr(self.multiworld, "zillion_logic_cache", {})
-
         zz_op, item_counts = validate(self.options)
 
         if zz_op.early_scope:
@@ -163,6 +160,8 @@ class ZillionWorld(World):
         assert self.zz_system.randomizer, "generate_early hasn't been called"
         assert self.id_to_zz_item, "generate_early hasn't been called"
         p = self.player
+        logic_cache = ZillionLogicCache(p, self.zz_system.randomizer, self.id_to_zz_item)
+        self.logic_cache = logic_cache
         w = self.multiworld
         self.my_locations = []
 
@@ -201,15 +200,12 @@ class ZillionWorld(World):
                 if not zz_loc.item:
 
                     def access_rule_wrapped(zz_loc_local: ZzLocation,
-                                            p: int,
+                                            lc: ZillionLogicCache,
-                                            zz_r: ZzRandomizer,
-                                            id_to_zz_item: Dict[int, ZzItem],
                                             cs: CollectionState) -> bool:
-                        accessible = cs_to_zz_locs(cs, p, zz_r, id_to_zz_item)
+                        accessible = lc.cs_to_zz_locs(cs)
                         return zz_loc_local in accessible
 
-                    access_rule = functools.partial(access_rule_wrapped,
+                    access_rule = functools.partial(access_rule_wrapped, zz_loc, logic_cache)
-                                                    zz_loc, self.player, self.zz_system.randomizer, self.id_to_zz_item)
 
                     loc_name = self.zz_system.randomizer.loc_name_2_pretty[zz_loc.name]
                     loc = ZillionLocation(zz_loc, self.player, loc_name, here)
@@ -402,13 +398,6 @@ class ZillionWorld(World):
         game = self.zz_system.get_game()
         return get_slot_info(game.regions, game.char_order[0], game.loc_name_2_pretty)
 
-    # def modify_multidata(self, multidata: Dict[str, Any]) -> None:
-    #     """For deeper modification of server multidata."""
-    #     # not modifying multidata, just want to call this at the end of the generation process
-    #     cache = getattr(self.multiworld, "zillion_logic_cache")
-    #     import sys
-    #     print(sys.getsizeof(cache))
-
     # end of ordered Main.py calls
 
     def create_item(self, name: str) -> Item:

worlds/zillion/logic.py

@@ -1,4 +1,4 @@
-from typing import Dict, FrozenSet, Tuple, List, Counter as _Counter
+from typing import Dict, FrozenSet, Mapping, Tuple, List, Counter as _Counter
 
 from BaseClasses import CollectionState
 
@@ -44,38 +44,51 @@ def item_counts(cs: CollectionState, p: int) -> Tuple[Tuple[str, int], ...]:
     return tuple((item_name, cs.count(item_name, p)) for item_name in item_name_to_id)
 
 
-LogicCacheType = Dict[int, Tuple[Dict[int, _Counter[str]], FrozenSet[Location]]]
+_cache_miss: Tuple[None, FrozenSet[Location]] = (None, frozenset())
-""" { hash: (cs.prog_items, accessible_locations) } """
 
 
-def cs_to_zz_locs(cs: CollectionState, p: int, zz_r: Randomizer, id_to_zz_item: Dict[int, Item]) -> FrozenSet[Location]:
+class ZillionLogicCache:
-    """
+    _cache: Dict[int, Tuple[_Counter[str], FrozenSet[Location]]]
-    given an Archipelago `CollectionState`,
+    """ `{ hash: (counter_from_prog_items, accessible_zz_locations) }` """
-    returns frozenset of accessible zilliandomizer locations
+    _player: int
-    """
+    _zz_r: Randomizer
-    # caching this function because it would be slow
+    _id_to_zz_item: Mapping[int, Item]
-    logic_cache: LogicCacheType = getattr(cs.multiworld, "zillion_logic_cache", {})
-    _hash = set_randomizer_locs(cs, p, zz_r)
-    counts = item_counts(cs, p)
-    _hash += hash(counts)
 
-    if _hash in logic_cache and logic_cache[_hash][0] == cs.prog_items:
+    def __init__(self, player: int, zz_r: Randomizer, id_to_zz_item: Mapping[int, Item]) -> None:
-        # print("cache hit")
+        self._cache = {}
-        return logic_cache[_hash][1]
+        self._player = player
+        self._zz_r = zz_r
+        self._id_to_zz_item = id_to_zz_item
 
-    # print("cache miss")
+    def cs_to_zz_locs(self, cs: CollectionState) -> FrozenSet[Location]:
-    have_items: List[Item] = []
+        """
-    for name, count in counts:
+        given an Archipelago `CollectionState`,
-        have_items.extend([id_to_zz_item[item_name_to_id[name]]] * count)
+        returns frozenset of accessible zilliandomizer locations
-    # have_req is the result of converting AP CollectionState to zilliandomizer collection state
+        """
-    have_req = zz_r.make_ability(have_items)
+        # caching this function because it would be slow
+        _hash = set_randomizer_locs(cs, self._player, self._zz_r)
+        counts = item_counts(cs, self._player)
+        _hash += hash(counts)
 
-    # This `get_locations` is where the core of the logic comes in.
+        cntr, locs = self._cache.get(_hash, _cache_miss)
-    # It takes a zilliandomizer collection state (a set of the abilities that I have)
+        if cntr == cs.prog_items[self._player]:
-    # and returns list of all the zilliandomizer locations I can access with those abilities.
+            # print("cache hit")
-    tr = frozenset(zz_r.get_locations(have_req))
+            return locs
 
-    # save result in cache
+        # print("cache miss")
-    logic_cache[_hash] = (cs.prog_items.copy(), tr)
+        have_items: List[Item] = []
+        for name, count in counts:
+            have_items.extend([self._id_to_zz_item[item_name_to_id[name]]] * count)
+        # have_req is the result of converting AP CollectionState to zilliandomizer collection state
+        have_req = self._zz_r.make_ability(have_items)
+        # print(f"{have_req=}")
 
-    return tr
+        # This `get_locations` is where the core of the logic comes in.
+        # It takes a zilliandomizer collection state (a set of the abilities that I have)
+        # and returns list of all the zilliandomizer locations I can access with those abilities.
+        tr = frozenset(self._zz_r.get_locations(have_req))
+
+        # save result in cache
+        self._cache[_hash] = (cs.prog_items[self._player].copy(), tr)
+
+        return tr