Refactoring AST lowering transform (#343)

finally decide to just do it... gonna supercede #315 

This PR refactors the AST lowering framework to make it more generic. So
that other langauge AST can reuse the framework to generate SSA IR (e.g
QASM2, etc.). We will also try to clean up the APIs in the process.

### Highlights

- new error message of lowering

<img width="682" alt="image"
src="https://github.com/user-attachments/assets/b5591b7a-a3dc-44fa-b977-538a2049d691"
/>

- removing the `push_frame` and `pop_frame` in `lowering.State`. Now you
can just write

```python
with state.frame(stmts) as body_frame:
   ...
```

The main changes are mostly splitting the APIs into several classes. We
still need to improve the precision of error reports here, but we will
do that in a few following PRs.

cc: @weinbe58

Author: Roger-luo

src/kirin/__init__.py

@@ -1,7 +1,6 @@
 # re-exports the public API of the kirin package
 from kirin import ir
-from kirin.decl import info, statement
 
 from . import types as types
 
-__all__ = ["ir", "types", "statement", "info"]
+__all__ = ["ir", "types"]

src/kirin/dialects/debug.py

@@ -2,27 +2,25 @@
 
 import rich
 
-from kirin import ir, decl, types, interp, lowering, exceptions
+from kirin import ir, decl, types, interp, lowering
 
 dialect = ir.Dialect("debug")
 
 
-class InfoLowering(ir.FromPythonCall):
+class InfoLowering(lowering.FromPythonCall):
 
     def lower(
-        self, stmt: type, state: lowering.LoweringState, node: ast.Call
+        self, stmt: type, state: lowering.State, node: ast.Call
     ) -> lowering.Result:
         if len(node.args) == 0:
-            raise exceptions.DialectLoweringError(
-                "info() requires at least one argument"
-            )
+            raise lowering.BuildError("info() requires at least one argument")
 
-        msg = state.visit(node.args[0]).expect_one()
+        msg = state.lower(node.args[0]).expect_one()
         if len(node.args) > 1:
-            inputs = tuple(state.visit(arg).expect_one() for arg in node.args[1:])
+            inputs = tuple(state.lower(arg).expect_one() for arg in node.args[1:])
         else:
             inputs = ()
-        return lowering.Result(state.append_stmt(Info(msg=msg, inputs=inputs)))
+        state.current_frame.push(Info(msg=msg, inputs=inputs))
 
 
 @decl.statement(dialect=dialect)

src/kirin/dialects/ilist/_wrapper.py

@@ -1,6 +1,6 @@
 import typing
 
-from kirin.lowering import wraps
+from kirin import lowering
 
 from . import stmts
 from .runtime import IList
@@ -25,42 +25,42 @@ def range(start: int, stop: int) -> IList[int, typing.Any]: ...
 def range(start: int, stop: int, step: int) -> IList[int, typing.Any]: ...
 
 
-@wraps(stmts.Range)
+@lowering.wraps(stmts.Range)
 def range(start: int, stop: int, step: int) -> IList[int, typing.Any]: ...
 
 
-@wraps(stmts.Map)
+@lowering.wraps(stmts.Map)
 def map(
     fn: typing.Callable[[ElemT], OutElemT],
     collection: IList[ElemT, LenT] | list[ElemT],
 ) -> IList[OutElemT, LenT]: ...
 
 
-@wraps(stmts.Foldr)
+@lowering.wraps(stmts.Foldr)
 def foldr(
     fn: typing.Callable[[ElemT, OutElemT], OutElemT],
     collection: IList[ElemT, LenT] | list[ElemT],
     init: OutElemT,
 ) -> OutElemT: ...
 
 
-@wraps(stmts.Foldl)
+@lowering.wraps(stmts.Foldl)
 def foldl(
     fn: typing.Callable[[OutElemT, ElemT], OutElemT],
     collection: IList[ElemT, LenT] | list[ElemT],
     init: OutElemT,
 ) -> OutElemT: ...
 
 
-@wraps(stmts.Scan)
+@lowering.wraps(stmts.Scan)
 def scan(
     fn: typing.Callable[[OutElemT, ElemT], tuple[OutElemT, ResultT]],
     collection: IList[ElemT, LenT] | list[ElemT],
     init: OutElemT,
 ) -> tuple[OutElemT, IList[ResultT, LenT]]: ...
 
 
-@wraps(stmts.ForEach)
+@lowering.wraps(stmts.ForEach)
 def for_each(
     fn: typing.Callable[[ElemT], typing.Any],
     collection: IList[ElemT, LenT] | list[ElemT],

src/kirin/dialects/ilist/lowering.py

@@ -1,17 +1,16 @@
 import ast
 
-from kirin import types
-from kirin.lowering import Result, FromPythonAST, LoweringState
+from kirin import types, lowering
 
 from . import stmts as ilist
 from ._dialect import dialect
 
 
 @dialect.register
-class PythonLowering(FromPythonAST):
+class PythonLowering(lowering.FromPythonAST):
 
-    def lower_List(self, state: LoweringState, node: ast.List) -> Result:
-        elts = tuple(state.visit(each).expect_one() for each in node.elts)
+    def lower_List(self, state: lowering.State, node: ast.List) -> lowering.Result:
+        elts = tuple(state.lower(each).expect_one() for each in node.elts)
 
         if len(elts):
             typ = elts[0].type
@@ -20,6 +19,4 @@ def lower_List(self, state: LoweringState, node: ast.List) -> Result:
         else:
             typ = types.Any
 
-        stmt = ilist.New(values=tuple(elts))
-        state.append_stmt(stmt)
-        return Result(stmt)
+        return state.current_frame.push(ilist.New(values=tuple(elts)))

src/kirin/dialects/ilist/stmts.py

@@ -1,6 +1,6 @@
 from typing import Sequence
 
-from kirin import ir, types
+from kirin import ir, types, lowering
 from kirin.decl import info, statement
 
 from .runtime import IList
@@ -14,7 +14,7 @@
 @statement(dialect=dialect)
 class Range(ir.Statement):
     name = "range"
-    traits = frozenset({ir.Pure(), ir.FromPythonRangeLike()})
+    traits = frozenset({ir.Pure(), lowering.FromPythonRangeLike()})
     start: ir.SSAValue = info.argument(types.Int)
     stop: ir.SSAValue = info.argument(types.Int)
     step: ir.SSAValue = info.argument(types.Int)
@@ -23,7 +23,7 @@ class Range(ir.Statement):
 
 @statement(dialect=dialect, init=False)
 class New(ir.Statement):
-    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
     values: tuple[ir.SSAValue, ...] = info.argument(ElemT)
     result: ir.ResultValue = info.result(IListType[ElemT])
 
@@ -49,7 +49,7 @@ def __init__(
 
 @statement(dialect=dialect)
 class Push(ir.Statement):
-    traits = frozenset({ir.Pure(), ir.FromPythonCall()})
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
     lst: ir.SSAValue = info.argument(IListType[ElemT])
     value: ir.SSAValue = info.argument(IListType[ElemT])
     result: ir.ResultValue = info.result(IListType[ElemT])
@@ -60,7 +60,7 @@ class Push(ir.Statement):
 
 @statement(dialect=dialect)
 class Map(ir.Statement):
-    traits = frozenset({ir.MaybePure(), ir.FromPythonCall()})
+    traits = frozenset({ir.MaybePure(), lowering.FromPythonCall()})
     purity: bool = info.attribute(default=False)
     fn: ir.SSAValue = info.argument(types.MethodType[[ElemT], OutElemT])
     collection: ir.SSAValue = info.argument(IListType[ElemT, ListLen])
@@ -69,7 +69,7 @@ class Map(ir.Statement):
 
 @statement(dialect=dialect)
 class Foldr(ir.Statement):
-    traits = frozenset({ir.MaybePure(), ir.FromPythonCall()})
+    traits = frozenset({ir.MaybePure(), lowering.FromPythonCall()})
     purity: bool = info.attribute(default=False)
     fn: ir.SSAValue = info.argument(
         types.Generic(ir.Method, [ElemT, OutElemT], OutElemT)
@@ -81,7 +81,7 @@ class Foldr(ir.Statement):
 
 @statement(dialect=dialect)
 class Foldl(ir.Statement):
-    traits = frozenset({ir.MaybePure(), ir.FromPythonCall()})
+    traits = frozenset({ir.MaybePure(), lowering.FromPythonCall()})
     purity: bool = info.attribute(default=False)
     fn: ir.SSAValue = info.argument(
         types.Generic(ir.Method, [OutElemT, ElemT], OutElemT)
@@ -97,7 +97,7 @@ class Foldl(ir.Statement):
 
 @statement(dialect=dialect)
 class Scan(ir.Statement):
-    traits = frozenset({ir.MaybePure(), ir.FromPythonCall()})
+    traits = frozenset({ir.MaybePure(), lowering.FromPythonCall()})
     purity: bool = info.attribute(default=False)
     fn: ir.SSAValue = info.argument(
         types.Generic(ir.Method, [OutElemT, ElemT], types.Tuple[OutElemT, ResultT])
@@ -111,7 +111,7 @@ class Scan(ir.Statement):
 
 @statement(dialect=dialect)
 class ForEach(ir.Statement):
-    traits = frozenset({ir.MaybePure(), ir.FromPythonCall()})
+    traits = frozenset({ir.MaybePure(), lowering.FromPythonCall()})
     purity: bool = info.attribute(default=False)
     fn: ir.SSAValue = info.argument(types.Generic(ir.Method, [ElemT], types.NoneType))
     collection: ir.SSAValue = info.argument(IListType[ElemT])

src/kirin/dialects/lowering/call.py

@@ -2,7 +2,6 @@
 
 from kirin import ir, types, lowering
 from kirin.dialects import func
-from kirin.exceptions import DialectLoweringError
 
 dialect = ir.Dialect("lowering.call")
 
@@ -11,38 +10,38 @@
 class Lowering(lowering.FromPythonAST):
 
     def lower_Call_local(
-        self, state: lowering.LoweringState, callee: ir.SSAValue, node: ast.Call
+        self, state: lowering.State, callee: ir.SSAValue, node: ast.Call
     ) -> lowering.Result:
         args, keywords = self.__lower_Call_args_kwargs(state, node)
         stmt = func.Call(callee, args, kwargs=keywords)
-        return lowering.Result(state.append_stmt(stmt))
+        return state.current_frame.push(stmt)
 
     def lower_Call_global_method(
         self,
-        state: lowering.LoweringState,
+        state: lowering.State,
         method: ir.Method,
         node: ast.Call,
     ) -> lowering.Result:
         args, keywords = self.__lower_Call_args_kwargs(state, node)
         stmt = func.Invoke(args, callee=method, kwargs=keywords)
         stmt.result.type = method.return_type or types.Any
-        return lowering.Result(state.append_stmt(stmt))
+        return state.current_frame.push(stmt)
 
     def __lower_Call_args_kwargs(
         self,
-        state: lowering.LoweringState,
+        state: lowering.State,
         node: ast.Call,
     ):
         args: list[ir.SSAValue] = []
         for arg in node.args:
             if isinstance(arg, ast.Starred):  # TODO: support *args
-                raise DialectLoweringError("starred arguments are not supported")
+                raise lowering.BuildError("starred arguments are not supported")
             else:
-                args.append(state.visit(arg).expect_one())
+                args.append(state.lower(arg).expect_one())
 
         keywords = []
         for kw in node.keywords:
             keywords.append(kw.arg)
-            args.append(state.visit(kw.value).expect_one())
+            args.append(state.lower(kw.value).expect_one())
 
         return tuple(args), tuple(keywords)