Stim REPEAT support

_typos.toml

@@ -14,4 +14,5 @@ mch = "mch"
 IY = "IY"
 ket = "ket"
 typ = "typ"
+anc_qs = "anc_qs" # "ancilla qubits" variable abbreivation - used to be autocorrected to AND_QS! :sigh:
 OT = "OT"

src/bloqade/analysis/measure_id/__init__.py

@@ -1,4 +1,4 @@
-from . import impls as impls
+from . import util as util, impls as impls
 from .analysis import (
     MeasureIDFrame as MeasureIDFrame,
     MeasurementIDAnalysis as MeasurementIDAnalysis,

src/bloqade/analysis/measure_id/accumulator.py

@@ -0,0 +1,110 @@
+"""
+Helper functions for detecting accumulator semantics
+(chiefly the pattern of using lists to store a growing number of measurements)
+and expanding them with the correct number of RawMeasureId objects in the correct order.
+"""
+
+from kirin import ir
+from kirin.dialects import py, scf
+
+from .util import get_scf_for_repeat_count
+from .lattice import (
+    MeasureId,
+    MutableIdx,
+    AnyMeasureId,
+    RawMeasureId,
+    MeasureIdTuple,
+)
+from .analysis import MeasureIDFrame
+
+
+def is_growing_accumulator(
+    init_var: MeasureId,
+    first_var: MeasureId,
+    second_var: MeasureId,
+) -> bool:
+    if not (
+        isinstance(init_var, MeasureIdTuple)
+        and isinstance(first_var, MeasureIdTuple)
+        and isinstance(second_var, MeasureIdTuple)
+    ):
+        return False
+
+    first_growth = len(first_var.data) - len(init_var.data)
+    second_growth = len(second_var.data) - len(first_var.data)
+    return first_growth > 0 and first_growth == second_growth
+
+
+def detect_append_order_from_ir(
+    stmt: scf.stmts.For,
+    var_idx: int,
+) -> bool | None:
+    """Detect append vs prepend by inspecting the py.Add in the loop body.
+
+    Returns True for append (acc + ms), False for prepend (ms + acc),
+    or None if the pattern cannot be determined from the IR.
+    """
+    body_block = stmt.body.blocks[0]
+    acc_block_arg = body_block.args[var_idx + 1]
+
+    yield_stmt = body_block.last_stmt
+    if not isinstance(yield_stmt, scf.stmts.Yield):
+        return None
+
+    yielded_val = yield_stmt.values[var_idx]
+    if not isinstance(yielded_val, ir.ResultValue):
+        return None
+
+    add_stmt = yielded_val.owner
+    if not isinstance(add_stmt, py.Add):
+        return None
+
+    if add_stmt.lhs is acc_block_arg:
+        return True
+    if add_stmt.rhs is acc_block_arg:
+        return False
+    return None
+
+
+def expand_accumulator(
+    stmt: scf.stmts.For,
+    frame: MeasureIDFrame,
+    init_var: MeasureIdTuple,
+    first_var: MeasureIdTuple,
+    is_append: bool,
+) -> MeasureId:
+    """Expand a growing accumulator to its final concrete MeasureIdTuple.
+
+    Uses the statically known loop count to compute the total number of
+    accumulated measurements and constructs the final tuple with correct
+    MutableIdx values registered in the parent GlobalRecordState.
+    """
+    num_iterations = get_scf_for_repeat_count(stmt)
+    if num_iterations is None:
+        return AnyMeasureId()
+
+    measurements_per_iter = len(first_var.data) - len(init_var.data)
+    total_new_measurements = num_iterations * measurements_per_iter
+
+    frame.global_record_state.offset_existing_records(total_new_measurements)
+
+    new_mutable_idxs = [MutableIdx(-i) for i in range(total_new_measurements, 0, -1)]
+    frame.global_record_state.buffer.extend(new_mutable_idxs)
+    new_raw_ids = [RawMeasureId(idx) for idx in new_mutable_idxs]
+
+    if is_append:
+        final_data = init_var.data + tuple(new_raw_ids)
+    else:
+        # you're prepending here, so your ordering is "staggered"
+        # ex: acc = ms + acc where acc is (-3, -2, -1)
+        # the ms happens on 3 qubits that bumps your existing acc to (-6, -5, -4)
+        # then you need to add your (-3, -2, -1) to the FRONT of the tuple so it looks like (-3, -2, -1, -6, -5, -4)
+        chunks: list[RawMeasureId] = []
+        for i in range(num_iterations - 1, -1, -1):
+            chunk = new_raw_ids[
+                i * measurements_per_iter : (i + 1) * measurements_per_iter
+            ]
+            chunks.extend(chunk)
+        final_data = tuple(chunks) + init_var.data
+
+    return MeasureIdTuple(data=final_data, obj_type=first_var.obj_type)

src/bloqade/analysis/measure_id/analysis.py

@@ -1,24 +1,95 @@
+from typing import Iterable
 from dataclasses import field, dataclass
 
 from kirin import ir
 from kirin.analysis import ForwardExtra
+from kirin.dialects import ilist
 from typing_extensions import Self
 from kirin.analysis.forward import ForwardFrame
 
-from .lattice import MeasureId, NotMeasureId
+from .lattice import (
+    MeasureId,
+    MutableIdx,
+    NotMeasureId,
+    RawMeasureId,
+    MeasureIdTuple,
+)
+
+
+@dataclass
+class GlobalRecordState:
+    type_for_scf_conds: dict[ir.Statement, MeasureId] = field(default_factory=dict)
+    buffer: list[MutableIdx] = field(default_factory=list)
+
+    def add_record_idxs(self, num_new_records: int) -> MeasureIdTuple:
+        # Adjust all previous indices
+        for mutable_idx in self.buffer:
+            mutable_idx.value -= num_new_records
+
+        # Generate new MutableIdx entries and add to buffer
+        new_mutable_idxs = [MutableIdx(-i) for i in range(num_new_records, 0, -1)]
+        self.buffer += new_mutable_idxs
+
+        # Create RawMeasureIds referencing these MutableIdxs
+        new_record_idxs = [RawMeasureId(idx) for idx in new_mutable_idxs]
+        return MeasureIdTuple(data=tuple(new_record_idxs), obj_type=ilist.IList)
+
+    def clone_measure_id_tuple(
+        self, measure_id_tuple: MeasureIdTuple
+    ) -> MeasureIdTuple:
+        cloned_members = []
+        for measure_id in measure_id_tuple.data:
+            cloned_measure_id = self.clone_measure_ids(measure_id)
+            cloned_members.append(cloned_measure_id)
+        return MeasureIdTuple(
+            data=tuple(cloned_members),
+            obj_type=measure_id_tuple.obj_type,
+        )
+
+    def clone_raw_measure_id(self, raw_measure_id: RawMeasureId) -> RawMeasureId:
+        # Create new MutableIdx for independent tracking
+        new_mutable_idx = MutableIdx(raw_measure_id.idx)
+        self.buffer.append(new_mutable_idx)
+
+        # Return RawMeasureId with same predicate, referencing new MutableIdx
+        return RawMeasureId(new_mutable_idx, predicate=raw_measure_id.predicate)
+
+    def clone_measure_ids(self, measure_id_type: MeasureId) -> MeasureId:
+        if isinstance(measure_id_type, RawMeasureId):
+            return self.clone_raw_measure_id(measure_id_type)
+        elif isinstance(measure_id_type, MeasureIdTuple):
+            return self.clone_measure_id_tuple(measure_id_type)
+        return None
+
+    def offset_existing_records(self, offset: int):
+        for mutable_idx in self.buffer:
+            mutable_idx.value -= offset
+
+    def add_unique_mutable_idxs(self, mutable_idxs: Iterable[MutableIdx]) -> None:
+        """Add MutableIdx objects to buffer, skipping duplicates."""
+        existing = set(self.buffer)
+        for idx in mutable_idxs:
+            if idx not in existing:
+                existing.add(idx)
+                self.buffer.append(idx)
 
 
 @dataclass
 class MeasureIDFrame(ForwardFrame[MeasureId]):
-    num_measures_at_stmt: dict[ir.Statement, int] = field(default_factory=dict)
+    # Keeps track of every measurement's record index, containing methods to
+    # update and offset record indices as new measurements/loops are encountered.
+    global_record_state: GlobalRecordState = field(default_factory=GlobalRecordState)
+    # Associates with each scf.IfElse statement encountered the lattice element
+    # of the condition, which is detached from the GlobalRecordState to
+    # prevent further mutation as more measurements are encountered.
+    type_for_scf_conds: dict[ir.Statement, MeasureId] = field(default_factory=dict)
+    measure_count_offset: int = 0
 
 
 class MeasurementIDAnalysis(ForwardExtra[MeasureIDFrame, MeasureId]):
 
     keys = ["measure_id"]
     lattice = MeasureId
-    # for every kind of measurement encountered, increment this
-    # then use this to generate the negative values for target rec indices
     measure_count = 0
     detector_count = 0
     observable_count = 0
@@ -34,8 +105,6 @@ def initialize_frame(
     ) -> MeasureIDFrame:
         return MeasureIDFrame(node, has_parent_access=has_parent_access)
 
-    # Still default to bottom,
-    # but let constants return the softer "NoMeasureId" type from impl
     def eval_fallback(
         self, frame: ForwardFrame[MeasureId], node: ir.Statement
     ) -> tuple[MeasureId, ...]: