From c35388b1b60867048c2b3f65cf56f5930b06149e Mon Sep 17 00:00:00 2001
From: Zannick <jokeserver@gmail.com>
Date: Sun, 13 Oct 2024 19:33:22 -0700
Subject: [PATCH] pokemon_rb: Cache which mons can learn HMs

This significantly improves the performance of `can_learn_hm`, which previously checked for the presence of each individual pokemon.
---
 worlds/pokemon_rb/__init__.py |  4 +++-
 worlds/pokemon_rb/logic.py    |  7 ++-----
 worlds/pokemon_rb/pokemon.py  | 16 ++++++++++++++--
 3 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/worlds/pokemon_rb/__init__.py b/worlds/pokemon_rb/__init__.py
index 2065507e0d59..0558ee199cea 100644
--- a/worlds/pokemon_rb/__init__.py
+++ b/worlds/pokemon_rb/__init__.py
@@ -19,7 +19,7 @@
 from .rom_addresses import rom_addresses
 from .text import encode_text
 from .rom import generate_output, get_base_rom_bytes, get_base_rom_path, RedDeltaPatch, BlueDeltaPatch
-from .pokemon import process_pokemon_data, process_move_data, verify_hm_moves
+from .pokemon import process_pokemon_data, process_move_data, rebuild_hm_mons_cache, verify_hm_moves
 from .encounters import process_pokemon_locations, process_trainer_data
 from .rules import set_rules
 from .level_scaling import level_scaling
@@ -101,6 +101,7 @@ def __init__(self, multiworld: MultiWorld, player: int):
         self.type_chart = None
         self.local_poke_data = None
         self.local_move_data = None
+        self.local_hm_mons_cache = None
         self.local_tms = None
         self.learnsets = None
         self.trainer_name = None
@@ -261,6 +262,7 @@ def encode_name(name, t):
 
         process_move_data(self)
         process_pokemon_data(self)
+        rebuild_hm_mons_cache(self)
 
         self.dexsanity_table = [
             *(True for _ in range(round(self.options.dexsanity.value))),
diff --git a/worlds/pokemon_rb/logic.py b/worlds/pokemon_rb/logic.py
index 03e3fa3dfad0..a486c830ac3c 100644
--- a/worlds/pokemon_rb/logic.py
+++ b/worlds/pokemon_rb/logic.py
@@ -33,11 +33,8 @@ def can_flash(state, world, player):
 
 
 def can_learn_hm(state, world, move, player):
-    for pokemon, data in world.local_poke_data.items():
-        if state.has(pokemon, player) and data["tms"][6] & 1 << (["Cut", "Fly", "Surf", "Strength",
-                                                                  "Flash"].index(move) + 2):
-            return True
-    return False
+    hm_mons = world.local_hm_mons_cache[move]
+    return hm_mons and state.has_any(hm_mons, player)
 
 
 def can_get_hidden_items(state, world, player):
diff --git a/worlds/pokemon_rb/pokemon.py b/worlds/pokemon_rb/pokemon.py
index 32c0e36869da..3638544de028 100644
--- a/worlds/pokemon_rb/pokemon.py
+++ b/worlds/pokemon_rb/pokemon.py
@@ -345,10 +345,21 @@ def roll_tm_compat(roll_move):
     self.learnsets = learnsets
 
 
+def rebuild_hm_mons_cache(world):
+    # Build a cache of which Pokémon can learn certain HMs
+    hm_mons_cache = {hm: [] for hm in poke_data.hm_moves}
+    for i, hm in enumerate(poke_data.hm_moves):
+        flag = i + 2  # same as can_learn_hm
+        for mon, data in world.local_poke_data.items():
+            if data["tms"][6] & flag:
+                hm_mons_cache[hm].append(mon)
+
+    world.local_hm_mons_cache = hm_mons_cache
+
+
 def verify_hm_moves(multiworld, world, player):
     def intervene(move, test_state):
-        move_bit = pow(2, poke_data.hm_moves.index(move) + 2)
-        viable_mons = [mon for mon in world.local_poke_data if world.local_poke_data[mon]["tms"][6] & move_bit]
+        viable_mons = world.local_hm_mons_cache[move]
         if world.options.randomize_wild_pokemon and viable_mons:
             accessible_slots = [loc for loc in multiworld.get_reachable_locations(test_state, player) if
                                 loc.type == "Wild Encounter"]
@@ -427,6 +438,7 @@ def number_of_zones(mon):
             if intervene_move == last_intervene:
                 raise Exception(f"Caught in infinite loop attempting to ensure {intervene_move} is available to player {player}")
             intervene(intervene_move, test_state)
+            rebuild_hm_mons_cache(world)
             last_intervene = intervene_move
         else:
             break