explicit Python ast.Call lowering trait (#139)

this PR removes the default behaviour of `@statement` which creates a
lowering transform from `ast.Call` to the statement call. This is
because `Statement` constructor is not necessarily compatible with
`ast.Call` syntax, e.g when `Statement` contains a `Region` or `Block`,
while most `Statement` is just a special `call`.

The change here requires one explicitly declares `FromPythonCall` trait
in the statement definition, which thus triggers the lowering transform
to match a compatible syntax in Python AST. In the future, I think we
can consider improving duplication in trait declaration by using
inheritance, e.g how `BinOp` in `py.binop` dialect is currently
implemented.

This change is important for the following support of `with xxx as f`
syntax where a given statement could have a similar roughly meaning of
`with` syntax via a trait e.g `FromPythonWithSingleContext` indicates
the statement has a compatible syntax with `with ... as f:` syntax where
the statement's only region is treated as the with body.

Author: Roger-luo

src/kirin/decl/__init__.py

@@ -15,7 +15,6 @@
 from kirin.decl.emit.dialect import EmitDialect
 from kirin.decl.emit.property import EmitProperty
 from kirin.decl.emit.typecheck import EmitTypeCheck
-from kirin.decl.emit.from_python_call import EmitFromPythonCall
 
 
 class StatementDecl(
@@ -29,7 +28,6 @@ class StatementDecl(
     EmitTraits,
     EmitVerify,
     EmitTypeCheck,
-    EmitFromPythonCall,
 ):
     pass
 
@@ -63,5 +61,5 @@ def wrap(cls):
     return wrap(cls)
 
 
-def fields(cls) -> info.StatementFields:
+def fields(cls: type[Statement]) -> info.StatementFields:
     return getattr(cls, ScanFields._FIELDS)

src/kirin/decl/emit/from_python_call.py

@@ -11,4 +11,3 @@ def emit_traits(self):
             if hasattr(base, "traits"):
                 return
         set_new_attribute(self.cls, "traits", frozenset({}))
-        return

src/kirin/decl/emit/traits.py

@@ -5,6 +5,7 @@
 
 @statement(dialect=dialect)
 class Foldl(ir.Statement):
+    traits = frozenset({ir.FromPythonCall()})
     fn: ir.SSAValue = info.argument(ir.types.PyClass(ir.Method))
     coll: ir.SSAValue = info.argument(ir.types.Any)  # TODO: make this more precise
     init: ir.SSAValue = info.argument(ir.types.Any)
@@ -13,6 +14,7 @@ class Foldl(ir.Statement):
 
 @statement(dialect=dialect)
 class Foldr(ir.Statement):
+    traits = frozenset({ir.FromPythonCall()})
     fn: ir.SSAValue = info.argument(ir.types.PyClass(ir.Method))
     coll: ir.SSAValue = info.argument(ir.types.Any)
     init: ir.SSAValue = info.argument(ir.types.Any)
@@ -39,6 +41,7 @@ def main():
     ```
     """
 
+    traits = frozenset({ir.FromPythonCall()})
     fn: ir.SSAValue = info.argument(ir.types.PyClass(ir.Method))
     """The kernel function to apply. The function should have signature `fn(x: int) -> Any`."""
     coll: ir.SSAValue = info.argument(ir.types.Any)
@@ -49,6 +52,7 @@ def main():
 
 @statement(dialect=dialect)
 class Scan(ir.Statement):
+    traits = frozenset({ir.FromPythonCall()})
     fn: ir.SSAValue = info.argument(ir.types.PyClass(ir.Method))
     init: ir.SSAValue = info.argument(ir.types.Any)
     coll: ir.SSAValue = info.argument(ir.types.List)

src/kirin/dialects/fcf/stmts.py

@@ -1,5 +1,4 @@
 "math dialect, modeling functions in python's `math` stdlib"  # This file is generated by gen.py
-from kirin.dialects.math import interp as interp
 from kirin.dialects.math.stmts import (
     cos as cos,
     erf as erf,
@@ -38,4 +37,5 @@
     isfinite as isfinite,
     remainder as remainder,
 )
+from kirin.dialects.math.interp import MathMethodTable as MathMethodTable
 from kirin.dialects.math.dialect import dialect as dialect

src/kirin/dialects/math/__init__.py

@@ -65,7 +65,7 @@ class {name}(ir.Statement):
     \"\"\"{name} statement, wrapping the math.{name} function
     \"\"\"
     name = "{name}"
-    traits = frozenset({{ir.Pure()}})
+    traits = frozenset({{ir.Pure(), ir.FromPythonCall()}})
 {fields}
     result: ir.ResultValue = info.result(Float)
 """
@@ -76,10 +76,9 @@ class {name}(ir.Statement):
 with open(os.path.join(os.path.dirname(__file__), "interp.py"), "w") as f:
     f.write("# This file is generated by gen.py\n")
     f.write("import math\n")
-    f.write("from typing import Any\n")
     f.write("from kirin.dialects.math.dialect import dialect\n")
     f.write("from kirin.dialects.math import stmts\n")
-    f.write("from kirin.interp import MethodTable, Frame, Result, impl\n")
+    f.write("from kirin.interp import MethodTable, Frame, impl\n")
     f.write("\n")
 
     implements = []
@@ -90,7 +89,7 @@ class {name}(ir.Statement):
         implements.append(
             f"""
     @impl(stmts.{name})
-    def {name}(self, interp, frame: Frame, stmt: stmts.{name}) -> Result[Any]:
+    def {name}(self, interp, frame: Frame, stmt: stmts.{name}):
         values = frame.get_values(stmt.args)
         return (math.{name}({fields}),)"""
         )
@@ -114,7 +113,9 @@ class MathMethodTable(MethodTable):
     for name, obj, sig in builtin_math_functions():
         f.write(f"    {name} as {name},\n")
     f.write(")\n")
-    f.write("from kirin.dialects.math.interp import Interpreter as Interpreter\n")
+    f.write(
+        "from kirin.dialects.math.interp import MathMethodTable as MathMethodTable\n"
+    )
     f.write("\n")
 
 for file in ["__init__.py", "interp.py", "stmts.py"]: