rework error handling in compiler (#377)

this PR should be considered non-breaking and let's try to backport it
to 0.16 and 0.17 because it recovers the original error behaviour when
an error was raised inside the validation or interpreter.

Author: Roger-luo

src/kirin/dialects/py/assertion.py

@@ -66,9 +66,9 @@ def assert_stmt(
             return ()
 
         if stmt.message:
-            raise interp.WrapException(AssertionError(frame.get(stmt.message)))
+            raise AssertionError(frame.get(stmt.message))
         else:
-            raise interp.WrapException(AssertionError("Assertion failed"))
+            raise AssertionError("Assertion failed")
 
 
 @dialect.register(key="typeinfer")

src/kirin/dialects/py/assign.py

@@ -99,9 +99,7 @@ def type_assert(self, interp_, frame: interp.Frame, stmt: TypeAssert):
         got = frame.get(stmt.got)
         got_type = types.PyClass(type(got))
         if not got_type.is_subseteq(stmt.expected):
-            raise interp.WrapException(
-                TypeError(f"Expected {stmt.expected}, got {got_type}")
-            )
+            raise TypeError(f"Expected {stmt.expected}, got {got_type}")
         return (frame.get(stmt.got),)
 
 

src/kirin/exception.py

@@ -2,44 +2,176 @@
 for the Kirin-based compilers.
 """
 
-import abc
+from __future__ import annotations
+
+import os
 import sys
+import math
 import types
-
-stacktrace = False
-
-
-class NoPythonStackTrace(Exception):
-    pass
-
-
-class CustomStackTrace(Exception):
-
-    @abc.abstractmethod
-    def print_stacktrace(self) -> None: ...
+import shutil
+import textwrap
+from typing import TYPE_CHECKING
+
+from rich.console import Console
+
+if TYPE_CHECKING:
+    from kirin import interp
+    from kirin.source import SourceInfo
+
+KIRIN_INTERP_STATE = "__kirin_interp_state"
+KIRIN_PYTHON_STACKTRACE = os.environ.get("KIRIN_PYTHON_STACKTRACE", "0") == "1"
+KIRIN_STATIC_CHECK_LINENO = os.environ.get("KIRIN_STATIC_CHECK_LINENO", "1") == "1"
+KIRIN_STATIC_CHECK_INDENT = int(os.environ.get("KIRIN_STATIC_CHECK_INDENT", "2"))
+KIRIN_STATIC_CHECK_MAX_LINES = int(os.environ.get("KIRIN_STATIC_CHECK_MAX_LINES", "3"))
+
+
+class StaticCheckError(Exception):
+    def __init__(self, *messages: str, help: str | None = None) -> None:
+        super().__init__(*messages)
+        self.help: str | None = help
+        self.source: SourceInfo | None = None
+        self.lines: list[str] | None = None
+        self.indent: int = KIRIN_STATIC_CHECK_INDENT
+        self.max_lines: int = KIRIN_STATIC_CHECK_MAX_LINES
+        self.show_lineno: bool = KIRIN_STATIC_CHECK_LINENO
+
+    def hint(self):
+        help = self.help or ""
+        source = self.source or SourceInfo(0, 0, 0, 0)
+        lines = self.lines or []
+        begin = max(0, source.lineno - self.max_lines)
+        end = max(
+            max(source.lineno + self.max_lines, source.end_lineno or 1),
+            1,
+        )
+        end = min(len(lines), end)  # make sure end is within bounds
+        lines = lines[begin:end]
+        error_lineno = source.lineno + source.lineno_begin
+        error_lineno_len = len(str(error_lineno))
+        code_indent = min(map(self.__get_indent, lines), default=0)
+
+        console = Console(force_terminal=True)
+        with console.capture() as capture:
+            console.print(
+                f"  {source or 'stdin'}",
+                markup=True,
+                highlight=False,
+            )
+            for lineno, line in enumerate(lines, begin):
+                line = " " * self.indent + line[code_indent:]
+                if self.show_lineno:
+                    if lineno + 1 == source.lineno:
+                        line = f"{error_lineno}[dim]│[/dim]" + line
+                    else:
+                        line = "[dim]" + " " * (error_lineno_len) + "│[/dim]" + line
+                console.print("  " + line, markup=True, highlight=False)
+                if lineno + 1 == source.lineno:
+                    console.print(
+                        "  "
+                        + self.__arrow(
+                            source,
+                            code_indent,
+                            error_lineno_len,
+                            help,
+                            self.indent,
+                            self.show_lineno,
+                        ),
+                        markup=True,
+                        highlight=False,
+                    )
+        return capture.get()
+
+    def __arrow(
+        self,
+        source: SourceInfo,
+        code_indent: int,
+        error_lineno_len: int,
+        help,
+        indent: int,
+        show_lineno: bool,
+    ) -> str:
+        ret = " " * (source.col_offset - code_indent)
+        if source.end_col_offset:
+            ret += "^" * (source.end_col_offset - source.col_offset)
+        else:
+            ret += "^"
+
+        ret = " " * indent + "[red]" + ret
+        if help:
+            hint_indent = len(ret) - len("[ret]") + len(" help: ")
+            terminal_width = math.floor(shutil.get_terminal_size().columns * 0.7)
+            terminal_width = max(terminal_width - hint_indent, 10)
+            wrapped = textwrap.fill(str(help), width=terminal_width)
+            lines = wrapped.splitlines()
+            ret += " help: " + lines[0] + "[/red]"
+            for line in lines[1:]:
+                ret += (
+                    "\n"
+                    + " " * (error_lineno_len + indent)
+                    + "[dim]│[/dim]"
+                    + " " * hint_indent
+                    + "[red]"
+                    + line
+                    + "[/red]"
+                )
+        if show_lineno:
+            ret = " " * error_lineno_len + "[dim]│[/dim]" + ret
+        return ret
+
+    @staticmethod
+    def __get_indent(line: str) -> int:
+        if len(line) == 0:
+            return int(1e9)  # very large number
+        return len(line) - len(line.lstrip())
 
 
 def enable_stracetrace():
     """Enable the stacktrace for all exceptions."""
-    global stacktrace
-    stacktrace = True
+    global KIRIN_PYTHON_STACKTRACE
+    KIRIN_PYTHON_STACKTRACE = True
 
 
 def disable_stracetrace():
     """Disable the stacktrace for all exceptions."""
-    global stacktrace
-    stacktrace = False
+    global KIRIN_PYTHON_STACKTRACE
+    KIRIN_PYTHON_STACKTRACE = False
+
+
+def print_stacktrace(exception: Exception, state: interp.InterpreterState):
+    frame: interp.FrameABC | None = state.current_frame
+    print(
+        "==== Python stacktrace has been disabled for simplicity, set KIRIN_PYTHON_STACKTRACE=1 to enable it ===="
+    )
+    print(f"{type(exception).__name__}: {exception}", file=sys.stderr)
+    print("Traceback (most recent call last):", file=sys.stderr)
+    frames: list[interp.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)
 
 
 def exception_handler(exc_type, exc_value, exc_tb: types.TracebackType):
     """Custom exception handler to format and print exceptions."""
-    if not stacktrace and issubclass(exc_type, NoPythonStackTrace):
-        print("".join(msg for msg in exc_value.args), file=sys.stderr)
+    if not KIRIN_PYTHON_STACKTRACE and issubclass(exc_type, StaticCheckError):
+        console = Console(force_terminal=True)
+        with console.capture() as capture:
+            console.print(f"[bold red]{exc_type.__name__}:[/bold red]", end="")
+        print(capture.get(), *exc_value.args, file=sys.stderr)
+        print("Source Traceback:", file=sys.stderr)
+        print(exc_value.hint(), file=sys.stderr, end="")
         return
 
-    if not stacktrace and issubclass(exc_type, CustomStackTrace):
+    if (
+        not KIRIN_PYTHON_STACKTRACE
+        and (state := getattr(exc_value, KIRIN_INTERP_STATE, None)) is not None
+    ):
         # Handle custom stack trace exceptions
-        exc_value.print_stacktrace()
+        print_stacktrace(exc_value, state)
         return
 
     # Call the default exception handler
@@ -51,7 +183,7 @@ def exception_handler(exc_type, exc_value, exc_tb: types.TracebackType):
 
 
 def custom_exc(shell, etype, evalue, tb, tb_offset=None):
-    if issubclass(etype, NoPythonStackTrace):
+    if issubclass(etype, StaticCheckError):
         # Handle BuildError exceptions
         print(evalue, file=sys.stderr)
         return

src/kirin/interp/__init__.py

@@ -18,6 +18,7 @@
 from .base import BaseInterpreter as BaseInterpreter
 from .impl import ImplDef as ImplDef, Signature as Signature, impl as impl
 from .frame import Frame as Frame, FrameABC as FrameABC
+from .state import InterpreterState as InterpreterState
 from .table import MethodTable as MethodTable
 from .value import (
     Successor as Successor,
@@ -32,8 +33,6 @@
 )
 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

@@ -16,14 +16,14 @@
 
 from typing_extensions import Self, deprecated
 
-from kirin.ir import Block, Region, Statement, DialectGroup, traits
-from kirin.ir.method import Method
+from kirin.ir import Block, Method, Region, Statement, DialectGroup, traits
+from kirin.exception import KIRIN_INTERP_STATE
 
 from .impl import Signature
 from .frame import FrameABC
 from .state import InterpreterState
 from .value import Successor, ReturnValue, SpecialValue, StatementResult
-from .exceptions import IntepreterExit, InterpreterError
+from .exceptions import InterpreterError
 
 if TYPE_CHECKING:
     from kirin.registry import StatementImpl, InterpreterRegistry
@@ -155,9 +155,10 @@ def run(
         try:
             _, results = self.run_method(mt, args)
         except Exception as e:
-            # NOTE: initialize will create new State
-            # so we don't need to copy the frames.
-            raise IntepreterExit(e, self.state) from e
+            # NOTE: insert the interpreter state into the exception
+            # so we can print the stack trace
+            setattr(e, KIRIN_INTERP_STATE, self.state)
+            raise e
         finally:
             self._eval_lock = False
             sys.setrecursionlimit(current_recursion_limit)
@@ -175,12 +176,10 @@ def run_stmt(
         Returns:
             StatementResult[ValueType]: the result of the statement.
         """
-        frame = self.initialize_frame(stmt)
-        self.state.push_frame(frame)
-        frame.set_values(stmt.args, args)
-        results = self.eval_stmt(frame, stmt)
-        self.state.pop_frame()
-        return results
+        with self.new_frame(stmt) as frame:
+            frame.set_values(stmt.args, args)
+            results = self.eval_stmt(frame, stmt)
+            return results
 
     @abstractmethod
     def run_method(

src/kirin/interp/exceptions.py

@@ -1,56 +1,16 @@
 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):
     """Generic interpreter error.
 
-    This is the base class for all interpreter errors. Interpreter
-    errors will be catched by the interpreter and handled appropriately
-    as an error with stack trace (of Kirin, not Python) from the interpreter.
+    This is the base class for all interpreter errors.
     """
 
     pass
 
 
-@dataclass
-class WrapException(InterpreterError):
-    """A special interpreter error that wraps a Python exception."""
-
-    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/ir/__init__.py

@@ -43,7 +43,6 @@
 from kirin.ir.attrs.py import PyAttr as PyAttr
 from kirin.ir.attrs.abc import Attribute as Attribute, AttributeMeta as AttributeMeta
 from kirin.ir.exception import (
-    HintedError as HintedError,
     CompilerError as CompilerError,
     TypeCheckError as TypeCheckError,
     ValidationError as ValidationError,