polish lowering (#362)

change the return of `lowering.run` from `Statement` to `Region` so we
can keep the original result as much as possible. We return
`ir.Statement` only for `python_function` because a function is mapped
to a `Statement`.

Author: Roger-luo

src/kirin/ir/dialect.py

@@ -110,7 +110,7 @@ def wrapper(node: type[T]) -> type[T]:
             elif issubclass(node, MethodTable):
                 if key in self.interps:
                     raise ValueError(
-                        f"Cannot register {node} to Dialect, key {key} exists"
+                        f"Cannot register {node} to Dialect, key {key} exists in {self}"
                     )
                 self.interps[key] = node()
             elif issubclass(node, FromPythonAST):

src/kirin/lowering/abc.py

@@ -4,7 +4,7 @@
 from typing import TYPE_CHECKING, Any, Generic, TypeVar, TypeAlias
 from dataclasses import dataclass
 
-from kirin.ir import SSAValue, Statement, DialectGroup
+from kirin.ir import Region, SSAValue, Statement, DialectGroup
 
 from .exception import BuildError
 
@@ -49,7 +49,7 @@ def run(
         lineno_offset: int = 0,
         col_offset: int = 0,
         compactify: bool = True,
-    ) -> Statement: ...
+    ) -> Region: ...
 
     @abstractmethod
     def visit(self, state: State[ASTNodeType], node: ASTNodeType) -> Result:

src/kirin/lowering/frame.py

@@ -65,7 +65,7 @@ def _push_stmt(self, stmt: StmtType) -> StmtType:
             raise BuildError(f"unexpected builtin statement {stmt.name}")
         elif stmt.dialect not in self.state.parent.dialects:
             raise BuildError(
-                f"Unsupported dialect `{stmt.dialect.name}` in statement {stmt.name}"
+                f"Unsupported dialect `{stmt.dialect.name}` from statement {stmt.name}"
             )
         self.curr_block.stmts.append(stmt)
         if stmt.source is None:

src/kirin/lowering/python/glob.py

@@ -23,9 +23,14 @@ class GlobalExprEval(ast.NodeVisitor):
     frame: Frame
 
     def generic_visit(self, node: ast.AST) -> Any:
+        if isinstance(node, ast.AST):
+            raise GlobalEvalError(
+                node,
+                f"Cannot lower global {node.__class__.__name__} node: {ast.dump(node)}",
+            )
         raise GlobalEvalError(
             node,
-            f"Cannot lower global {node.__class__.__name__} node: {ast.dump(node)}",
+            f"Unexpected global `{node.__class__.__name__}` node: {repr(node)} is not an AST node",
         )
 
     def visit_Name(self, node: ast.Name) -> Any:

src/kirin/lowering/python/lowering.py

@@ -81,7 +81,7 @@ def python_function(
         except Exception:
             nonlocals = {}
         globals.update(nonlocals)
-        return self.run(
+        region = self.run(
             ast.parse(source).body[0],
             source=source,
             globals=globals,
@@ -90,6 +90,13 @@ def python_function(
             col_offset=col_offset,
             compactify=compactify,
         )
+        if not region.blocks:
+            raise ValueError("No block generated")
+
+        code = region.blocks[0].first_stmt
+        if code is None:
+            raise ValueError("No code generated")
+        return code
 
     def run(
         self,
@@ -101,7 +108,7 @@ def run(
         lineno_offset: int = 0,
         col_offset: int = 0,
         compactify: bool = True,
-    ) -> ir.Statement:
+    ) -> ir.Region:
         source = source or ast.unparse(stmt)
         state = State(
             self,
@@ -132,18 +139,12 @@ def run(
                     raise e
 
             region = frame.curr_region
-            if not region.blocks:
-                raise ValueError("No block generated")
-
-            code = region.blocks[0].first_stmt
-            if code is None:
-                raise ValueError("No code generated")
 
         if compactify:
             from kirin.rewrite import Walk, CFGCompactify
 
-            Walk(CFGCompactify()).rewrite(code)
-        return code
+            Walk(CFGCompactify()).rewrite(region)
+        return region
 
     def lower_literal(self, state: State[ast.AST], value) -> ir.SSAValue:
         return state.lower(ast.Constant(value=value)).expect_one()

src/kirin/lowering/state.py

@@ -50,7 +50,7 @@ class State(Generic[ASTNodeType]):
     "current frame being lowered"
 
     def __repr__(self) -> str:
-        return f"lowering.State({self.current_frame})"
+        return f"lowering.State(current_frame={self._current_frame})"
 
     @property
     def code(self):
@@ -178,14 +178,19 @@ def frame(
         entr_block = entr_block or Block()
         region.blocks.append(entr_block)
 
+        if self._current_frame is not None:
+            globals = globals or self.current_frame.globals
+        else:
+            globals = globals or {}
+
         frame = Frame(
             state=self,
             stream=stmts,
             curr_region=region or Region(entr_block),
             entr_block=entr_block,
             curr_block=entr_block,
             next_block=next_block or Block(),
-            globals=globals or self.current_frame.globals,
+            globals=globals,
             capture_callback=capture_callback,
         )
         self.push_frame(frame)