HK: Abstract and default grub counts (#4336)

worlds/hk/__init__.py

@@ -181,6 +181,7 @@ class HKWorld(World):
     charm_costs: typing.List[int]
     cached_filler_items = {}
     grub_count: int
+    grub_player_count: typing.Dict[int, int]
 
     def __init__(self, multiworld, player):
         super(HKWorld, self).__init__(multiworld, player)
@@ -190,7 +191,6 @@ class HKWorld(World):
         self.ranges = {}
         self.created_shop_items = 0
         self.vanilla_shop_costs = deepcopy(vanilla_shop_costs)
-        self.grub_count = 0
 
     def generate_early(self):
         options = self.options
@@ -206,6 +206,13 @@ class HKWorld(World):
         self.multiworld.push_precollected(HKItem(starts[options.StartLocation.current_key],
                                           True, None, "Event", self.player))
 
+        # defaulting so completion condition isn't incorrect before pre_fill
+        self.grub_count = (
+            46 if options.GrubHuntGoal == GrubHuntGoal.special_range_names["all"]
+            else options.GrubHuntGoal
+            )
+        self.grub_player_count = {self.player: self.grub_count}
+
     def white_palace_exclusions(self):
         exclusions = set()
         wp = self.options.WhitePalace
@@ -469,25 +476,20 @@ class HKWorld(World):
         elif goal == Goal.option_godhome_flower:
             multiworld.completion_condition[player] = lambda state: state.count("Godhome_Flower_Quest", player)
         elif goal == Goal.option_grub_hunt:
-            pass  # will set in stage_pre_fill()
+            multiworld.completion_condition[player] = lambda state: self.can_grub_goal(state)
         else:
             # Any goal
             multiworld.completion_condition[player] = lambda state: _hk_siblings_ending(state, player) and \
-                _hk_can_beat_radiance(state, player) and state.count("Godhome_Flower_Quest", player)
+                _hk_can_beat_radiance(state, player) and state.count("Godhome_Flower_Quest", player) and \
+                self.can_grub_goal(state)
 
         set_rules(self)
 
+    def can_grub_goal(self, state: CollectionState) -> bool:
+        return all(state.has("Grub", owner, count) for owner, count in self.grub_player_count.items())
+
     @classmethod
     def stage_pre_fill(cls, multiworld: "MultiWorld"):
-        def set_goal(player, grub_rule: typing.Callable[[CollectionState], bool]):
-            world = multiworld.worlds[player]
-
-            if world.options.Goal == "grub_hunt":
-                multiworld.completion_condition[player] = grub_rule
-            else:
-                old_rule = multiworld.completion_condition[player]
-                multiworld.completion_condition[player] = lambda state: old_rule(state) and grub_rule(state)
-
         worlds = [world for world in multiworld.get_game_worlds(cls.game) if world.options.Goal in ["any", "grub_hunt"]]
         if worlds:
             grubs = [item for item in multiworld.get_items() if item.name == "Grub"]
@@ -525,13 +527,13 @@ class HKWorld(World):
 
             for player, grub_player_count in per_player_grubs_per_player.items():
                 if player in all_grub_players:
-                    set_goal(player, lambda state, g=grub_player_count: all(state.has("Grub", owner, count) for owner, count in g.items()))
+                    multiworld.worlds[player].grub_player_count = grub_player_count
 
         for world in worlds:
             if world.player not in all_grub_players:
                 world.grub_count = world.options.GrubHuntGoal.value
                 player = world.player
-                set_goal(player, lambda state, p=player, c=world.grub_count: state.has("Grub", p, c))
+                world.grub_player_count = {player: world.grub_count}
 
     def fill_slot_data(self):
         slot_data = {}