added stacktrace for error message

LuisaGroup · Feb 4, 2025 · 161517d · 161517d
1 parent fd66c6f
commit 161517d
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 24 deletions.
diff --git a/luisa_lang/hir.py b/luisa_lang/hir.py
@@ -1166,12 +1166,15 @@ def __str__(self) -> str:
             return f"Template matching error:\n\t{self.message}"
         return f"Template matching error at {self.span}:\n\t{self.message}"
 
+
 class ComptimeCallStack:
     pass
 
+
 class SpannedError(Exception):
     span: Span | None
     message: str
+    stack_trace: str | None
 
     def __init__(self, node: Node | Span | ast.AST | None, message: str) -> None:
         if node is not None:
@@ -1185,27 +1188,32 @@ def __init__(self, node: Node | Span | ast.AST | None, message: str) -> None:
         else:
             self.span = None
         self.message = message
+        self.stack_trace = None
+
+
+def _pretty_print_error(err_kind: str, stack_trace: Optional[str], span: Optional[Span], message: str) -> str:
+    s = f"{err_kind} error"
+    if span is not None:
+        s += f" at {span}"
+    s += f":\n\t{message}"
+    if stack_trace is not None:
+        s += f"\n{stack_trace}"
+    return s
 
 
 class ParsingError(SpannedError):
     def __str__(self) -> str:
-        if self.span is None:
-            return f"Parsing error:\n\t{self.message}"
-        return f"Parsing error at {self.span}:\n\t{self.message}"
+        return _pretty_print_error("Parsing", self.stack_trace, self.span, self.message)
 
 
 class InlineError(SpannedError):
     def __str__(self) -> str:
-        if self.span is None:
-            return f"Inline error:\n\t{self.message}"
-        return f"Inline error at {self.span}:\n\t{self.message}"
+        return _pretty_print_error("Inline", self.stack_trace, self.span, self.message)
 
 
 class TypeInferenceError(SpannedError):
     def __str__(self) -> str:
-        if self.span is None:
-            return f"Type inference error:\n\t{self.message}"
-        return f"Type inference error at {self.span}:\n\t{self.message}"
+        return _pretty_print_error("Type inference", self.stack_trace, self.span, self.message)
 
 
 class Assign(Node):

diff --git a/luisa_lang/parse.py b/luisa_lang/parse.py
@@ -251,6 +251,13 @@ def push(self, f: 'FuncParser') -> None:
 
     def pop(self) -> 'FuncParser':
         return self.st.pop()
+
+    def dump_stack(self)->str:
+        trace = []
+        for i, f in enumerate(reversed(self.st)):
+            span = hir.Span.from_ast(f.func_def)
+            trace.append(f"{i}: {f.name} at {span if span else 'unknown'}")
+        return '\n'.join(trace)
 
 
 FUNC_STACK = FuncStack()
@@ -1371,21 +1378,27 @@ def parse_anno_ty() -> hir.Type:
 
     def parse_body(self):
         FUNC_STACK.push(self)
-        assert self.parsed_func is not None
-        body = self.func_def.body
-        entry = hir.BasicBlock(hir.Span.from_ast(self.func_def))
-        self.bb_stack.append(entry)
-        for stmt in body:
-            self.parse_stmt(stmt)
-        assert len(self.bb_stack) == 1
-        self.parsed_func.body = entry
-        self.parsed_func.locals = list(
-            [x for x in self.vars.values() if isinstance(x, hir.Var)])
-        if not self.parsed_func.return_type:
-            self.parsed_func.return_type = hir.UnitType()
-        self.parsed_func.complete = True
-        assert FUNC_STACK.pop() is self
-        return self.parsed_func
+        try:
+            assert self.parsed_func is not None
+            body = self.func_def.body
+            entry = hir.BasicBlock(hir.Span.from_ast(self.func_def))
+            self.bb_stack.append(entry)
+            for stmt in body:
+                self.parse_stmt(stmt)
+            assert len(self.bb_stack) == 1
+            self.parsed_func.body = entry
+            self.parsed_func.locals = list(
+                [x for x in self.vars.values() if isinstance(x, hir.Var)])
+            if not self.parsed_func.return_type:
+                self.parsed_func.return_type = hir.UnitType()
+            self.parsed_func.complete = True
+            assert FUNC_STACK.pop() is self
+            return self.parsed_func
+        except hir.SpannedError as e:
+            if e.stack_trace is None:
+                e.stack_trace =  FUNC_STACK.dump_stack()
+            raise e from e
+
 
 
 UNARY_OP_TO_METHOD_NAMES: Dict[type, str] = {