From 9381c25f82167b45256b51752ed483b6e5f3547b Mon Sep 17 00:00:00 2001
From: alwaysintreble <mmmcheese158@gmail.com>
Date: Mon, 29 Jul 2024 22:14:23 -0500
Subject: [PATCH] Core: add a hook for worlds to modify early_locations

---
 Fill.py             | 37 ++++++++++++++++++++++---------------
 worlds/AutoWorld.py | 14 ++++++++++++++
 2 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/Fill.py b/Fill.py
index 5185bbb60ee4..ae23baa62a01 100644
--- a/Fill.py
+++ b/Fill.py
@@ -358,9 +358,9 @@ def distribute_early_items(multiworld: MultiWorld,
         early_priority_locations: typing.List[Location] = []
         loc_indexes_to_remove: typing.Set[int] = set()
         base_state = multiworld.state.copy()
-        base_state.sweep_for_events(locations=(loc for loc in multiworld.get_filled_locations() if loc.address is None))
+        per_player_early_locations = get_early_locations(multiworld)
         for i, loc in enumerate(fill_locations):
-            if loc.can_reach(base_state):
+            if loc in per_player_early_locations[loc.player]:
                 if loc.progress_type == LocationProgressType.PRIORITY:
                     early_priority_locations.append(loc)
                 else:
@@ -821,17 +821,6 @@ def failed(warning: str, force: typing.Union[bool, str]) -> None:
         else:
             warn(warning, force)
 
-    swept_state = multiworld.state.copy()
-    swept_state.sweep_for_events()
-    reachable = frozenset(multiworld.get_reachable_locations(swept_state))
-    early_locations: typing.Dict[int, typing.List[str]] = collections.defaultdict(list)
-    non_early_locations: typing.Dict[int, typing.List[str]] = collections.defaultdict(list)
-    for loc in multiworld.get_unfilled_locations():
-        if loc in reachable:
-            early_locations[loc.player].append(loc.name)
-        else:  # not reachable with swept state
-            non_early_locations[loc.player].append(loc.name)
-
     world_name_lookup = multiworld.world_name_lookup
 
     block_value = typing.Union[typing.List[str], typing.Dict[str, typing.Any], str]
@@ -911,13 +900,22 @@ def failed(warning: str, force: typing.Union[bool, str]) -> None:
                 locations = block['locations']
             if isinstance(locations, str):
                 locations = [locations]
-
-            if isinstance(locations, dict):
+            elif isinstance(locations, dict):
                 location_list = []
                 for key, value in locations.items():
                     location_list += [key] * value
                 locations = location_list
 
+            if "early_locations" in locations or "non_early_locations" in locations:
+                per_player_early_locations = get_early_locations(multiworld)
+                early_locations: typing.Dict[int, typing.List[str]] = collections.defaultdict(list)
+                non_early_locations: typing.Dict[int, typing.List[str]] = collections.defaultdict(list)
+                for loc in multiworld.get_unfilled_locations():
+                    if loc in per_player_early_locations[loc.player]:
+                        early_locations[loc.player].append(loc.name)
+                    else:  # not reachable with swept state
+                        non_early_locations[loc.player].append(loc.name)
+
             if "early_locations" in locations:
                 locations.remove("early_locations")
                 for target_player in worlds:
@@ -1015,3 +1013,12 @@ def failed(warning: str, force: typing.Union[bool, str]) -> None:
         except Exception as e:
             raise Exception(
                 f"Error running plando for player {player} ({multiworld.player_name[player]})") from e
+
+
+def get_early_locations(multiworld: MultiWorld) -> typing.Dict[int, typing.List[Location]]:
+    early_locations = collections.defaultdict(list)
+    base_state = multiworld.state.copy()
+    for location in multiworld.get_reachable_locations(base_state):
+        early_locations[location.player].append(location)
+    call_all(multiworld, "modify_early_locations", early_locations)
+    return early_locations
diff --git a/worlds/AutoWorld.py b/worlds/AutoWorld.py
index af067e5cb8a6..6f01219e5198 100644
--- a/worlds/AutoWorld.py
+++ b/worlds/AutoWorld.py
@@ -378,6 +378,19 @@ def pre_fill(self) -> None:
         """Optional method that is supposed to be used for special fill stages. This is run *after* plando."""
         pass
 
+    def modify_early_locations(self, early_locations: Dict[int, List["Location"]]) -> None:
+        """
+        Gets called as part of distribute_early_items and early items plando.
+        Can be used to modify which locations are considered early.
+        """
+        # make a clean copy of state with just start inventory
+        sweep_state = self.multiworld.state.copy()
+        # collect any events from our already reachable locations
+        for location in early_locations[self.player]:
+            if location.is_event:
+                sweep_state.collect(location.item, location=location)
+        early_locations[self.player] += self.multiworld.get_reachable_locations(sweep_state, self.player)
+
     def fill_hook(self,
                   progitempool: List["Item"],
                   usefulitempool: List["Item"],
@@ -391,6 +404,7 @@ def post_fill(self) -> None:
         Optional Method that is called after regular fill. Can be used to do adjustments before output generation.
         This happens before progression balancing, so the items may not be in their final locations yet.
         """
+        pass
 
     def generate_output(self, output_directory: str) -> None:
         """