latest attempt to try to reconcile type lattices

Author: johnzl-777

src/bloqade/analysis/measure_id/analysis.py

@@ -4,12 +4,17 @@
 from kirin.analysis import ForwardExtra
 from kirin.analysis.forward import ForwardFrame
 
-from .lattice import MeasureId, NotMeasureId, KnownMeasureId, MeasureIdTuple
+from .lattice import (
+    MeasureId,
+    NotMeasureId,
+    RawMeasureId,
+    MeasureIdTuple,
+)
 
 
 @dataclass
 class GlobalRecordState:
-    buffer: list[KnownMeasureId] = field(default_factory=list)
+    buffer: list[RawMeasureId] = field(default_factory=list)
 
     # assume that this KnownMeasureId will always be -1
     def add_record_idxs(self, num_new_records: int) -> MeasureIdTuple:
@@ -18,7 +23,7 @@ def add_record_idxs(self, num_new_records: int) -> MeasureIdTuple:
             record_idx.idx -= num_new_records
 
         # generate new indices and add them to the buffer
-        new_record_idxs = [KnownMeasureId(-i) for i in range(num_new_records, 0, -1)]
+        new_record_idxs = [RawMeasureId(-i) for i in range(num_new_records, 0, -1)]
         self.buffer += new_record_idxs
         # Return for usage, idxs linked to the global state
         return MeasureIdTuple(data=tuple(new_record_idxs))
@@ -30,13 +35,13 @@ def add_record_idxs(self, num_new_records: int) -> MeasureIdTuple:
     # buffer
     def clone_record_idxs(self, measure_id_tuple: MeasureIdTuple) -> MeasureIdTuple:
         cloned_members = []
-        for known_measure_id in measure_id_tuple.data:
-            assert isinstance(known_measure_id, KnownMeasureId)
-            cloned_known_measure_id = KnownMeasureId(known_measure_id.idx)
+        for raw_measure_id in measure_id_tuple.data:
+            assert isinstance(raw_measure_id, RawMeasureId)
+            cloned_raw_measure_id = RawMeasureId(raw_measure_id.idx)
             # put into the global buffer but also
             # return an analysis-facing copy
-            self.buffer.append(cloned_known_measure_id)
-            cloned_members.append(cloned_known_measure_id)
+            self.buffer.append(cloned_raw_measure_id)
+            cloned_members.append(cloned_raw_measure_id)
         return MeasureIdTuple(data=tuple(cloned_members))
 
     def offset_existing_records(self, offset: int):

src/bloqade/analysis/measure_id/impls.py

@@ -10,11 +10,12 @@
     Predicate,
     AnyMeasureId,
     NotMeasureId,
-    KnownMeasureId,
+    RawMeasureId,
     MeasureIdTuple,
     ConstantCarrier,
     InvalidMeasureId,
     ImmutableMeasureIds,
+    PredicatedMeasureId,
 )
 from .analysis import MeasureIDFrame, MeasurementIDAnalysis
 
@@ -61,7 +62,7 @@ def measurement_predicate(
     ):
         original_measure_id_tuple = frame.get(stmt.measurements)
         if not all(
-            isinstance(measure_id, KnownMeasureId)
+            isinstance(measure_id, RawMeasureId)
             for measure_id in original_measure_id_tuple.data
         ):
             return (InvalidMeasureId(),)
@@ -76,10 +77,10 @@ def measurement_predicate(
             return (InvalidMeasureId(),)
 
         predicate_measure_ids = [
-            KnownMeasureId(measure_id.idx, predicate)
+            PredicatedMeasureId(measure_id.idx, predicate)
             for measure_id in original_measure_id_tuple.data
         ]
-        return (MeasureIdTuple(data=tuple(predicate_measure_ids)),)
+        return (ImmutableMeasureIds(data=tuple(predicate_measure_ids)),)
 
 
 @gemini.logical.dialect.register(key="measure_id")
@@ -100,11 +101,13 @@ def terminal_measurement(
             return (AnyMeasureId(),)
 
         measure_id_bools = []
-        for _ in range(num_qubits.data):
-            interp.measure_count += 1
-            measure_id_bools.append(RawMeasureId(interp.measure_count))
+        for i in range(num_qubits.data):
+            measure_id_bools.append(RawMeasureId(idx=-(i + 1)))
 
-        return (MeasureIdTuple(data=tuple(measure_id_bools)),)
+        # Immutable usually desired for stim generation
+        # but we can reuse it here to indicate
+        # the measurement ids should not change anymore.
+        return (ImmutableMeasureIds(data=tuple(measure_id_bools)),)
 
 
 @annotate.dialect.register(key="measure_id")
@@ -121,7 +124,9 @@ def consumes_measurements(
 
         if not (
             isinstance(measure_id_tuple_at_stmt, MeasureIdTuple)
-            and kirin_types.is_tuple_of(measure_id_tuple_at_stmt.data, KnownMeasureId)
+            and kirin_types.is_tuple_of(
+                measure_id_tuple_at_stmt.data, PredicatedMeasureId
+            )
         ):
             return (InvalidMeasureId(),)
 
@@ -241,7 +246,7 @@ def invoke(
 
 
 @scf.dialect.register(key="measure_id")
-class LoopHandling(interp.MethodTable):
+class ScfHandling(interp.MethodTable):
     @interp.impl(scf.stmts.For)
     def for_loop(
         self, interp_: MeasurementIDAnalysis, frame: MeasureIDFrame, stmt: scf.stmts.For
@@ -323,7 +328,7 @@ def for_loop(
             if isinstance(var, MeasureIdTuple):
                 for member in var.data:
                     if (
-                        isinstance(member, KnownMeasureId)
+                        isinstance(member, RawMeasureId)
                         and member.idx not in witnessed_record_idxs
                     ):
                         witnessed_record_idxs.add(member.idx)

src/bloqade/analysis/measure_id/lattice.py

@@ -68,12 +68,23 @@ def is_subseteq(self, other: MeasureId) -> bool:
 
 @final
 @dataclass
-class KnownMeasureId(MeasureId):
+class RawMeasureId(MeasureId):
+    idx: int
+
+    def is_subseteq(self, other: MeasureId) -> bool:
+        if isinstance(other, RawMeasureId):
+            return self.idx == other.idx
+        return False
+
+
+@final
+@dataclass
+class PredicatedMeasureId(MeasureId):
     idx: int
     predicate: Predicate
 
     def is_subseteq(self, other: MeasureId) -> bool:
-        if isinstance(other, KnownMeasureId):
+        if isinstance(other, PredicatedMeasureId):
             return self.idx == other.idx and self.predicate == other.predicate
         return False
 
@@ -92,7 +103,10 @@ def is_subseteq(self, other: MeasureId) -> bool:
 @final
 @dataclass
 class ImmutableMeasureIds(MeasureId):
-    data: tuple[KnownMeasureId, ...]
+    # SetDetector happily consumes RawMeasureIds, but
+    # for scf.IfElse rewrite with predicates I need to allow
+    # PredicatedMeasureIds as well.
+    data: tuple[PredicatedMeasureId | RawMeasureId, ...]
 
     def is_subseteq(self, other: MeasureId) -> bool:
         if isinstance(other, ImmutableMeasureIds):

src/bloqade/stim/passes/squin_to_stim.py

@@ -29,6 +29,8 @@
 
 from ..rewrite import IfToStim, SetDetectorToStim, SetObservableToStim
 
+# from bloqade.stim.passes.soft_flatten import SoftFlatten
+
 
 @dataclass
 class SquinToStimPass(Pass):

src/bloqade/stim/rewrite/get_record_util.py

@@ -2,7 +2,7 @@
 from kirin.dialects import py
 
 from bloqade.stim.dialects import auxiliary
-from bloqade.analysis.measure_id.lattice import KnownMeasureId, MeasureIdTuple
+from bloqade.analysis.measure_id.lattice import MeasureIdTuple, PredicatedMeasureId
 
 
 def insert_get_records(node: ir.Statement, measure_id_tuple: MeasureIdTuple):
@@ -11,7 +11,7 @@ def insert_get_records(node: ir.Statement, measure_id_tuple: MeasureIdTuple):
     """
     get_record_ssas = []
     for known_measure_id in measure_id_tuple.data:
-        assert isinstance(known_measure_id, KnownMeasureId)
+        assert isinstance(known_measure_id, PredicatedMeasureId)
         idx_stmt = py.constant.Constant(known_measure_id.idx)
         idx_stmt.insert_before(node)
         get_record_stmt = auxiliary.GetRecord(idx_stmt.result)

src/bloqade/stim/rewrite/ifs_to_stim.py

@@ -13,7 +13,7 @@
 from bloqade.stim.dialects.gate import CX as stim_CX, CY as stim_CY, CZ as stim_CZ
 from bloqade.analysis.measure_id import MeasureIDFrame
 from bloqade.stim.dialects.auxiliary import GetRecord
-from bloqade.analysis.measure_id.lattice import Predicate, KnownMeasureId
+from bloqade.analysis.measure_id.lattice import Predicate, PredicatedMeasureId
 
 
 @dataclass
@@ -140,7 +140,7 @@ def rewrite_IfElse(self, stmt: scf.IfElse) -> RewriteResult:
         condition_type = self.measure_frame.entries.get(stmt.cond)
         # Check the condition is a singular MeasurementIdBool and that
         # it was generated by querying if the measurement is equivalent to the one state
-        if not isinstance(condition_type, KnownMeasureId):
+        if not isinstance(condition_type, PredicatedMeasureId):
             return RewriteResult()
 
         if condition_type.predicate != Predicate.IS_ONE:
@@ -162,12 +162,8 @@ def rewrite_IfElse(self, stmt: scf.IfElse) -> RewriteResult:
             return RewriteResult()
 
         # generate get record statement
-        num_measures = self.measure_frame.num_measures_at_stmt.get(stmt)
-        if num_measures is None:
-            return RewriteResult()
-
-        measure_id_idx_stmt = py.Constant((condition_type.idx - 1) - num_measures)
-        get_record_stmt = GetRecord(id=measure_id_idx_stmt.result)  # noqa: F841
+        measure_id_idx_stmt = py.Constant(condition_type.idx)
+        get_record_stmt = GetRecord(id=measure_id_idx_stmt.result)
 
         address_attr = stmts[0].qubits.hints.get("address")