add alias support to measure_id analysis (#412)

@ehua7365 encountered a problem where reassigning an IList of measure
results would cause `SquinToStim` to fail:

```
        q = qubit.new(4)
        ms = qubit.measure(q)
        new_ms = ms

        if new_ms[0]:
            qubit.apply(op.z(), q[0])

```
(This is an intentionally contrived example to keep things short).
Originally I had directed Eric to include kirin's InlineAlias and
InlineGetItem rules which also got rid of the issue until I took a
closer look and realized the problem is because when an alias happens
with an IList of measure results, the analysis does not propagate the
corresponding data.

As a result, further Stim-related rewrites have problems because the
MeasureIDAnalysis data is not available for the subsequent SSA Values.

---------

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

Author: johnzl-777

src/bloqade/analysis/measure_id/impls.py

@@ -114,6 +114,15 @@ def getitem(
             return (InvalidMeasureId(),)
 
 
+@py.assign.dialect.register(key="measure_id")
+class PyAssign(interp.MethodTable):
+    @interp.impl(py.Alias)
+    def alias(
+        self, interp: MeasurementIDAnalysis, frame: interp.Frame, stmt: py.assign.Alias
+    ):
+        return (frame.get(stmt.value),)
+
+
 @py.binop.dialect.register(key="measure_id")
 class PyBinOp(interp.MethodTable):
     @interp.impl(py.Add)

src/bloqade/stim/passes/squin_to_stim.py

@@ -16,6 +16,7 @@
 from kirin.passes.abc import Pass
 from kirin.rewrite.abc import RewriteResult
 from kirin.passes.inline import InlinePass
+from kirin.rewrite.alias import InlineAlias
 
 from bloqade.stim.rewrite import (
     SquinWireToStim,
@@ -89,6 +90,9 @@ def unsafe_run(self, mt: Method) -> RewriteResult:
         rewrite_result = (
             Walk(Fixpoint(CFGCompactify())).rewrite(mt.code).join(rewrite_result)
         )
+
+        Walk(InlineAlias()).rewrite(mt.code).join(rewrite_result)
+
         rewrite_result = (
             StimSimplifyIfs(mt.dialects, no_raise=self.no_raise)
             .unsafe_run(mt)

test/analysis/measure_id/test_measure_id.py

@@ -39,6 +39,37 @@ def test():
     assert measure_id_tuples[-1] == expected_measure_id_tuple
 
 
+def test_measure_alias():
+
+    @kernel
+    def test():
+        ql = qubit.new(5)
+        ml = qubit.measure(ql)
+        ml_alias = ml
+
+        return ml_alias
+
+    frame, _ = MeasurementIDAnalysis(test.dialects).run_analysis(test)
+
+    test.print(analysis=frame.entries)
+
+    # Collect MeasureIdTuples
+    measure_id_tuples = [
+        value for value in frame.entries.values() if isinstance(value, MeasureIdTuple)
+    ]
+
+    # construct expected MeasureIdTuple
+    expected_measure_id_tuple = MeasureIdTuple(
+        data=tuple([MeasureIdBool(idx=i) for i in range(1, 6)])
+    )
+
+    assert len(measure_id_tuples) == 2
+    assert all(
+        measure_id_tuple == expected_measure_id_tuple
+        for measure_id_tuple in measure_id_tuples
+    )
+
+
 def test_measure_count_at_if_else():
 
     @kernel

test/stim/passes/stim_reference_programs/qubit/alias_with_measure_list.stim

@@ -0,0 +1,3 @@
+
+MZ(0.00000000) 0 1 2 3
+CZ rec[-4] 0

test/stim/passes/test_squin_meas_to_stim.py

@@ -53,6 +53,25 @@ def main():
     assert base_stim_prog == codegen(main)
 
 
+def test_alias_with_measure_list():
+
+    @squin.kernel
+    def main():
+
+        q = qubit.new(4)
+        ms = qubit.measure(q)
+        new_ms = ms
+
+        if new_ms[0]:
+            qubit.apply(op.z(), q[0])
+
+    SquinToStimPass(main.dialects)(main)
+
+    base_stim_prog = load_reference_program("alias_with_measure_list.stim")
+
+    assert base_stim_prog == codegen(main)
+
+
 def test_record_index_order():
 
     @squin.kernel