fix nested list handling in SquinToStimPass (#405)

Fixes #404 and also caught a mistake I introduced which is that
AggressiveUnroll does not `join` the rewrite_result chain.

---------

Co-authored-by: David Plankensteiner <[email protected]>

Author: johnzl-777

src/bloqade/stim/passes/simplify_ifs.py

@@ -9,6 +9,7 @@
     ConstantFold,
     CommonSubexpressionElimination,
 )
+from kirin.dialects.ilist.passes import ConstList2IList
 
 from ..rewrite.ifs_to_stim import StimLiftThenBody, StimSplitIfStmts
 
@@ -23,8 +24,16 @@ def unsafe_run(self, mt: ir.Method):
             Walk(StimSplitIfStmts()),
         ).rewrite(mt.code)
 
+        # because nested python lists don't have their
+        # member lists converted to ILists, ConstantFold
+        # can add python lists that can't be hashed, causing
+        # issues with CSE. ConstList2IList remedies that problem here.
         result = (
-            Fixpoint(Walk(Chain(ConstantFold(), CommonSubexpressionElimination())))
+            Chain(
+                Fixpoint(Walk(ConstantFold())),
+                Walk(ConstList2IList()),
+                Walk(CommonSubexpressionElimination()),
+            )
             .rewrite(mt.code)
             .join(result)
         )

src/bloqade/stim/passes/squin_to_stim.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 
-from kirin.passes import Fold, HintConst
+from kirin.passes import Fold, HintConst, TypeInfer
 from kirin.rewrite import (
     Walk,
     Chain,
@@ -33,6 +33,7 @@
 from bloqade.rewrite.passes import CanonicalizeIList
 from bloqade.analysis.address import AddressAnalysis
 from bloqade.analysis.measure_id import MeasurementIDAnalysis
+from bloqade.squin.rewrite.desugar import ApplyDesugarRule
 
 from .simplify_ifs import StimSimplifyIfs
 from ..rewrite.ifs_to_stim import IfToStim
@@ -79,9 +80,11 @@ def unsafe_run(self, mt: Method) -> RewriteResult:
             dialects=mt.dialects, no_raise=self.no_raise
         ).unsafe_run(mt)
 
-        rewrite_result = AggressiveForLoopUnroll(
-            dialects=mt.dialects, no_raise=self.no_raise
-        ).fixpoint(mt)
+        rewrite_result = (
+            AggressiveForLoopUnroll(dialects=mt.dialects, no_raise=self.no_raise)
+            .fixpoint(mt)
+            .join(rewrite_result)
+        )
 
         rewrite_result = (
             Walk(Fixpoint(CFGCompactify())).rewrite(mt.code).join(rewrite_result)
@@ -105,6 +108,9 @@ def unsafe_run(self, mt: Method) -> RewriteResult:
             .join(rewrite_result)
         )
 
+        TypeInfer(dialects=mt.dialects, no_raise=self.no_raise).unsafe_run(mt)
+        Walk(ApplyDesugarRule()).rewrite(mt.code)
+
         # after this the program should be in a state where it is analyzable
         # -------------------------------------------------------------------
 

test/stim/passes/stim_reference_programs/qubit/nested_list.stim

@@ -0,0 +1,3 @@
+
+H 0
+H 2

test/stim/passes/test_squin_qubit_to_stim.py

@@ -161,7 +161,7 @@ def main():
         squin.qubit.apply(squin.op.h(), q[0])
         cx = squin.op.cx()
         for i in range(2):
-            squin.qubit.apply(cx, [q[i], q[i + 1]])
+            squin.qubit.apply(cx, q[i], q[i + 1])
 
     SquinToStimPass(main.dialects)(main)
     base_stim_prog = load_reference_program("for_loop_nontrivial_index.stim")
@@ -178,9 +178,34 @@ def main():
         cx = squin.op.cx()
         for i in range(2):
             for j in range(2, 3):
-                squin.qubit.apply(cx, [q[i], q[j]])
+                squin.qubit.apply(cx, q[i], q[j])
 
     SquinToStimPass(main.dialects)(main)
     base_stim_prog = load_reference_program("nested_for_loop.stim")
 
     assert codegen(main) == base_stim_prog.rstrip()
+
+
+def test_nested_list():
+
+    # NOTE: While SquinToStim now has the ability to handle
+    # the nested list outside of the kernel in this test,
+    # in general it will be necessary to explicitly
+    # annotate it as an IList so type inference can work
+    # properly. Otherwise its global, mutable nature means
+    # we cannot assume a static type.
+
+    pairs = [[0, 1], [2, 3]]
+
+    @squin.kernel
+    def main():
+        q = qubit.new(10)
+        h = squin.op.h()
+        for i in range(2):
+            squin.qubit.apply(h, q[pairs[i][0]])
+
+    SquinToStimPass(main.dialects)(main)
+
+    base_stim_prog = load_reference_program("nested_list.stim")
+
+    assert codegen(main) == base_stim_prog.rstrip()