perf(search): cache ConfigurationNode derived occupancy and lookup state (#334)

* perf(search): cache ConfigurationNode derived occupancy and lookup state

Precompute configuration key, occupied locations, and location-to-qubit mappings in ConfigurationNode to avoid repeated recomputation in search hot paths. Add a stability regression test for repeated derived-view and lookup accesses.

Made-with: Cursor

* perf(search): cache lane indexes for source and triplet lookups (#327)

Precompute lane indexes in ConfigurationTree to avoid rebuilding lane mappings in hot search loops. This adds constant-time helpers for triplet and source-based lane queries used during traversal and scoring.

Made-with: Cursor

Co-authored-by: Jason Ludmir <jasonludmir@JLudmir-Mac-QDNH.local>
Co-authored-by: Phillip Weinberg <weinbe58@gmail.com>

* perf(search): use cached_property for ConfigurationNode accessors

Align ConfigurationNode caching with review guidance by using cached_property on derived accessors while preserving existing cached values.

Made-with: Cursor

---------

Co-authored-by: Jason Ludmir <jasonludmir@JLudmir-Mac-QDNH.local>
Co-authored-by: Phillip Weinberg <weinbe58@gmail.com>

Author: 3 people

python/bloqade/lanes/search/configuration.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
+from functools import cached_property
 
 from bloqade.lanes.layout import LaneAddress, LocationAddress
 
@@ -37,30 +38,40 @@ class ConfigurationNode:
     depth: int = 0
     """Distance from the root node."""
 
-    @property
+    _config_key: Configuration = field(init=False, repr=False, compare=False)
+    _occupied_locations: frozenset[LocationAddress] = field(
+        init=False, repr=False, compare=False
+    )
+    _qubit_at_location: dict[LocationAddress, int] = field(
+        init=False, repr=False, compare=False
+    )
+
+    def __post_init__(self) -> None:
+        self._config_key = frozenset(self.configuration.items())
+        self._occupied_locations = frozenset(self.configuration.values())
+        self._qubit_at_location = {loc: qid for qid, loc in self.configuration.items()}
+
+    @cached_property
     def config_key(self) -> Configuration:
         """Canonical hashable key for this configuration.
 
         Two nodes with the same atom placement (regardless of history)
         produce the same config_key.
         """
-        return frozenset(self.configuration.items())
+        return self._config_key
 
-    @property
+    @cached_property
     def occupied_locations(self) -> frozenset[LocationAddress]:
         """The set of physical locations currently occupied by atoms."""
-        return frozenset(self.configuration.values())
+        return self._occupied_locations
 
     def is_occupied(self, location: LocationAddress) -> bool:
         """Check whether a physical location has an atom."""
-        return location in self.configuration.values()
+        return location in self._occupied_locations
 
     def get_qubit_at(self, location: LocationAddress) -> int | None:
         """Return the qubit ID at a location, or None if empty."""
-        for qid, loc in self.configuration.items():
-            if loc == location:
-                return qid
-        return None
+        return self._qubit_at_location.get(location)
 
     def path_to_root(self) -> list[frozenset[LaneAddress]]:
         """Walk from this node to the root, returning move sets in root-to-leaf order.

python/bloqade/lanes/search/tree.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from collections import defaultdict
 from collections.abc import Iterator
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
@@ -41,10 +42,70 @@ class ConfigurationTree:
     seen: dict[Configuration, ConfigurationNode] = field(
         default_factory=dict, init=False, repr=False
     )
+    _lanes_by_triplet: dict[
+        tuple[MoveType, int, Direction], tuple[LaneAddress, ...]
+    ] = field(default_factory=dict, init=False, repr=False)
+    _lane_by_src: dict[
+        tuple[MoveType, int, Direction], dict[LocationAddress, LaneAddress]
+    ] = field(default_factory=dict, init=False, repr=False)
+    _outgoing_lanes_by_src: dict[LocationAddress, tuple[LaneAddress, ...]] = field(
+        default_factory=dict, init=False, repr=False
+    )
 
     def __post_init__(self) -> None:
         self.path_finder = PathFinder(self.arch_spec)
         self.seen[self.root.config_key] = self.root
+        self._build_lane_indexes()
+
+    def _build_lane_indexes(self) -> None:
+        """Precomputes all lane mappings once (meant to be called at tree construction time)."""
+        lanes_by_triplet: dict[tuple[MoveType, int, Direction], list[LaneAddress]] = {}
+        lane_by_src: dict[
+            tuple[MoveType, int, Direction], dict[LocationAddress, LaneAddress]
+        ] = {}
+        outgoing: dict[LocationAddress, list[LaneAddress]] = defaultdict(list)
+
+        for mt in (MoveType.SITE, MoveType.WORD):
+            buses = (
+                self.arch_spec.site_buses
+                if mt == MoveType.SITE
+                else self.arch_spec.word_buses
+            )
+            for bus_id, bus in enumerate(buses):
+                for direction in (Direction.FORWARD, Direction.BACKWARD):
+                    key = (mt, bus_id, direction)
+                    lanes_for_key: list[LaneAddress] = []
+                    src_map: dict[LocationAddress, LaneAddress] = {}
+                    if mt == MoveType.SITE:
+                        for word_id in self.arch_spec.has_site_buses:
+                            for site_id in bus.src:
+                                lane = LaneAddress(
+                                    mt, word_id, site_id, bus_id, direction
+                                )
+                                src, _ = self.arch_spec.get_endpoints(lane)
+                                lanes_for_key.append(lane)
+                                src_map[src] = lane
+                                outgoing[src].append(lane)
+                    else:
+                        for word_id in bus.src:
+                            for site_id in self.arch_spec.has_word_buses:
+                                lane = LaneAddress(
+                                    mt, word_id, site_id, bus_id, direction
+                                )
+                                src, _ = self.arch_spec.get_endpoints(lane)
+                                lanes_for_key.append(lane)
+                                src_map[src] = lane
+                                outgoing[src].append(lane)
+                    lanes_by_triplet[key] = lanes_for_key
+                    lane_by_src[key] = src_map
+
+        self._lanes_by_triplet = {
+            key: tuple(values) for key, values in lanes_by_triplet.items()
+        }
+        self._lane_by_src = lane_by_src
+        self._outgoing_lanes_by_src = {
+            src: tuple(values) for src, values in outgoing.items()
+        }
 
     @classmethod
     def from_initial_placement(
@@ -80,16 +141,21 @@ def lanes_for(
         Yields:
             LaneAddress values.
         """
-        if move_type == MoveType.SITE:
-            bus = self.arch_spec.site_buses[bus_id]
-            for w in self.arch_spec.has_site_buses:
-                for s in bus.src:
-                    yield LaneAddress(move_type, w, s, bus_id, direction)
-        else:
-            bus = self.arch_spec.word_buses[bus_id]
-            for w in bus.src:
-                for s in self.arch_spec.has_word_buses:
-                    yield LaneAddress(move_type, w, s, bus_id, direction)
+        yield from self._lanes_by_triplet[(move_type, bus_id, direction)]
+
+    def lane_for_source(
+        self,
+        move_type: MoveType,
+        bus_id: int,
+        direction: Direction,
+        source: LocationAddress,
+    ) -> LaneAddress | None:
+        """Resolve one lane by source for a specific triplet."""
+        return self._lane_by_src[(move_type, bus_id, direction)].get(source)
+
+    def outgoing_lanes(self, source: LocationAddress) -> tuple[LaneAddress, ...]:
+        """Return all precomputed outgoing lanes from source."""
+        return self._outgoing_lanes_by_src.get(source, ())
 
     def valid_lanes(
         self,

python/tests/search/test_configuration.py

@@ -77,6 +77,29 @@ def test_occupied_locations():
     assert node.occupied_locations == frozenset({loc0, loc1})
 
 
+def test_cached_derived_views_and_lookup_methods_are_stable():
+    loc0 = LocationAddress(0, 0)
+    loc1 = LocationAddress(1, 0)
+    empty_loc = LocationAddress(0, 5)
+    node = ConfigurationNode(configuration={0: loc0, 1: loc1})
+
+    first_key = node.config_key
+    first_occupied = node.occupied_locations
+
+    # Repeated access should return stable cached objects.
+    assert node.config_key is first_key
+    assert node.occupied_locations is first_occupied
+
+    # Lookup helpers should remain stable across repeated calls.
+    for _ in range(3):
+        assert node.get_qubit_at(loc0) == 0
+        assert node.get_qubit_at(loc1) == 1
+        assert node.get_qubit_at(empty_loc) is None
+        assert node.is_occupied(loc0) is True
+        assert node.is_occupied(loc1) is True
+        assert node.is_occupied(empty_loc) is False
+
+
 def test_path_to_root_at_root():
     node = ConfigurationNode(configuration={0: LocationAddress(0, 0)})
     assert node.path_to_root() == []