From 2d83765098dc839c3042ec62ab9429a25e4f6f4a Mon Sep 17 00:00:00 2001 From: SamOgon-one <114773312+SamOgon-one@users.noreply.github.com> Date: Fri, 13 Nov 2026 09:20:42 +0100 Subject: [PATCH 1/3] perf: optimize `_abilities_count_and_build_progress` ### Optimize `_abilities_count_and_build_progress` #### Description Optimizes `_abilities_count_and_build_progress` by eliminating intermediate Python object allocations and reducing attribute lookup overhead in the main loop over `all_own_units`. #### How it optimizes & speeds up: 1. **Bypasses `UnitOrder` instantiation:** Instead of iterating over `unit.orders` (which creates a new `UnitOrder` instance and converts target coordinates to `Point2` for every single order), we read directly from `unit._proto.orders`. 2. **Fast Enum resolution:** Uses `AbilityId._value2member_map_` cached in a local variable for direct $O(1)$ dict lookups of `order.ability_id`. 3. **Loop-invariant hoisting:** Hoists `is_terran = self.race == Race.Terran` and `structure_type_ids = self.game_data.structure_type_ids` out of the unit loop. 4. **Direct protobuf field access:** Reads `proto.build_progress` and `proto.unit_type` directly to avoid property overhead where applicable. #### Benchmark / Profiling: According to `py-spy` profiling across a 10-minute Protoss game: * **Before:** `_abilities_count_and_build_progress` accounted for **3.5%** of total CPU runtime. * **After:** Reduced to **2.6%** of total CPU runtime (~26% relative reduction for this method). --- sc2/bot_ai_internal.py | 61 +++++++++++++++++++++++++----------------- 1 file changed, 37 insertions(+), 24 deletions(-) diff --git a/sc2/bot_ai_internal.py b/sc2/bot_ai_internal.py index 8bf8fa9f..797fb7a9 100644 --- a/sc2/bot_ai_internal.py +++ b/sc2/bot_ai_internal.py @@ -429,31 +429,44 @@ def _abilities_count_and_build_progress(self) -> tuple[Counter[AbilityId], dict[ all units in production and all structures, and all morphs""" abilities_amount: Counter[AbilityId] = Counter() max_build_progress: dict[AbilityId, float] = {} + ability_by_id = AbilityId._value2member_map_ + structure_type_ids = self.game_data.structure_type_ids + is_terran = self.race == Race.Terran + unit: Unit - 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 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: - if unit.type_id == UnitTypeId.ARCHON: - # Hotfix for archons in morph state - creation_ability_id = AbilityId.ARCHON_WARP_TARGET - abilities_amount[creation_ability_id] += 2 - else: - # Hotfix for rich geysirs - creation_ability_id = CREATION_ABILITY_FIX[unit.type_id] - abilities_amount[creation_ability_id] += 1 - else: - creation_ability = self.game_data.units[unit.type_id.value].creation_ability - if creation_ability is None: - continue - creation_ability_id = creation_ability.exact_id - abilities_amount[creation_ability_id] += 1 - max_build_progress[creation_ability_id] = max( - max_build_progress.get(creation_ability_id, 0), unit.build_progress - ) + for unit in self.all_own_units: + proto = unit._proto + + # Read the raw protobuf orders: this avoids constructing a + # UnitOrder and a Point2 for every order. + for order in proto.orders: + abilities_amount[ability_by_id[order.ability_id]] += 1 + + if proto.build_progress == 1: + continue + # An SCV constructing a Terran structure is already counted by its + # build order; do not count the structure a second time. + if is_terran and proto.unit_type in structure_type_ids: + continue + + unit_type = unit.type_id + creation_ability = CREATION_ABILITY_FIX.get(unit_type) + if creation_ability is None: + creation_ability_data = self.game_data.units[proto.unit_type].creation_ability + if creation_ability_data is None: + continue + creation_ability = creation_ability_data.exact_id + abilities_amount[creation_ability] += 1 + elif unit_type == UnitTypeId.ARCHON: + creation_ability = AbilityId.ARCHON_WARP_TARGET + abilities_amount[creation_ability] += 2 + else: + # Hotfix for rich geysirs, cocoons, interceptors, etc. + abilities_amount[creation_ability] += 1 + + max_build_progress[creation_ability] = max( + max_build_progress.get(creation_ability, 0), proto.build_progress + ) return abilities_amount, max_build_progress From 396557d7a19e6c014dd0e308702500957d807197 Mon Sep 17 00:00:00 2001 From: SamOgon-one <114773312+SamOgon-one@users.noreply.github.com> Date: Fri, 13 Nov 2026 09:47:24 +0100 Subject: [PATCH 2/3] Cast ability_by_id to Dict[int, AbilityId] Cast ability_by_id to the correct type for better type safety. --- sc2/bot_ai_internal.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sc2/bot_ai_internal.py b/sc2/bot_ai_internal.py index 797fb7a9..3d0d211e 100644 --- a/sc2/bot_ai_internal.py +++ b/sc2/bot_ai_internal.py @@ -8,7 +8,7 @@ from collections import Counter from collections.abc import Generator, Iterable from contextlib import suppress -from typing import TYPE_CHECKING, Any, final +from typing import TYPE_CHECKING, Any, cast, final import numpy as np from loguru import logger @@ -429,7 +429,7 @@ def _abilities_count_and_build_progress(self) -> tuple[Counter[AbilityId], dict[ all units in production and all structures, and all morphs""" abilities_amount: Counter[AbilityId] = Counter() max_build_progress: dict[AbilityId, float] = {} - ability_by_id = AbilityId._value2member_map_ + ability_by_id = cast(Dict[int, AbilityId], AbilityId._value2member_map_) structure_type_ids = self.game_data.structure_type_ids is_terran = self.race == Race.Terran From 8bfe2784726469cf760084e0913b7af872bdcedf Mon Sep 17 00:00:00 2001 From: SamOgon-one <114773312+SamOgon-one@users.noreply.github.com> Date: Fri, 13 Nov 2026 09:57:38 +0100 Subject: [PATCH 3/3] Fix type hint for ability_by_id variable --- sc2/bot_ai_internal.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sc2/bot_ai_internal.py b/sc2/bot_ai_internal.py index 3d0d211e..11c75ff1 100644 --- a/sc2/bot_ai_internal.py +++ b/sc2/bot_ai_internal.py @@ -429,7 +429,7 @@ def _abilities_count_and_build_progress(self) -> tuple[Counter[AbilityId], dict[ all units in production and all structures, and all morphs""" abilities_amount: Counter[AbilityId] = Counter() max_build_progress: dict[AbilityId, float] = {} - ability_by_id = cast(Dict[int, AbilityId], AbilityId._value2member_map_) + ability_by_id = cast(dict[int, AbilityId], AbilityId._value2member_map_) structure_type_ids = self.game_data.structure_type_ids is_terran = self.race == Race.Terran