perf: use raw unit type IDs for unit attributes

sc2/bot_ai_internal.py

@@ -433,7 +433,9 @@ def _abilities_count_and_build_progress(self) -> tuple[Counter[AbilityId], dict[
         for unit in self.units + self.structures:
             for order in unit.orders:
                 abilities_amount[order.ability.exact_id] += 1
-            if not unit.is_ready and (self.race != Race.Terran or not unit.is_structure):
+            if not unit.is_ready and (
+                self.race != Race.Terran or unit._proto.unit_type not in self.game_data.structure_type_ids
+            ):
                 # If an SCV is constructing a building, already_pending would count this structure twice
                 # (once from the SCV order, and once from "not structure.is_ready")
                 if unit.type_id in CREATION_ABILITY_FIX:
@@ -722,7 +724,7 @@ def _prepare_units(self: BotAI) -> None:
         self.reactor_tags: set[int] = set()
 
         worker_types: set[UnitTypeId] = {UnitTypeId.DRONE, UnitTypeId.DRONEBURROWED, UnitTypeId.SCV, UnitTypeId.PROBE}
-        structure_types: frozenset[int] = self.game_data.structure_types
+        structure_type_ids: frozenset[int] = self.game_data.structure_type_ids
 
         index: int = 0
         for unit in self.state.observation_raw.units:
@@ -761,7 +763,7 @@ def _prepare_units(self: BotAI) -> None:
                 elif alliance == 1:
                     self.all_own_units.append(unit_obj)
                     unit_id: UnitTypeId = unit_obj.type_id
-                    if unit_type in structure_types:
+                    if unit_type in structure_type_ids:
                         self.structures.append(unit_obj)
                         if unit_id in race_townhalls[self.race]:
                             self.townhalls.append(unit_obj)
@@ -791,7 +793,7 @@ def _prepare_units(self: BotAI) -> None:
                 # Alliance.Enemy.value = 4
                 elif alliance == 4:
                     self.all_enemy_units.append(unit_obj)
-                    if unit_type in structure_types:
+                    if unit_type in structure_type_ids:
                         self.enemy_structures.append(unit_obj)
                     else:
                         self.enemy_units.append(unit_obj)

sc2/game_data.py

@@ -31,10 +31,45 @@ def __init__(self, data: sc2api_pb2.ResponseData) -> None:
         }
         self.units: dict[int, UnitTypeData] = {u.unit_id: UnitTypeData(self, u) for u in data.units if u.available}
         self.upgrades: dict[int, UpgradeData] = {u.upgrade_id: UpgradeData(self, u) for u in data.upgrades}
-        # Precomputed structure type IDs (raw ints) for O(1) membership checks
-        self.structure_types: frozenset[int] = frozenset(
-            u.unit_id for u in data.units if u.available and Attribute.Structure.value in u.attributes
-        )
+        # Precompute unit type IDs (raw ints) per attribute for instant O(1) checks
+        armored = set()
+        biological = set()
+        light = set()
+        massive = set()
+        mechanical = set()
+        psionic = set()
+        structures = set()
+        unit_attributes: dict[int, frozenset[int]] = {}
+
+        for u in data.units:
+            if not u.available:
+                continue
+            uid = u.unit_id
+            attrs = u.attributes
+            unit_attributes[uid] = frozenset(attrs)
+            if Attribute.Armored.value in attrs:
+                armored.add(uid)
+            if Attribute.Biological.value in attrs:
+                biological.add(uid)
+            if Attribute.Light.value in attrs:
+                light.add(uid)
+            if Attribute.Massive.value in attrs:
+                massive.add(uid)
+            if Attribute.Mechanical.value in attrs:
+                mechanical.add(uid)
+            if Attribute.Psionic.value in attrs:
+                psionic.add(uid)
+            if Attribute.Structure.value in attrs:
+                structures.add(uid)
+
+        self.armored_type_ids: frozenset[int] = frozenset(armored)
+        self.biological_type_ids: frozenset[int] = frozenset(biological)
+        self.light_type_ids: frozenset[int] = frozenset(light)
+        self.massive_type_ids: frozenset[int] = frozenset(massive)
+        self.mechanical_type_ids: frozenset[int] = frozenset(mechanical)
+        self.psionic_type_ids: frozenset[int] = frozenset(psionic)
+        self.structure_type_ids: frozenset[int] = frozenset(structures)
+        self.unit_attributes: dict[int, frozenset[int]] = unit_attributes
         # Cached UnitTypeIds so that conversion does not take long. This needs to be moved elsewhere if a new GameData object is created multiple times per game
 
     @lru_cache(maxsize=256)

sc2/unit.py

@@ -11,9 +11,7 @@
 from sc2.constants import (
     CAN_BE_ATTACKED,
     DAMAGE_BONUS_PER_UPGRADE,
-    IS_ARMORED,
     IS_ATTACKING,
-    IS_BIOLOGICAL,
     IS_CARRYING_MINERALS,
     IS_CARRYING_RESOURCES,
     IS_CARRYING_VESPENE,
@@ -24,12 +22,9 @@
     IS_ENEMY,
     IS_GATHERING,
     IS_LIGHT,
-    IS_MASSIVE,
-    IS_MECHANICAL,
     IS_MINE,
     IS_PATROLLING,
     IS_PLACEHOLDER,
-    IS_PSIONIC,
     IS_REPAIRING,
     IS_RETURNING,
     IS_REVEALED,
@@ -161,37 +156,37 @@ def tag(self) -> int:
     @property
     def is_structure(self) -> bool:
         """Checks if the unit is a structure."""
-        return self._proto.unit_type in self._bot_object.game_data.structure_types
+        return self._proto.unit_type in self._bot_object.game_data.structure_type_ids
 
     @property
     def is_light(self) -> bool:
         """Checks if the unit has the 'light' attribute."""
-        return IS_LIGHT in self._type_data._proto.attributes
+        return self._proto.unit_type in self._bot_object.game_data.light_type_ids
 
     @property
     def is_armored(self) -> bool:
         """Checks if the unit has the 'armored' attribute."""
-        return IS_ARMORED in self._type_data._proto.attributes
+        return self._proto.unit_type in self._bot_object.game_data.armored_type_ids
 
     @property
     def is_biological(self) -> bool:
         """Checks if the unit has the 'biological' attribute."""
-        return IS_BIOLOGICAL in self._type_data._proto.attributes
+        return self._proto.unit_type in self._bot_object.game_data.biological_type_ids
 
     @property
     def is_mechanical(self) -> bool:
         """Checks if the unit has the 'mechanical' attribute."""
-        return IS_MECHANICAL in self._type_data._proto.attributes
+        return self._proto.unit_type in self._bot_object.game_data.mechanical_type_ids
 
     @property
     def is_massive(self) -> bool:
         """Checks if the unit has the 'massive' attribute."""
-        return IS_MASSIVE in self._type_data._proto.attributes
+        return self._proto.unit_type in self._bot_object.game_data.massive_type_ids
 
     @property
     def is_psionic(self) -> bool:
         """Checks if the unit has the 'psionic' attribute."""
-        return IS_PSIONIC in self._type_data._proto.attributes
+        return self._proto.unit_type in self._bot_object.game_data.psionic_type_ids
 
     @cached_property
     def tech_alias(self) -> list[UnitTypeId] | None:
@@ -703,7 +698,7 @@ def calculate_damage_vs_target(
             # TODO: hardcode hellbats when they have blueflame or attack upgrades
             for bonus in weapon.damage_bonus:
                 # More about damage bonus https://github.com/Blizzard/s2client-proto/blob/b73eb59ac7f2c52b2ca585db4399f2d3202e102a/s2clientprotocol/data.proto#L55
-                if bonus.attribute in target._type_data._proto.attributes:
+                if bonus.attribute in self._bot_object.game_data.unit_attributes.get(target._proto.unit_type, ()):
                     bonus_damage_per_upgrade = (
                         0
                         if not self.attack_upgrade_level