Fix address analysis

Author: david-pl

src/bloqade/analysis/address/analysis.py

@@ -119,6 +119,7 @@ def run_lattice(
             case ConstResult(const.Value(ir.Method() as method)):
                 _, ret = self.call(
                     method.code,
+                    self.method_self(method),
                     *inputs,
                     # **kwargs,
                     # **{k: v for k, v in zip(kwargs, frame.get_values(stmt.kwargs))},

src/bloqade/analysis/address/impls.py

@@ -180,13 +180,10 @@ def invoke(
         frame: ForwardFrame[Address],
         stmt: func.Invoke,
     ):
-
-        args = interp_.permute_values(
-            stmt.callee.arg_names, frame.get_values(stmt.inputs), stmt.kwargs
-        )
-        _, ret = interp_.run_method(
-            stmt.callee,
-            args,
+        _, ret = interp_.call(
+            stmt.callee.code,
+            interp_.method_self(stmt.callee),
+            *frame.get_values(stmt.inputs),
         )
 
         return (ret,)
@@ -319,26 +316,28 @@ def ifelse(
         ):
             body = stmt.then_body if const_cond.data else stmt.else_body
             with interp_.new_frame(stmt, has_parent_access=True) as body_frame:
-                ret = interp_.run_ssacfg_region(body_frame, body, (address_cond,))
+                ret = interp_.frame_call_region(body_frame, stmt, body, address_cond)
                 # interp_.set_values(frame, body_frame.entries.keys(), body_frame.entries.values())
                 return ret
         else:
             # run both branches
             with interp_.new_frame(stmt, has_parent_access=True) as then_frame:
-                then_results = interp_.run_ssacfg_region(
-                    then_frame, stmt.then_body, (address_cond,)
-                )
-                interp_.set_values(
-                    frame, then_frame.entries.keys(), then_frame.entries.values()
+                then_results = interp_.frame_call_region(
+                    then_frame,
+                    stmt,
+                    stmt.then_body,
+                    address_cond,
                 )
+                frame.set_values(then_frame.entries.keys(), then_frame.entries.values())
 
             with interp_.new_frame(stmt, has_parent_access=True) as else_frame:
-                else_results = interp_.run_ssacfg_region(
-                    else_frame, stmt.else_body, (address_cond,)
-                )
-                interp_.set_values(
-                    frame, else_frame.entries.keys(), else_frame.entries.values()
+                else_results = interp_.frame_call_region(
+                    else_frame,
+                    stmt,
+                    stmt.else_body,
+                    address_cond,
                 )
+                frame.set_values(else_frame.entries.keys(), else_frame.entries.values())
             # TODO: pick the non-return value
             if isinstance(then_results, interp.ReturnValue) and isinstance(
                 else_results, interp.ReturnValue
@@ -364,12 +363,12 @@ def for_loop(
         iter_type, iterable = interp_.unpack_iterable(frame.get(stmt.iterable))
 
         if iter_type is None:
-            return interp_.eval_stmt_fallback(frame, stmt)
+            return interp_.eval_fallback(frame, stmt)
 
         for value in iterable:
             with interp_.new_frame(stmt, has_parent_access=True) as body_frame:
-                loop_vars = interp_.run_ssacfg_region(
-                    body_frame, stmt.body, (value,) + loop_vars
+                loop_vars = interp_.frame_call_region(
+                    body_frame, stmt, stmt.body, value, *loop_vars
                 )
 
             if loop_vars is None:

test/analysis/address/test_qubit_analysis.py

@@ -21,7 +21,7 @@ def test():
         return (q1[1], q2)
 
     address_analysis = address.AddressAnalysis(test.dialects)
-    frame, _ = address_analysis.run_analysis(test, no_raise=False)
+    frame, _ = address_analysis.run(test)
     address_types = collect_address_types(frame, address.PartialTuple)
 
     test.print(analysis=frame.entries)
@@ -116,7 +116,7 @@ def main():
         return q
 
     address_analysis = address.AddressAnalysis(main.dialects)
-    address_analysis.run_analysis(main, no_raise=False)
+    address_analysis.run(main)
 
 
 def test_new_qubit():
@@ -125,7 +125,7 @@ def main():
         return squin.qubit.new()
 
     address_analysis = address.AddressAnalysis(main.dialects)
-    _, result = address_analysis.run_analysis(main, no_raise=False)
+    _, result = address_analysis.run(main)
     assert result == address.AddressQubit(0)
 
 
@@ -139,8 +139,9 @@ def main(n: int):
         return qreg
 
     address_analysis = address.AddressAnalysis(main.dialects)
-    frame, result = address_analysis.run_analysis(
-        main, args=(address.ConstResult(const.Unknown()),), no_raise=False
+    frame, result = address_analysis.run(
+        main,
+        address.ConstResult(const.Unknown()),
     )
     assert result == address.AddressReg(data=tuple(range(4)))
 
@@ -155,7 +156,7 @@ def main(n: int):
         return qreg
 
     address_analysis = address.AddressAnalysis(main.dialects)
-    frame, result = address_analysis.run_analysis(main, no_raise=False)
+    frame, result = address_analysis.run(main)
     assert result == address.AddressReg(data=tuple(range(4)))
 
 
@@ -165,7 +166,7 @@ def main(n: int):
         return (0, 1) + (2, n)
 
     address_analysis = address.AddressAnalysis(main.dialects)
-    frame, result = address_analysis.run_analysis(main, no_raise=False)
+    frame, result = address_analysis.run(main)
 
     assert result == address.PartialTuple(
         data=(
@@ -183,7 +184,7 @@ def main(n: int):
         return (0, 1) + [2, n]  # type: ignore
 
     address_analysis = address.AddressAnalysis(main.dialects)
-    frame, result = address_analysis.run_analysis(main, no_raise=False)
+    frame, result = address_analysis.run(main)
 
     assert result == address.Bottom()
 
@@ -194,7 +195,7 @@ def main(n: tuple[int, ...]):
         return (0, 1) + n
 
     address_analysis = address.AddressAnalysis(main.dialects)
-    frame, result = address_analysis.run_analysis(main, no_raise=False)
+    frame, result = address_analysis.run(main)
 
     assert result == address.Unknown()
 
@@ -207,7 +208,7 @@ def main(q: qubit.Qubit):
         return (0, q, 2, q)[1::2]
 
     address_analysis = address.AddressAnalysis(main.dialects)
-    frame, result = address_analysis.run_analysis(main, no_raise=False)
+    frame, result = address_analysis.run(main)
     assert result == address.UnknownReg()
 
 
@@ -219,7 +220,7 @@ def main(n: int):
 
     main.print()
     address_analysis = address.AddressAnalysis(main.dialects)
-    frame, result = address_analysis.run_analysis(main, no_raise=False)
+    frame, result = address_analysis.run(main)
     main.print(analysis=frame.entries)
     assert (
         result == address.UnknownReg()
@@ -260,7 +261,7 @@ def main():
 
     func = main
     analysis = address.AddressAnalysis(squin.kernel)
-    _, ret = analysis.run_analysis(func, no_raise=False)
+    _, ret = analysis.run(func)
 
     assert ret == address.AddressReg(data=tuple(range(20)))
     assert analysis.qubit_count == 20