refactor interp to clean up APIs + improve stacktrace (#369)

this PR cleans up the double `new_frame` API in the interpreter
framework (closes #363). We also refactor the `run_ssacfg_region` and
deprecate `run_block` interface to make linter happier (it was not
consistent in `AbstractInterpreter` vs `BaseInterpreter`).

We still need some further thinking on what is a value of `region` and
`block` (determine by their terminating statement/block?) so that
codegen framework can share more codebase with interpreter.

Author: Roger-luo

src/kirin/analysis/const/prop.py

@@ -52,16 +52,18 @@ def initialize(self):
         self._interp.initialize()
         return self
 
-    def new_frame(self, code: ir.Statement) -> Frame:
-        return Frame.from_func_like(code)
+    def initialize_frame(
+        self, code: ir.Statement, *, has_parent_access: bool = False
+    ) -> Frame:
+        return Frame(code, has_parent_access=has_parent_access)
 
     def try_eval_const_pure(
         self,
         frame: Frame,
         stmt: ir.Statement,
         values: tuple[Value, ...],
     ) -> interp.StatementResult[Result]:
-        _frame = self._interp.new_frame(frame.code)
+        _frame = self._interp.initialize_frame(frame.code)
         _frame.set_values(stmt.args, tuple(x.data for x in values))
         method = self._interp.lookup_registry(frame, stmt)
         if method is not None:

src/kirin/analysis/forward.py

@@ -69,7 +69,7 @@ def run_analysis(
             # so we don't need to copy the frames.
             if not no_raise:
                 raise e
-            return self.new_frame(method.code), self.lattice.bottom()
+            return self.state.current_frame, self.lattice.bottom()
         finally:
             self._eval_lock = False
             sys.setrecursionlimit(current_recursion_limit)
@@ -103,5 +103,7 @@ class Forward(ForwardExtra[ForwardFrame[LatticeElemType], LatticeElemType], ABC)
     [`ForwardExtra`][kirin.analysis.forward.ForwardExtra] instead.
     """
 
-    def new_frame(self, code: ir.Statement) -> ForwardFrame[LatticeElemType]:
-        return ForwardFrame.from_func_like(code)
+    def initialize_frame(
+        self, code: ir.Statement, *, has_parent_access: bool = False
+    ) -> ForwardFrame[LatticeElemType]:
+        return ForwardFrame(code, has_parent_access=has_parent_access)

src/kirin/dialects/cf/constprop.py

@@ -9,7 +9,7 @@ class DialectConstProp(MethodTable):
 
     @impl(Branch)
     def branch(self, interp: const.Propagate, frame: const.Frame, stmt: Branch):
-        interp.state.current_frame().worklist.append(
+        interp.state.current_frame.worklist.append(
             Successor(stmt.successor, *frame.get_values(stmt.arguments))
         )
         return ()
@@ -21,7 +21,7 @@ def conditional_branch(
         frame: const.Frame,
         stmt: ConditionalBranch,
     ):
-        frame = interp.state.current_frame()
+        frame = interp.state.current_frame
         cond = frame.get(stmt.cond)
         if isinstance(cond, const.Value):
             else_successor = Successor(

src/kirin/dialects/func/emit.py

@@ -18,12 +18,12 @@ def emit_function(
         self, interp: EmitJulia[IO_t], frame: emit.EmitStrFrame, stmt: Function
     ):
         fn_args = stmt.body.blocks[0].args[1:]
-        argnames = frame.get_values(fn_args)
+        argnames = tuple(interp.ssa_id[arg] for arg in fn_args)
         argtypes = tuple(interp.emit_attribute(x.type) for x in fn_args)
         args = [f"{name}::{type}" for name, type in zip(argnames, argtypes)]
         interp.write(f"function {stmt.sym_name}({', '.join(args)})")
         frame.indent += 1
-        interp.run_ssacfg_region(frame, stmt.body)
+        interp.run_ssacfg_region(frame, stmt.body, (stmt.sym_name,) + argnames)
         frame.indent -= 1
         interp.writeln(frame, "end")
         return ()
@@ -63,12 +63,11 @@ def emit_lambda(
         self, interp: EmitJulia[IO_t], frame: emit.EmitStrFrame, stmt: Lambda
     ):
         args = tuple(interp.ssa_id[x] for x in stmt.body.blocks[0].args[1:])
-        frame.set_values(stmt.body.blocks[0].args, args)
-        frame.set_values((stmt.body.blocks[0].args[0],), (stmt.sym_name,))
+        frame.set_values(stmt.body.blocks[0].args, (stmt.sym_name,) + args)
         frame.captured[stmt.body.blocks[0].args[0]] = frame.get_values(stmt.captured)
         interp.writeln(frame, f"function {stmt.sym_name}({', '.join(args[1:])})")
         frame.indent += 1
-        interp.run_ssacfg_region(frame, stmt.body)
+        interp.run_ssacfg_region(frame, stmt.body, args)
         frame.indent -= 1
         interp.writeln(frame, "end")
         return (stmt.sym_name,)

src/kirin/dialects/scf/absint.py

@@ -56,9 +56,7 @@ def _infer_if_else_cond(
             frame.worklist.append(interp.Successor(body_block, frame.get(stmt.cond)))
             return
 
-        with interp_.state.new_frame(interp_.new_frame(stmt)) as body_frame:
-            body_frame.entries.update(frame.entries)
-            body_frame.set(body_block.args[0], frame.get(stmt.cond))
-            ret = interp_.run_ssacfg_region(body_frame, body)
+        with interp_.new_frame(stmt, has_parent_access=True) as body_frame:
+            ret = interp_.run_ssacfg_region(body_frame, body, (frame.get(stmt.cond),))
             frame.entries.update(body_frame.entries)
-            return ret
+        return ret

src/kirin/dialects/scf/constprop.py

@@ -83,11 +83,9 @@ def _prop_const_cond_ifelse(
         cond: const.Value,
         body: ir.Region,
     ):
-        with interp_.state.new_frame(interp_.new_frame(stmt)) as body_frame:
-            body_frame.entries.update(frame.entries)
-            body_frame.set(body.blocks[0].args[0], cond)
-            results = interp_.run_ssacfg_region(body_frame, body)
-        return body_frame, results
+        with interp_.new_frame(stmt, has_parent_access=True) as body_frame:
+            results = interp_.run_ssacfg_region(body_frame, body, (cond,))
+            return body_frame, results
 
     @interp.impl(For)
     def for_loop(
@@ -116,17 +114,12 @@ def _prop_const_iterable_forloop(
             )
 
         loop_vars = frame.get_values(stmt.initializers)
-        body_block = stmt.body.blocks[0]
-        block_args = body_block.args
 
         for value in iterable.data:
-            with interp_.state.new_frame(interp_.new_frame(stmt)) as body_frame:
-                body_frame.entries.update(frame.entries)
-                body_frame.set_values(
-                    block_args,
-                    (const.Value(value),) + loop_vars,
+            with interp_.new_frame(stmt, has_parent_access=True) as body_frame:
+                loop_vars = interp_.run_ssacfg_region(
+                    body_frame, stmt.body, (const.Value(value),) + loop_vars
                 )
-                loop_vars = interp_.run_ssacfg_region(body_frame, stmt.body)
 
             if body_frame.frame_is_not_pure:
                 frame_is_not_pure = True

src/kirin/dialects/scf/interp.py

@@ -18,16 +18,16 @@ def if_else(self, interp_: interp.Interpreter, frame: interp.Frame, stmt: IfElse
             body = stmt.then_body
         else:
             body = stmt.else_body
-        return interp_.run_ssacfg_region(frame, body)
+        return interp_.run_ssacfg_region(frame, body, (cond,))
 
     @interp.impl(For)
     def for_loop(self, interpreter: interp.Interpreter, frame: interp.Frame, stmt: For):
         iterable = frame.get(stmt.iterable)
         loop_vars = frame.get_values(stmt.initializers)
-        block_args = stmt.body.blocks[0].args
         for value in iterable:
-            frame.set_values(block_args, (value,) + loop_vars)
-            loop_vars = interpreter.run_ssacfg_region(frame, stmt.body)
+            loop_vars = interpreter.run_ssacfg_region(
+                frame, stmt.body, (value,) + loop_vars
+            )
             if isinstance(loop_vars, interp.ReturnValue):
                 return loop_vars
             elif loop_vars is None:

src/kirin/dialects/scf/typeinfer.py

@@ -45,9 +45,10 @@ def for_loop(
             frame.worklist.append(interp.Successor(body_block, item, *loop_vars))
             return  # if terminate is Return, there is no result
 
-        with interp_.state.new_frame(interp_.new_frame(stmt)) as body_frame:
-            body_frame.entries.update(frame.entries)
-            loop_vars_ = interp_.run_ssacfg_region(body_frame, stmt.body)
+        with interp_.new_frame(stmt, has_parent_access=True) as body_frame:
+            loop_vars_ = interp_.run_ssacfg_region(
+                body_frame, stmt.body, (iterable,) + loop_vars
+            )
 
         frame.entries.update(body_frame.entries)
         if isinstance(loop_vars_, interp.ReturnValue):

src/kirin/emit/abc.py

@@ -21,7 +21,11 @@ class EmitFrame(interp.Frame[ValueType]):
 class EmitABC(interp.BaseInterpreter[FrameType, ValueType], ABC):
 
     def run_callable_region(
-        self, frame: FrameType, code: ir.Statement, region: ir.Region
+        self,
+        frame: FrameType,
+        code: ir.Statement,
+        region: ir.Region,
+        args: tuple[ValueType, ...],
     ) -> ValueType:
         results = self.eval_stmt(frame, code)
         if isinstance(results, tuple):
@@ -32,12 +36,11 @@ def run_callable_region(
         raise interp.InterpreterError(f"Unexpected results {results}")
 
     def run_ssacfg_region(
-        self, frame: FrameType, region: ir.Region
+        self, frame: FrameType, region: ir.Region, args: tuple[ValueType, ...]
     ) -> tuple[ValueType, ...]:
-        frame.worklist.append(
-            interp.Successor(region.blocks[0], frame.get_values(region.blocks[0].args))
-        )
+        frame.worklist.append(interp.Successor(region.blocks[0], *args))
         while (succ := frame.worklist.pop()) is not None:
+            frame.set_values(succ.block.args, succ.block_args)
             block_header = self.emit_block(frame, succ.block)
             frame.block_ref[succ.block] = block_header
         return ()

src/kirin/emit/abc.pyi

@@ -15,10 +15,14 @@ FrameType = TypeVar("FrameType", bound=EmitFrame)
 
 class EmitABC(interp.BaseInterpreter[FrameType, ValueType]):
     def run_callable_region(
-        self, frame: FrameType, code: ir.Statement, region: ir.Region
+        self,
+        frame: FrameType,
+        code: ir.Statement,
+        region: ir.Region,
+        args: tuple[ValueType, ...],
     ) -> ValueType: ...
     def run_ssacfg_region(
-        self, frame: FrameType, region: ir.Region
+        self, frame: FrameType, region: ir.Region, args: tuple[ValueType, ...]
     ) -> tuple[ValueType, ...]: ...
     def emit_attribute(self, attr: ir.Attribute) -> ValueType: ...
     def emit_type_Any(self, attr: types.AnyType) -> ValueType: ...