add module dialect to support (de)serialization (#151)

adding the module dialect to support (de)serialization so that functions
can be compiled into a serializable format (without `Method` object in
the IR).

To support runing this dialect, interpreter now contains a symbol table.
Note that the symbol table does not implements a namespace, developers
need to mangle the function name when rewrite your function into a
serializable `module` dialect.

Author: Roger-luo

src/kirin/dialects/_pprint_helper.py

@@ -0,0 +1,61 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from kirin.ir import Statement
+    from kirin.print.printer import Printer
+
+
+def pprint_calllike(
+    invoke_or_call: "Statement", callee: str, printer: "Printer"
+) -> None:
+    with printer.rich(style="red"):
+        printer.print_name(invoke_or_call)
+    printer.plain_print(" ")
+
+    n_total = len(invoke_or_call.args)
+    # if (callee := getattr(invoke_or_call, "callee", None)) is None:
+    #     raise ValueError(f"{invoke_or_call} does not have a callee")
+
+    # if isinstance(callee, SSAValue):
+    #     printer.plain_print(printer.state.ssa_id[callee])
+    # elif isinstance(callee, Method):
+    #     printer.plain_print(callee.sym_name)
+    # elif isinstance(callee, str):
+    #     printer.plain_print(callee)
+    # else:
+    #     raise ValueError(f"Unknown callee type {type(callee)}")
+
+    printer.plain_print(callee)
+    if (inputs := getattr(invoke_or_call, "inputs", None)) is None:
+        raise ValueError(f"{invoke_or_call} does not have inputs")
+
+    if not isinstance(inputs, tuple):
+        raise ValueError(f"inputs of {invoke_or_call} is not a tuple")
+
+    if (kwargs := getattr(invoke_or_call, "kwargs", None)) is None:
+        raise ValueError(f"{invoke_or_call} does not have kwargs")
+
+    if not isinstance(kwargs, tuple):
+        raise ValueError(f"kwargs of {invoke_or_call} is not a tuple")
+
+    positional = inputs[: n_total - len(kwargs)]
+    kwargs = dict(
+        zip(
+            kwargs,
+            inputs[n_total - len(kwargs) :],
+        )
+    )
+
+    printer.plain_print("(")
+    printer.print_seq(positional)
+    if kwargs and positional:
+        printer.plain_print(", ")
+    printer.print_mapping(kwargs, delim=", ")
+    printer.plain_print(")")
+
+    with printer.rich(style="black"):
+        printer.plain_print(" : ")
+        printer.print_seq(
+            [result.type for result in invoke_or_call._results],
+            delim=", ",
+        )

src/kirin/dialects/func/stmts.py

@@ -1,5 +1,3 @@
-from typing import Union
-
 from kirin.ir import (
     Pure,
     Method,
@@ -23,43 +21,7 @@
 from kirin.dialects.func.attrs import Signature, MethodType
 from kirin.dialects.func.dialect import dialect
 
-
-def _print_invoke_or_call(
-    invoke_or_call: Union["Invoke", "Call"], printer: Printer
-) -> None:
-    with printer.rich(style="red"):
-        printer.print_name(invoke_or_call)
-    printer.plain_print(" ")
-
-    n_total = len(invoke_or_call.args)
-    callee = invoke_or_call.callee
-    if isinstance(callee, SSAValue):
-        printer.plain_print(printer.state.ssa_id[callee])
-    else:
-        printer.plain_print(callee.sym_name)
-
-    inputs = invoke_or_call.inputs
-    positional = inputs[: n_total - len(invoke_or_call.kwargs)]
-    kwargs = dict(
-        zip(
-            invoke_or_call.kwargs,
-            inputs[n_total - len(invoke_or_call.kwargs) :],
-        )
-    )
-
-    printer.plain_print("(")
-    printer.print_seq(positional)
-    if kwargs and positional:
-        printer.plain_print(", ")
-    printer.print_mapping(kwargs, delim=", ")
-    printer.plain_print(")")
-
-    with printer.rich(style="black"):
-        printer.plain_print(" : ")
-        printer.print_seq(
-            [result.type for result in invoke_or_call._results],
-            delim=", ",
-        )
+from .._pprint_helper import pprint_calllike
 
 
 class FuncOpCallableInterface(CallableStmtInterface["Function"]):
@@ -116,7 +78,7 @@ class Call(Statement):
     result: ResultValue = info.result()
 
     def print_impl(self, printer: Printer) -> None:
-        _print_invoke_or_call(self, printer)
+        pprint_calllike(self, printer.state.ssa_id[self.callee], printer)
 
 
 @statement(dialect=dialect)
@@ -242,7 +204,7 @@ class Invoke(Statement):
     result: ResultValue = info.result()
 
     def print_impl(self, printer: Printer) -> None:
-        _print_invoke_or_call(self, printer)
+        pprint_calllike(self, self.callee.sym_name, printer)
 
     def verify(self) -> None:
         if self.kwargs:

src/kirin/dialects/module.py

@@ -0,0 +1,139 @@
+"""Module dialect provides a simple module
+that is roughly a list of function statements.
+
+This dialect provides the dialect necessary for compiling a function into
+lower-level IR with all its callee functions.
+"""
+
+from kirin import ir, interp
+from kirin.decl import info, statement
+from kirin.print import Printer
+from kirin.analysis import TypeInference
+from kirin.exceptions import InterpreterError, VerificationError
+
+from ._pprint_helper import pprint_calllike
+
+dialect = ir.Dialect("module")
+
+
+@statement(dialect=dialect)
+class Module(ir.Statement):
+    traits = frozenset(
+        {ir.IsolatedFromAbove(), ir.SymbolTable(), ir.SymbolOpInterface()}
+    )
+    sym_name: str = info.attribute(property=True)
+    entry: str = info.attribute(property=True)
+    body: ir.Region = info.region(multi=False)
+
+
+@statement(dialect=dialect)
+class Invoke(ir.Statement):
+    """A special statement that represents
+    a function calling functions by symbol name.
+
+    Note:
+        This statement is here for completeness, for interpretation,
+        it is recommended to rewrite this statement into a `func.Invoke`
+        after looking up the symbol table.
+    """
+
+    callee: str = info.attribute(property=True)
+    inputs: tuple[ir.SSAValue, ...] = info.argument()
+    kwargs: tuple[str, ...] = info.attribute(property=True)
+    result: ir.ResultValue = info.result()
+
+    def print_impl(self, printer: Printer) -> None:
+        pprint_calllike(self, self.callee, printer)
+
+    def verify(self) -> None:
+        if self.kwargs:
+            for name in self.kwargs:
+                if name not in self.callee:
+                    raise VerificationError(
+                        self,
+                        f"method {self.callee} does not have argument {name}",
+                    )
+        elif len(self.callee) - 1 != len(self.args):
+            raise VerificationError(
+                self,
+                f"expected {len(self.callee)} arguments, got {len(self.args)}",
+            )
+
+
+@dialect.register
+class Concrete(interp.MethodTable):
+
+    @interp.impl(Module)
+    def interp_Module(
+        self, interp: interp.Interpreter, frame: interp.Frame, stmt: Module
+    ):
+        for stmt_ in stmt.body.blocks[0].stmts:
+            if (trait := stmt.get_trait(ir.SymbolOpInterface)) is not None:
+                interp.symbol_table[trait.get_sym_name(stmt_).data] = stmt_
+        return ()
+
+    @interp.impl(Invoke)
+    def interp_Invoke(
+        self, interpreter: interp.Interpreter, frame: interp.Frame, stmt: Invoke
+    ):
+        callee = interpreter.symbol_table.get(stmt.callee)
+        if callee is None:
+            raise InterpreterError(f"symbol {stmt.callee} not found")
+
+        trait = callee.get_trait(ir.CallableStmtInterface)
+        if trait is None:
+            raise InterpreterError(
+                f"{stmt.callee} is not callable, got {callee.__class__.__name__}"
+            )
+
+        body = trait.get_callable_region(callee)
+        mt = ir.Method(
+            mod=None,
+            py_func=None,
+            sym_name=stmt.callee,
+            arg_names=[
+                arg.name or str(idx) for idx, arg in enumerate(body.blocks[0].args)
+            ],
+            dialects=interpreter.dialects,
+            code=stmt,
+        )
+        return interpreter.run_method(mt, frame.get_values(stmt.inputs))
+
+
+@dialect.register(key="typeinfer")
+class TypeInfer(interp.MethodTable):
+
+    @interp.impl(Module)
+    def typeinfer_Module(
+        self, interp: TypeInference, frame: interp.Frame, stmt: Module
+    ):
+        for stmt_ in stmt.body.blocks[0].stmts:
+            if (trait := stmt.get_trait(ir.SymbolOpInterface)) is not None:
+                interp.symbol_table[trait.get_sym_name(stmt_).data] = stmt_
+        return ()
+
+    @interp.impl(Invoke)
+    def typeinfer_Invoke(
+        self, interp: TypeInference, frame: interp.Frame, stmt: Invoke
+    ):
+        callee = interp.symbol_table.get(stmt.callee)
+        if callee is None:
+            return (ir.types.Bottom,)
+
+        trait = callee.get_trait(ir.CallableStmtInterface)
+        if trait is None:
+            return (ir.types.Bottom,)
+
+        body = trait.get_callable_region(callee)
+        mt = ir.Method(
+            mod=None,
+            py_func=None,
+            sym_name=stmt.callee,
+            arg_names=[
+                arg.name or str(idx) for idx, arg in enumerate(body.blocks[0].args)
+            ],
+            dialects=interp.dialects,
+            code=stmt,
+        )
+        interp.run_method(mt, mt.arg_types)
+        return tuple(result.type for result in callee.results)

src/kirin/interp/base.py

@@ -66,6 +66,7 @@ def __init__(
         self.bottom = bottom
 
         self.registry = self.dialects.registry.interpreter(keys=self.keys)
+        self.symbol_table: dict[str, Statement] = {}
         self.state: InterpreterState[FrameType] = InterpreterState()
         self.fuel = fuel
         self.max_depth = max_depth

src/kirin/symbol_table.py

@@ -0,0 +1,27 @@
+from typing import Generic, TypeVar
+from dataclasses import field, dataclass
+
+T = TypeVar("T")
+
+
+@dataclass
+class SymbolTable(Generic[T]):
+    names: dict[str, T] = field(default_factory=dict)
+    """The table that maps names to values."""
+    prefix: str = field(default="", kw_only=True)
+    name_count: dict[str, int] = field(default_factory=dict, kw_only=True)
+    """The count of names that have been requested."""
+
+    def __getitem__(self, name: str) -> T:
+        return self.names[name]
+
+    def __contains__(self, name: str) -> bool:
+        return name in self.names
+
+    def __setitem__(self, name: str, value: T) -> None:
+        count = self.name_count.setdefault(name, 0)
+        self.name_count[name] = count + 1
+        self.names[f"{self.prefix}_{name}_{count}"] = value
+
+    def __delitem__(self, name: str) -> None:
+        del self.names[name]