support stacktrace in interpreter (#374)

```python
from kirin.prelude import basic

@basic(typeinfer=True)
def some_code_will_error(x):
    """
    This function will raise an error.
    """
    return 1 / x

@basic(typeinfer=True)
def some_func():
    return some_code_will_error(0)

some_func()
```

<img width="839" alt="image"
src="https://github.com/user-attachments/assets/031700a6-87f7-4d72-bc54-a11871c66e64"
/>

Author: Roger-luo

src/kirin/analysis/forward.py

@@ -64,7 +64,7 @@ def run_analysis(
         sys.setrecursionlimit(self.max_python_recursion_depth)
         try:
             frame, ret = self.run_method(method, args)
-        except interp.InterpreterError as e:
+        except Exception as e:
             # NOTE: initialize will create new State
             # so we don't need to copy the frames.
             if not no_raise:

src/kirin/exception.py

@@ -2,6 +2,7 @@
 for the Kirin-based compilers.
 """
 
+import abc
 import sys
 import types
 
@@ -12,6 +13,12 @@ class NoPythonStackTrace(Exception):
     pass
 
 
+class CustomStackTrace(Exception):
+
+    @abc.abstractmethod
+    def print_stacktrace(self) -> None: ...
+
+
 def enable_stracetrace():
     """Enable the stacktrace for all exceptions."""
     global stacktrace
@@ -30,6 +37,11 @@ def exception_handler(exc_type, exc_value, exc_tb: types.TracebackType):
         print(exc_value, file=sys.stderr)
         return
 
+    if not stacktrace and issubclass(exc_type, CustomStackTrace):
+        # Handle custom stack trace exceptions
+        exc_value.print_stacktrace()
+        return
+
     # Call the default exception handler
     sys.__excepthook__(exc_type, exc_value, exc_tb)
 

src/kirin/interp/__init__.py

@@ -34,6 +34,7 @@
 from .concrete import Interpreter as Interpreter
 from .exceptions import (
     WrapException as WrapException,
+    IntepreterExit as IntepreterExit,
     InterpreterError as InterpreterError,
     FuelExhaustedError as FuelExhaustedError,
 )

src/kirin/interp/base.py

@@ -23,8 +23,7 @@
 from .frame import FrameABC
 from .state import InterpreterState
 from .value import Successor, ReturnValue, SpecialValue, StatementResult
-from .result import Ok, Err, Result
-from .exceptions import InterpreterError
+from .exceptions import IntepreterExit, InterpreterError
 
 if TYPE_CHECKING:
     from kirin.registry import StatementImpl, InterpreterRegistry
@@ -132,7 +131,7 @@ def run(
         mt: Method,
         args: tuple[ValueType, ...],
         kwargs: dict[str, ValueType] | None = None,
-    ) -> Result[ValueType]:
+    ) -> ValueType:
         """Run a method. This is the main entry point of the interpreter.
 
         Args:
@@ -155,14 +154,14 @@ def run(
         args = self.get_args(mt.arg_names[len(args) + 1 :], args, kwargs)
         try:
             _, results = self.run_method(mt, args)
-        except InterpreterError as e:
+        except Exception as e:
             # NOTE: initialize will create new State
             # so we don't need to copy the frames.
-            return Err(e, self.state)
+            raise IntepreterExit(e, self.state) from e
         finally:
             self._eval_lock = False
             sys.setrecursionlimit(current_recursion_limit)
-        return Ok(results)
+        return results
 
     def run_stmt(
         self, stmt: Statement, args: tuple[ValueType, ...]

src/kirin/interp/exceptions.py

@@ -1,5 +1,15 @@
+from __future__ import annotations
+
+import sys
+from typing import TYPE_CHECKING
 from dataclasses import dataclass
 
+from kirin.exception import CustomStackTrace
+
+if TYPE_CHECKING:
+    from .frame import FrameABC
+    from .state import InterpreterState
+
 
 # errors
 class InterpreterError(Exception):
@@ -20,6 +30,27 @@ class WrapException(InterpreterError):
     exception: Exception
 
 
+@dataclass
+class IntepreterExit(CustomStackTrace):
+    exception: Exception
+    state: InterpreterState
+
+    def print_stacktrace(self) -> None:
+        """Print the stacktrace of the interpreter."""
+        frame: FrameABC | None = self.state.current_frame
+        print(f"{type(self.exception).__name__}: {self.exception}", file=sys.stderr)
+        print("Traceback (most recent call last):", file=sys.stderr)
+        frames: list[FrameABC] = []
+        while frame is not None:
+            frames.append(frame)
+            frame = frame.parent
+        frames.reverse()
+        for frame in frames:
+            if stmt := frame.current_stmt:
+                print("  " + repr(stmt.source), file=sys.stderr)
+                print("     " + stmt.print_str(end=""), file=sys.stderr)
+
+
 class FuelExhaustedError(InterpreterError):
     """An error raised when the interpreter runs out of fuel."""
 

src/kirin/interp/frame.py

@@ -34,6 +34,8 @@ class FrameABC(ABC, Generic[ValueType]):
     has_parent_access: bool = field(default=False, kw_only=True, compare=True)
     """If we have access to the entries of the parent frame."""
 
+    lineno_offset: int = field(default=0, kw_only=True, compare=True)
+
     current_stmt: Statement | None = field(
         default=None, init=False, compare=False, repr=False
     )

src/kirin/ir/group.py

@@ -190,15 +190,15 @@ def __call__(
         Returns:
             Method: the method created from the python function.
         """
+        frame = inspect.currentframe()
 
         def wrapper(py_func: Callable) -> Method:
             if py_func.__name__ == "<lambda>":
                 raise ValueError("Cannot compile lambda functions")
 
             lineno_offset, file = 0, ""
-            frame = inspect.currentframe()
-            if frame and frame.f_back is not None and frame.f_back.f_back is not None:
-                call_site_frame = frame.f_back.f_back
+            if frame and frame.f_back is not None:
+                call_site_frame = frame.f_back
                 if py_func.__name__ in call_site_frame.f_locals:
                     raise CompilerError(
                         f"overwriting function definition of `{py_func.__name__}`"

src/kirin/ir/method.py

@@ -50,7 +50,7 @@ def __call__(self, *args: Param.args, **kwargs: Param.kwargs) -> RetType:
             raise ValueError("Incorrect number of arguments")
         # NOTE: multi-return values will be wrapped in a tuple for Python
         interp = Interpreter(self.dialects)
-        return interp.run(self, args=args, kwargs=kwargs).expect()
+        return interp.run(self, args=args, kwargs=kwargs)
 
     @property
     def args(self):
@@ -141,7 +141,6 @@ def __postprocess_validation_error(self, e: ValidationError):
                 source.splitlines(),
                 e,
                 file=self.file,
-                lineno_offset=self.lineno_offset,
             )
         else:
             msg += "\nNo source available"

src/kirin/lowering/python/lowering.py

@@ -112,7 +112,7 @@ def run(
         source = source or ast.unparse(stmt)
         state = State(
             self,
-            source=SourceInfo.from_ast(stmt, lineno_offset, col_offset),
+            source=SourceInfo.from_ast(stmt, lineno_offset, col_offset, file),
             file=file,
             lines=source.splitlines(),
             lineno_offset=lineno_offset,
@@ -156,7 +156,7 @@ def lower_global(self, state: State[ast.AST], node: ast.AST) -> LoweringABC.Resu
     def visit(self, state: State[ast.AST], node: ast.AST) -> Result:
         if hasattr(node, "lineno"):
             state.source = SourceInfo.from_ast(
-                node, state.lineno_offset, state.col_offset
+                node, state.lineno_offset, state.col_offset, state.file
             )
         name = node.__class__.__name__
         if name in self.registry.ast_table:
@@ -169,7 +169,7 @@ def generic_visit(self, state: State[ast.AST], node: ast.AST) -> Result:
     def visit_Call(self, state: State[ast.AST], node: ast.Call) -> Result:
         if hasattr(node.func, "lineno"):
             state.source = SourceInfo.from_ast(
-                node.func, state.lineno_offset, state.col_offset
+                node.func, state.lineno_offset, state.col_offset, state.file
             )
 
         global_callee_result = state.get_global(node.func, no_raise=True)

src/kirin/lowering/state.py

@@ -228,5 +228,4 @@ def error_hint(
             indent=indent,
             show_lineno=show_lineno,
             max_lines=max_lines,
-            lineno_offset=self.lineno_offset,
         )