Fix UnrollIfs pass for non-empty else bodies (#417)

Closes #284 and #286.

These issues were blocked by another bug in kirin that is fixed now.

Also, note that this change basically just keeps the rewrite from
raising an error, but ultimately things like non-empty else bodies
(including `elif`) is just not supported in QASM2, so you will not be
able to emit QASM2 from such a program.

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

Author: david-pl

src/bloqade/rewrite/rules/split_ifs.py

@@ -46,9 +46,13 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
         if not isinstance(node, scf.IfElse):
             return RewriteResult()
 
+        # NOTE: only empty else bodies are allowed in valid QASM2
+        if not self._has_empty_else(node):
+            return RewriteResult()
+
         *stmts, yield_or_return = node.then_body.stmts()
 
-        if len(stmts) == 1:
+        if len(stmts) <= 1:
             return RewriteResult()
 
         is_yield = isinstance(yield_or_return, scf.Yield)
@@ -71,3 +75,16 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
         node.delete()
 
         return RewriteResult(has_done_something=True)
+
+    def _has_empty_else(self, node: scf.IfElse) -> bool:
+        else_stmts = list(node.else_body.stmts())
+        if len(else_stmts) > 1:
+            return False
+
+        if len(else_stmts) == 0:
+            return True
+
+        if not isinstance(else_stmts[0], scf.Yield):
+            return False
+
+        return len(else_stmts[0].values) == 0

test/qasm2/passes/test_unroll_if.py

@@ -1,5 +1,6 @@
 from bloqade import qasm2
 from bloqade.qasm2.emit import QASM2
+from bloqade.qasm2.passes import QASM2Fold
 
 
 def test_unrolling_ifs():
@@ -103,3 +104,28 @@ def main():
     ast = target.emit(main)
 
     qasm2.parse.pprint(ast)
+
+
+def test_elif():
+
+    @qasm2.extended
+    def main():
+        q = qasm2.qreg(1)
+        c = qasm2.creg(1)
+        qasm2.h(q[0])
+        qasm2.measure(q, c)
+
+        parity = 0
+        if c[0] == 1 and parity == 0:
+            qasm2.x(q[0])
+            parity = 0
+        elif c[0] == 0:
+            parity = 1
+        elif c[0] == 2:
+            parity = 2
+
+        return parity
+
+    QASM2Fold(qasm2.extended, unroll_ifs=True, no_raise=False).fixpoint(main)
+
+    main.print()