refactor(python): port IR address validation from Python to Rust (#348)

* refactor(python): port IR address validation from Python to Rust

Replace pure-Python validation methods with Rust-backed group
validators. The IR validation pass now delegates entire groups of
addresses to Rust in a single call instead of validating one-by-one.

Changes:
- Add check_location_group() and check_lane_group() wrapper methods
  on ArchSpec that extract _inner Rust objects and call the existing
  check_locations/check_lanes PyO3 bindings
- Rewrite validate_location() and validate_lane() to delegate to
  Rust via single-element group calls (backwards compatible)
- Rewrite compatible_lane_error() and compatible_lanes() to delegate
  to Rust group validation
- Simplify validation/address.py: replace filter_by_error loops and
  pairwise compatibility checks with single group validation calls
- Remove ~80 lines of pure-Python validation logic from ArchSpec
- Update test expected error messages to match Rust format

Closes #287

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* fix(python): address PR review feedback for Rust address validation

- Tighten return types of check_location_group and check_lane_group
  to Sequence[LocationGroupError] and Sequence[LaneGroupError]
- Add TYPE_CHECKING import for the error types
- Include lane count in validation error messages for diagnostics

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* fix(python): use proper return types in _native.pyi for check_locations/check_lanes

Update type stubs to return list[LocationGroupError] and
list[LaneGroupError] instead of list[Exception], removing the need
for type: ignore comments in the wrapper methods.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: weinbe58

python/bloqade/lanes/bytecode/_native.pyi

@@ -2,6 +2,11 @@
 
 from typing import Optional, final
 
+from bloqade.lanes.bytecode.exceptions import (
+    LaneGroupError,
+    LocationGroupError,
+)
+
 # ── Enums ──
 
 @final
@@ -724,7 +729,9 @@ class ArchSpec:
         """
         ...
 
-    def check_locations(self, locations: list[LocationAddress]) -> list[Exception]:
+    def check_locations(
+        self, locations: list[LocationAddress]
+    ) -> list[LocationGroupError]:
         """Validate a group of location addresses against this architecture.
 
         Checks for duplicate addresses and invalid word/site combinations.
@@ -733,11 +740,11 @@ class ArchSpec:
             locations (list[LocationAddress]): Location addresses to validate.
 
         Returns:
-            list[Exception]: ``LocationGroupError`` subclass instances (empty if all valid).
+            list[LocationGroupError]: Error instances (empty if all valid).
         """
         ...
 
-    def check_lanes(self, lanes: list[LaneAddress]) -> list[Exception]:
+    def check_lanes(self, lanes: list[LaneAddress]) -> list[LaneGroupError]:
         """Validate a group of lane addresses against this architecture.
 
         Checks for duplicates, invalid addresses, bus consistency, and
@@ -747,7 +754,7 @@ class ArchSpec:
             lanes (list[LaneAddress]): Lane addresses to validate.
 
         Returns:
-            list[Exception]: ``LaneGroupError`` subclass instances (empty if all valid).
+            list[LaneGroupError]: Error instances (empty if all valid).
         """
         ...
 

python/bloqade/lanes/layout/arch.py

@@ -29,6 +29,8 @@
 if TYPE_CHECKING:
     from collections.abc import Iterator
 
+    from bloqade.lanes.bytecode.exceptions import LaneGroupError, LocationGroupError
+
 
 class ArchSpec:
     """Architecture specification for a quantum device."""
@@ -378,104 +380,62 @@ def show(
         )
         plt.show()
 
+    def check_location_group(
+        self, locations: Sequence[LocationAddress]
+    ) -> Sequence[LocationGroupError]:
+        """Validate a group of location addresses via Rust.
+
+        Returns a list of LocationGroupError exceptions (empty if all valid).
+        """
+        rust_addrs = [loc._inner for loc in locations]
+        return self._inner.check_locations(rust_addrs)
+
+    def check_lane_group(
+        self, lanes: Sequence[LaneAddress]
+    ) -> Sequence[LaneGroupError]:
+        """Validate a group of lane addresses via Rust.
+
+        Checks individual lane validity, group consistency (direction, bus_id,
+        move_type), bus membership, and AOD geometry constraints.
+        Returns a list of LaneGroupError exceptions (empty if all valid).
+        """
+        rust_addrs = [lane._inner for lane in lanes]
+        return self._inner.check_lanes(rust_addrs)
+
     def compatible_lane_error(self, lane1: LaneAddress, lane2: LaneAddress) -> set[str]:
         """Get error messages if two lanes are not compatible.
 
-        NOTE: this function assumes that both lanes are valid.
+        Delegates to Rust group validation.
         """
-        errors: set[str] = set()
-        if lane1.direction != lane2.direction:
-            errors.add("Lanes have different directions")
-
-        if lane1.move_type == MoveType.SITE and lane2.move_type == MoveType.SITE:
-            if lane1.bus_id != lane2.bus_id:
-                errors.add("Lanes are on different site-buses")
-            if lane1.word_id == lane2.word_id and lane1.site_id == lane2.site_id:
-                errors.add("Lanes are the same")
-        elif lane1.move_type == MoveType.WORD and lane2.move_type == MoveType.WORD:
-            if lane2.bus_id != lane1.bus_id:
-                errors.add("Lanes are on different word-buses")
-            if lane1.word_id == lane2.word_id and lane1.site_id == lane2.site_id:
-                errors.add("Lanes are the same")
-        else:
-            errors.add("Lanes have different move types")
-
-        return errors
+        errors = self.check_lane_group([lane1, lane2])
+        return {str(e) for e in errors}
 
     def compatible_lanes(self, lane1: LaneAddress, lane2: LaneAddress) -> bool:
         """Check if two lanes are compatible (can be executed in parallel)."""
-        return len(self.compatible_lane_error(lane1, lane2)) == 0
+        return len(self.check_lane_group([lane1, lane2])) == 0
 
     def validate_location(self, location_address: LocationAddress) -> set[str]:
-        """Check if a location address is valid in this architecture."""
-        errors: set[str] = set()
+        """Check if a location address is valid in this architecture.
 
-        num_words = len(self.words)
-        if location_address.word_id >= num_words:
-            errors.add(
-                f"Word id {location_address.word_id} out of range of {num_words}"
-            )
-            return errors
-
-        word = self.words[location_address.word_id]
+        Delegates to Rust validation.
+        """
+        errors = self.check_location_group([location_address])
+        return {str(e) for e in errors}
 
-        num_sites = len(word.site_indices)
-        if location_address.site_id >= num_sites:
-            errors.add(
-                f"Site id {location_address.site_id} out of range of {num_sites}"
-            )
+    def validate_lane(self, lane_address: LaneAddress) -> set[str]:
+        """Check if a lane address is valid in this architecture.
 
-        return errors
+        Delegates to Rust validation.
+        """
+        errors = self.check_lane_group([lane_address])
+        return {str(e) for e in errors}
 
     def get_lane_address(
         self, src: LocationAddress, dst: LocationAddress
     ) -> LaneAddress | None:
         """Given an input tuple of locations, gets the lane (w/direction)."""
         return self._lane_map.get((src, dst))
 
-    def validate_lane(self, lane_address: LaneAddress) -> set[str]:
-        """Check if a lane address is valid in this architecture."""
-        errors = self.validate_location(lane_address.src_site())
-
-        if lane_address.move_type == MoveType.WORD:
-            if lane_address.site_id not in self.has_word_buses:
-                errors.add(
-                    f"Site {lane_address.site_id} does not support word-bus moves"
-                )
-            num_word_buses = len(self.word_buses)
-            if lane_address.bus_id >= num_word_buses:
-                errors.add(
-                    f"Bus id {lane_address.bus_id} out of range of {num_word_buses}"
-                )
-                return errors
-
-            bus = self.word_buses[lane_address.bus_id]
-            if lane_address.word_id not in bus.src:
-                errors.add(f"Word {lane_address.word_id} not in bus source {bus.src}")
-
-        elif lane_address.move_type == MoveType.SITE:
-            if lane_address.word_id not in self.has_site_buses:
-                errors.add(
-                    f"Word {lane_address.word_id} does not support site-bus moves"
-                )
-
-            num_site_buses = len(self.site_buses)
-            if lane_address.bus_id >= num_site_buses:
-                errors.add(
-                    f"Bus id {lane_address.bus_id} out of range of {num_site_buses}"
-                )
-                return errors
-
-            bus = self.site_buses[lane_address.bus_id]
-            if lane_address.site_id not in bus.src:
-                errors.add(f"Site {lane_address.site_id} not in bus source {bus.src}")
-        else:
-            errors.add(
-                f"Unsupported move type {lane_address.move_type} for lane address"
-            )
-
-        return errors
-
     def get_endpoints(
         self, lane_address: LaneAddress
     ) -> tuple[LocationAddress, LocationAddress]:

python/bloqade/lanes/validation/address.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass, field
 from itertools import chain
-from typing import Any, Callable, Iterable, TypeVar
+from typing import Any
 
 from kirin import interp, ir
 from kirin.analysis.forward import Forward, ForwardFrame
@@ -9,7 +9,7 @@
 
 from bloqade.lanes.dialects import move
 from bloqade.lanes.layout.arch import ArchSpec
-from bloqade.lanes.layout.encoding import Encoder, LaneAddress
+from bloqade.lanes.layout.encoding import LaneAddress, LocationAddress
 
 
 @dataclass
@@ -25,29 +25,29 @@ def method_self(self, method: ir.Method) -> EmptyLattice:
     def eval_fallback(self, frame: ForwardFrame[EmptyLattice], node: ir.Statement):
         return tuple(EmptyLattice.bottom() for _ in node.results)
 
-    AddressType = TypeVar("AddressType", bound=Encoder)
-
-    def filter_by_error(
-        self,
-        addresses: Iterable[AddressType],
-        checker: Callable[[AddressType], set[str]],
+    def report_location_errors(
+        self, node: ir.Statement, locations: tuple[LocationAddress, ...]
     ):
-        """Apply a checker function to a sequence of addresses, yielding those with errors
-        along with their error messages.
-
-        Args:
-            addresses: A tuple of address objects to be checked.
-            checker: A function that takes an address and returns a set of error messages.
-                if the set is empty, the address is considered valid.
-        Yields:
-            Tuples of (address, error message) for each address that has an error.
-        """
-
-        def has_error(tup: tuple[Any, set[str]]) -> bool:
-            return len(tup[1]) > 0
+        """Validate a group of locations via Rust and report errors."""
+        for error in self.arch_spec.check_location_group(locations):
+            self.add_validation_error(
+                node,
+                ir.ValidationError(
+                    node,
+                    f"Invalid location address: {error}",
+                ),
+            )
 
-        error_checks = zip(addresses, map(checker, addresses))
-        yield from filter(has_error, error_checks)
+    def report_lane_errors(self, node: ir.Statement, lanes: tuple[LaneAddress, ...]):
+        """Validate a group of lanes via Rust and report errors."""
+        for error in self.arch_spec.check_lane_group(lanes):
+            self.add_validation_error(
+                node,
+                ir.ValidationError(
+                    node,
+                    f"Invalid lane group (count={len(lanes)}): {error}",
+                ),
+            )
 
 
 @move.dialect.register(key="move.address.validation")
@@ -59,46 +59,8 @@ def lane_checker(
         frame: ForwardFrame[EmptyLattice],
         node: move.Move,
     ):
-        if len(node.lanes) == 0:
-            return ()
-
-        invalid_lanes = []
-        for lane, error_msgs in _interp.filter_by_error(
-            node.lanes, _interp.arch_spec.validate_lane
-        ):
-            invalid_lanes.append(lane)
-            for error_msg in error_msgs:
-                _interp.add_validation_error(
-                    node,
-                    ir.ValidationError(
-                        node,
-                        f"Invalid lane address {lane!r}: {error_msg}",
-                    ),
-                )
-
-        valid_lanes = set(node.lanes) - set(invalid_lanes)
-        if len(valid_lanes) == 0:
-            return
-
-        first_lane = valid_lanes.pop()
-        incompatible_lanes = []
-
-        def validate_compatible_lane(lane: LaneAddress):
-            return _interp.arch_spec.compatible_lane_error(first_lane, lane)
-
-        for lane, error_msgs in _interp.filter_by_error(
-            valid_lanes, validate_compatible_lane
-        ):
-            incompatible_lanes.append(lane)
-            for error_msg in error_msgs:
-                _interp.add_validation_error(
-                    node,
-                    ir.ValidationError(
-                        node,
-                        f"Incompatible lane address {first_lane!r} with lane {lane!r}: {error_msg}",
-                    ),
-                )
-
+        if len(node.lanes) > 0:
+            _interp.report_lane_errors(node, node.lanes)
         return (EmptyLattice.bottom(),)
 
     @interp.impl(move.LogicalInitialize)
@@ -111,23 +73,7 @@ def location_checker(
         frame: ForwardFrame[EmptyLattice],
         node: move.LogicalInitialize | move.LocalR | move.LocalRz | move.Fill,
     ):
-        invalid_locations = list(
-            _interp.filter_by_error(
-                node.location_addresses,
-                _interp.arch_spec.validate_location,
-            )
-        )
-
-        for lane_address, error_msgs in invalid_locations:
-            for error_msg in error_msgs:
-                _interp.add_validation_error(
-                    node,
-                    ir.ValidationError(
-                        node,
-                        f"Invalid location address {lane_address!r}: {error_msg}",
-                    ),
-                )
-
+        _interp.report_location_errors(node, node.location_addresses)
         return (EmptyLattice.bottom(),)
 
     @interp.impl(move.GetFutureResult)
@@ -137,17 +83,7 @@ def location_checker_get_future(
         frame: ForwardFrame[EmptyLattice],
         node: move.GetFutureResult,
     ):
-        location_address = node.location_address
-        error_msgs = _interp.arch_spec.validate_location(location_address)
-
-        for error_msg in error_msgs:
-            _interp.add_validation_error(
-                node,
-                ir.ValidationError(
-                    node,
-                    f"Invalid location address {location_address!r}: {error_msg}",
-                ),
-            )
+        _interp.report_location_errors(node, (node.location_address,))
 
     @interp.impl(move.PhysicalInitialize)
     def location_checker_physical(
@@ -156,22 +92,8 @@ def location_checker_physical(
         frame: ForwardFrame[EmptyLattice],
         node: move.PhysicalInitialize,
     ):
-        invalid_locations = list(
-            _interp.filter_by_error(
-                chain.from_iterable(node.location_addresses),
-                _interp.arch_spec.validate_location,
-            )
-        )
-
-        for lane_address, error_msgs in invalid_locations:
-            for error_msg in error_msgs:
-                _interp.add_validation_error(
-                    node,
-                    ir.ValidationError(
-                        node,
-                        f"Invalid location address {lane_address!r}: {error_msg}",
-                    ),
-                )
+        all_locations = tuple(chain.from_iterable(node.location_addresses))
+        _interp.report_location_errors(node, all_locations)
 
 
 @dataclass

python/tests/arch/gemini/logical/test_spec.py

@@ -70,8 +70,8 @@ def plot():
 
 def invalid_locations():
     arch_spec = logical.get_arch_spec()
-    yield arch_spec, LocationAddress(16, 0), set(["Word id 16 out of range of 2"])
-    yield arch_spec, LocationAddress(0, 32), set(["Site id 32 out of range of 10"])
+    yield arch_spec, LocationAddress(16, 0), {"invalid location word_id=16, site_id=0"}
+    yield arch_spec, LocationAddress(0, 32), {"invalid location word_id=0, site_id=32"}
 
 
 @pytest.mark.parametrize("arch_spec, location_address, message", invalid_locations())