Speed up Progression Balancing, mostly by using generators and pre-sorts where the opportunity exists

In some cases multi-thousand element lists were created in-memory with near identical contents, per player, then discarded and reassembled.
Was testing against a case with 3 GB of additional memory use (50 players) which appeared to get stuck, but really was just very slow. This example case is fixed with these changes.
Additionally, progression balancing is now run after ShopSlotFill, so it is now "aware" of the changed progression shops can produce.
As well, special handling for keys was removed, as not all games will have the notion of keys.

BaseClasses.py

@@ -606,9 +606,9 @@ class CollectionState(object):
         new_locations = True
         while new_locations:
             reachable_events = {location for location in locations if location.event and
-                                (not key_only or (not self.world.keyshuffle[
+                                (not key_only or
-                                    location.item.player] and location.item.smallkey) or (not self.world.bigkeyshuffle[
+                                 (not self.world.keyshuffle[location.item.player] and location.item.smallkey)
-                                    location.item.player] and location.item.bigkey))
+                                 or (not self.world.bigkeyshuffle[location.item.player] and location.item.bigkey))
                                 and location.can_reach(self)}
             new_locations = reachable_events - self.events
             for event in new_locations:

Fill.py

@@ -1,5 +1,7 @@
 import logging
 import typing
+import collections
+import itertools
 
 from BaseClasses import CollectionState, PlandoItem, Location
 from Items import ItemFactory
@@ -243,12 +245,7 @@ def balance_multiworld_progression(world):
         unchecked_locations = world.get_locations().copy()
         world.random.shuffle(unchecked_locations)
 
-        reachable_locations_count = {player: 0 for player in range(1, world.players + 1)}
+        reachable_locations_count = {player: 0 for player in world.player_ids}
-
-        def event_key(location):
-            return location.event and (
-                    world.keyshuffle[location.item.player] or not location.item.smallkey) and (
-                    world.bigkeyshuffle[location.item.player] or not location.item.bigkey)
 
         def get_sphere_locations(sphere_state, locations):
             sphere_state.sweep_for_events(key_only=True, locations=locations)
@@ -269,33 +266,38 @@ def balance_multiworld_progression(world):
                     balancing_unchecked_locations = unchecked_locations.copy()
                     balancing_reachables = reachable_locations_count.copy()
                     balancing_sphere = sphere_locations.copy()
-                    candidate_items = []
+                    candidate_items = collections.defaultdict(list)
                     while True:
                         for location in balancing_sphere:
-                            if event_key(location):
+                            if location.event:
                                 balancing_state.collect(location.item, True, location)
-                                if location.item.player in balancing_players and not location.locked:
+                                player = location.item.player
-                                    candidate_items.append(location)
+                                # only replace items that end up in another player's world
+                                if not location.locked and player in balancing_players and location.player != player:
+                                    candidate_items[player].append(location)
                         balancing_sphere = get_sphere_locations(balancing_state, balancing_unchecked_locations)
                         for location in balancing_sphere:
                             balancing_unchecked_locations.remove(location)
                             balancing_reachables[location.player] += 1
                         if world.has_beaten_game(balancing_state) or all(
-                                [reachables >= threshold for reachables in balancing_reachables.values()]):
+                                reachables >= threshold for reachables in balancing_reachables.values()):
                             break
                         elif not balancing_sphere:
                             raise RuntimeError('Not all required items reachable. Something went terribly wrong here.')
-
+                    unlocked_locations = collections.defaultdict(list)
-                    unlocked_locations = [l for l in unchecked_locations if l not in balancing_unchecked_locations]
+                    for l in unchecked_locations:
+                        if l not in balancing_unchecked_locations:
+                            unlocked_locations[l.player].append(l)
                     items_to_replace = []
                     for player in balancing_players:
-                        locations_to_test = [l for l in unlocked_locations if l.player == player]
+                        locations_to_test = unlocked_locations[player]
-                        # only replace items that end up in another player's world
+                        items_to_test = candidate_items[player]
-                        items_to_test = [l for l in candidate_items if l.item.player == player and l.player != player]
                         while items_to_test:
                             testing = items_to_test.pop()
                             reducing_state = state.copy()
-                            for location in [*[l for l in items_to_replace if l.item.player == player], *items_to_test]:
+                            for location in itertools.chain((l for l in items_to_replace if l.item.player == player),
+                                                            items_to_test):
+
                                 reducing_state.collect(location.item, True, location)
 
                             reducing_state.sweep_for_events(locations=locations_to_test)
@@ -320,21 +322,20 @@ def balance_multiworld_progression(world):
                             new_location = replacement_locations.pop()
 
                         swap_location_item(old_location, new_location)
-
-                        new_location.event, old_location.event = True, False
                         logging.debug(f"Progression balancing moved {new_location.item} to {new_location}, "
-                                      f"displacing {old_location.item} in {old_location}")
+                                      f"displacing {old_location.item} into {old_location}")
                         state.collect(new_location.item, True, new_location)
                         replaced_items = True
+
                     if replaced_items:
-                        for location in get_sphere_locations(state, [l for l in unlocked_locations if
+                        unlocked = [fresh for player in balancing_players for fresh in unlocked_locations[player]]
-                                                                     l.player in balancing_players]):
+                        for location in get_sphere_locations(state, unlocked):
                             unchecked_locations.remove(location)
                             reachable_locations_count[location.player] += 1
                             sphere_locations.append(location)
 
             for location in sphere_locations:
-                if event_key(location):
+                if location.event:
                     state.collect(location.item, True, location)
             checked_locations.extend(sphere_locations)
 
@@ -345,7 +346,7 @@ def balance_multiworld_progression(world):
 
 
 def swap_location_item(location_1: Location, location_2: Location, check_locked=True):
-    """Swaps Items of locations. Does NOT swap flags like event, shop_slot or locked"""
+    """Swaps Items of locations. Does NOT swap flags like shop_slot or locked, but does swap event"""
     if check_locked:
         if location_1.locked:
             logging.warning(f"Swapping {location_1}, which is marked as locked.")
@@ -354,6 +355,7 @@ def swap_location_item(location_1: Location, location_2: Location, check_locked=
     location_2.item, location_1.item = location_1.item, location_2.item
     location_1.item.location = location_1
     location_2.item.location = location_2
+    location_1.event, location_2.event = location_2.event, location_1.event
 
 
 def distribute_planned(world):

Main.py

@@ -216,13 +216,13 @@ def main(args, seed=None):
     elif args.algorithm == 'balanced':
         distribute_items_restrictive(world, True)
 
-    if world.players > 1:
-        balance_multiworld_progression(world)
-
     logger.info("Filling Shop Slots")
 
     ShopSlotFill(world)
 
+    if world.players > 1:
+        balance_multiworld_progression(world)
+
     logger.info('Patching ROM.')
 
 

Mystery.py

@@ -199,10 +199,10 @@ def main(args=None, callback=ERmain):
         for option, player_settings in vars(erargs).items():
             if type(player_settings) == dict:
                 if all(type(value) != list for value in player_settings.values()):
-                    if len(frozenset(player_settings.values())) > 1:
+                    if len(player_settings.values()) > 1:
                         important[option] = {player: value for player, value in player_settings.items() if
                                              player <= args.yaml_output}
-                    elif len(frozenset(player_settings.values())) > 0:
+                    elif len(player_settings.values()) > 0:
                         important[option] = player_settings[1]
                     else:
                         logging.debug(f"No player settings defined for option '{option}'")