Support Slice propogation in MeasurementIDAnalysis (#438)

I was working on a more advanced unit test for SquinToStim when I ran
into this problem.

If you use slice on a list of measurements the analysis doesn't go
through. The solution was to just carry over the logic I figured out for
the address analysis. However I cooked up a case where even this is
insufficient and figured out the simplest addition to the logic to amend
this case.


Basically both measure and address have something like the following:
```
# container could be AddressTuple or MeasureIDTuple
getitem(container, slice) -> tuple(....)
```
Where tuple is literally a plain Python tuple. This doesn't present any
problems until you slice again:

```
getitem(tuple, slice) -> ???
```
in which case the method table impl is blind to how to handle a raw
tuple. I believe the more correct way is that slice should generate
ANOTHER container type instead of a naked python tuple.

---------

Co-authored-by: Kai-Hsin Wu <[email protected]>

Author: johnzl-777

src/bloqade/analysis/measure_id/analysis.py

@@ -1,7 +1,7 @@
 from typing import TypeVar
 from dataclasses import field, dataclass
 
-from kirin import ir, interp
+from kirin import ir
 from kirin.analysis import ForwardExtra, const
 from kirin.analysis.forward import ForwardFrame
 
@@ -37,20 +37,19 @@ def run_method(self, method: ir.Method, args: tuple[MeasureId, ...]):
         # NOTE: we do not support dynamic calls here, thus no need to propagate method object
         return self.run_callable(method.code, (self.lattice.bottom(),) + args)
 
-    T = TypeVar("T")
-
     # Xiu-zhe (Roger) Luo came up with this in the address analysis,
-    # reused here for convenience
+    # reused here for convenience (now modified to be a bit more graceful)
     # TODO: Remove this function once upgrade to kirin 0.18 happens,
     #       method is built-in to interpreter then
-    def get_const_value(self, input_type: type[T], value: ir.SSAValue) -> T:
+
+    T = TypeVar("T")
+
+    def get_const_value(
+        self, input_type: type[T], value: ir.SSAValue
+    ) -> type[T] | None:
         if isinstance(hint := value.hints.get("const"), const.Value):
             data = hint.data
             if isinstance(data, input_type):
                 return hint.data
-            raise interp.InterpreterError(
-                f"Expected constant value <type = {input_type}>, got {data}"
-            )
-        raise interp.InterpreterError(
-            f"Expected constant value <type = {input_type}>, got {value}"
-        )
+
+        return None

src/bloqade/analysis/measure_id/impls.py

@@ -103,10 +103,23 @@ class PyIndexing(interp.MethodTable):
     def getitem(
         self, interp: MeasurementIDAnalysis, frame: interp.Frame, stmt: py.GetItem
     ):
-        idx = interp.get_const_value(int, stmt.index)
+
+        idx_or_slice = interp.get_const_value((int, slice), stmt.index)
+        if idx_or_slice is None:
+            return (InvalidMeasureId(),)
+
+        # hint = stmt.index.hints.get("const")
+        # if hint is None or not isinstance(hint, const.Value):
+        #    return (InvalidMeasureId(),)
+
         obj = frame.get(stmt.obj)
         if isinstance(obj, MeasureIdTuple):
-            return (obj.data[idx],)
+            if isinstance(idx_or_slice, slice):
+                return (MeasureIdTuple(data=obj.data[idx_or_slice]),)
+            elif isinstance(idx_or_slice, int):
+                return (obj.data[idx_or_slice],)
+            else:
+                return (InvalidMeasureId(),)
         # just propagate these down the line
         elif isinstance(obj, (AnyMeasureId, NotMeasureId)):
             return (obj,)

test/analysis/measure_id/test_measure_id.py

@@ -7,6 +7,7 @@
     NotMeasureId,
     MeasureIdBool,
     MeasureIdTuple,
+    InvalidMeasureId,
 )
 
 
@@ -139,3 +140,78 @@ def test():
         NotMeasureId(),
         MeasureIdBool(idx=1),
     ]
+
+
+def test_slice():
+    @kernel
+    def test():
+        q = qubit.new(6)
+        qubit.apply(op.x(), q[2])
+
+        ms = qubit.measure(q)
+        msi = ms[1:]  # MeasureIdTuple becomes a python tuple
+        msi2 = msi[1:]  # slicing should still work on previous tuple
+        ms_final = msi2[::2]
+
+        return ms_final
+
+    frame, _ = MeasurementIDAnalysis(test.dialects).run_analysis(test)
+
+    test.print(analysis=frame.entries)
+
+    assert [frame.entries[result] for result in results_at(test, 0, 8)] == [
+        MeasureIdTuple(data=tuple(list(MeasureIdBool(idx=i) for i in range(2, 7))))
+    ]
+    assert [frame.entries[result] for result in results_at(test, 0, 10)] == [
+        MeasureIdTuple(data=tuple(list(MeasureIdBool(idx=i) for i in range(3, 7))))
+    ]
+    assert [frame.entries[result] for result in results_at(test, 0, 12)] == [
+        MeasureIdTuple(data=(MeasureIdBool(idx=3), MeasureIdBool(idx=5)))
+    ]
+
+
+def test_getitem_no_hint():
+    @kernel
+    def test(idx):
+        q = qubit.new(6)
+        ms = qubit.measure(q)
+
+        return ms[idx]
+
+    frame, _ = MeasurementIDAnalysis(test.dialects).run_analysis(test)
+
+    assert [frame.entries[result] for result in results_at(test, 0, 3)] == [
+        InvalidMeasureId(),
+    ]
+
+
+def test_getitem_invalid_hint():
+    @kernel
+    def test():
+        q = qubit.new(6)
+        ms = qubit.measure(q)
+
+        return ms["x"]
+
+    frame, _ = MeasurementIDAnalysis(test.dialects).run_analysis(test)
+
+    assert [frame.entries[result] for result in results_at(test, 0, 4)] == [
+        InvalidMeasureId()
+    ]
+
+
+def test_getitem_propagate_invalid_measure():
+
+    @kernel
+    def test():
+        q = qubit.new(6)
+        ms = qubit.measure(q)
+        # this will return an InvalidMeasureId
+        invalid_ms = ms["x"]
+        return invalid_ms[0]
+
+    frame, _ = MeasurementIDAnalysis(test.dialects).run_analysis(test)
+
+    assert [frame.entries[result] for result in results_at(test, 0, 6)] == [
+        InvalidMeasureId()
+    ]