rework interp (part I) (#380)

this PR cleans up the interpreter registry and simplifies `MethodTable`
and `impl`. We will do a rewrite of the interpreter interface in the
following PR.

Author: Roger-luo

src/kirin/analysis/typeinfer/analysis.py

@@ -3,7 +3,6 @@
 from kirin import ir, types, interp
 from kirin.decl import fields
 from kirin.analysis import const
-from kirin.interp.impl import Signature
 from kirin.analysis.forward import Forward, ForwardFrame
 
 from .solve import TypeResolution
@@ -41,15 +40,15 @@ def run_analysis(
     # value (which is a type) to determine the method dispatch.
     def build_signature(
         self, frame: ForwardFrame[types.TypeAttribute], stmt: ir.Statement
-    ) -> Signature:
+    ) -> interp.Signature:
         _args = ()
         for x in frame.get_values(stmt.args):
             # TODO: remove this after we have multiple dispatch...
             if isinstance(x, types.Generic):
                 _args += (x.body,)
             else:
                 _args += (x,)
-        return Signature(stmt.__class__, _args)
+        return interp.Signature(stmt.__class__, _args)
 
     def eval_stmt_fallback(
         self, frame: ForwardFrame[types.TypeAttribute], stmt: ir.Statement

src/kirin/dialects/func/emit.py

@@ -19,7 +19,7 @@ def emit_function(
     ):
         fn_args = stmt.body.blocks[0].args[1:]
         argnames = tuple(interp.ssa_id[arg] for arg in fn_args)
-        argtypes = tuple(interp.emit_attribute(x.type) for x in fn_args)
+        argtypes = tuple(interp.emit_attribute(frame, x.type) for x in fn_args)
         args = [f"{name}::{type}" for name, type in zip(argnames, argtypes)]
         interp.write(f"function {stmt.sym_name}({', '.join(args)})")
         frame.indent += 1

src/kirin/dialects/ilist/constprop.py

@@ -23,7 +23,9 @@ def detect_purity(
 
     @impl(Map)
     @impl(ForEach)
-    def one_args(self, interp_: const.Propagate, frame: const.Frame, stmt: Map):
+    def one_args(
+        self, interp_: const.Propagate, frame: const.Frame, stmt: ForEach | Map
+    ):
         fn, collection = frame.get(stmt.fn), frame.get(stmt.collection)
 
         # 1. if the function is a constant method, and the method is pure, then the map is pure

src/kirin/dialects/py/constant.py

@@ -81,4 +81,4 @@ class JuliaTable(interp.MethodTable):
 
     @interp.impl(Constant)
     def emit_Constant(self, emit: EmitJulia, frame: EmitStrFrame, stmt: Constant):
-        return (emit.emit_attribute(stmt.value),)
+        return (emit.emit_attribute(frame, stmt.value),)

src/kirin/dialects/scf/typeinfer.py

@@ -21,7 +21,7 @@ def if_else_(
         frame.set(
             stmt.cond, frame.get(stmt.cond).meet(types.Bool)
         )  # set cond backwards
-        return super().if_else(self, interp_, frame, stmt)
+        return super().if_else(interp_, frame, stmt)
 
     @interp.impl(For)
     def for_loop(

src/kirin/emit/abc.py

@@ -45,14 +45,16 @@ def run_ssacfg_region(
             frame.block_ref[succ.block] = block_header
         return ()
 
-    def emit_attribute(self, attr: ir.Attribute) -> ValueType:
+    def emit_attribute(self, frame: FrameType, attr: ir.Attribute) -> ValueType:
         return getattr(
             self, f"emit_type_{type(attr).__name__}", self.emit_attribute_fallback
-        )(attr)
+        )(frame, attr)
 
-    def emit_attribute_fallback(self, attr: ir.Attribute) -> ValueType:
-        if (method := self.registry.attributes.get(type(attr))) is not None:
-            return method(self, attr)
+    def emit_attribute_fallback(
+        self, frame: FrameType, attr: ir.Attribute
+    ) -> ValueType:
+        if (method := self.registry.get(interp.Signature(type(attr)))) is not None:
+            return method(self, frame, attr)
         raise NotImplementedError(f"Attribute {type(attr)} not implemented")
 
     def emit_stmt_begin(self, frame: FrameType, stmt: ir.Statement) -> None:

src/kirin/emit/abc.pyi

@@ -24,17 +24,21 @@ class EmitABC(interp.BaseInterpreter[FrameType, ValueType]):
     def run_ssacfg_region(
         self, frame: FrameType, region: ir.Region, args: tuple[ValueType, ...]
     ) -> tuple[ValueType, ...]: ...
-    def emit_attribute(self, attr: ir.Attribute) -> ValueType: ...
-    def emit_type_Any(self, attr: types.AnyType) -> ValueType: ...
-    def emit_type_Bottom(self, attr: types.BottomType) -> ValueType: ...
-    def emit_type_Literal(self, attr: types.Literal) -> ValueType: ...
-    def emit_type_Union(self, attr: types.Union) -> ValueType: ...
-    def emit_type_TypeVar(self, attr: types.TypeVar) -> ValueType: ...
-    def emit_type_Vararg(self, attr: types.Vararg) -> ValueType: ...
-    def emit_type_Generic(self, attr: types.Generic) -> ValueType: ...
-    def emit_type_PyClass(self, attr: types.PyClass) -> ValueType: ...
-    def emit_type_PyAttr(self, attr: ir.PyAttr) -> ValueType: ...
-    def emit_attribute_fallback(self, attr: ir.Attribute) -> ValueType: ...
+    def emit_attribute(self, frame: FrameType, attr: ir.Attribute) -> ValueType: ...
+    def emit_type_Any(self, frame: FrameType, attr: types.AnyType) -> ValueType: ...
+    def emit_type_Bottom(
+        self, frame: FrameType, attr: types.BottomType
+    ) -> ValueType: ...
+    def emit_type_Literal(self, frame: FrameType, attr: types.Literal) -> ValueType: ...
+    def emit_type_Union(self, frame: FrameType, attr: types.Union) -> ValueType: ...
+    def emit_type_TypeVar(self, frame: FrameType, attr: types.TypeVar) -> ValueType: ...
+    def emit_type_Vararg(self, frame: FrameType, attr: types.Vararg) -> ValueType: ...
+    def emit_type_Generic(self, frame: FrameType, attr: types.Generic) -> ValueType: ...
+    def emit_type_PyClass(self, frame: FrameType, attr: types.PyClass) -> ValueType: ...
+    def emit_type_PyAttr(self, frame: FrameType, attr: ir.PyAttr) -> ValueType: ...
+    def emit_attribute_fallback(
+        self, frame: FrameType, attr: ir.Attribute
+    ) -> ValueType: ...
     def emit_stmt_begin(self, frame: FrameType, stmt: ir.Statement) -> None: ...
     def emit_stmt_end(self, frame: FrameType, stmt: ir.Statement) -> None: ...
     def emit_block_begin(self, frame: FrameType, block: ir.Block) -> None: ...

src/kirin/emit/julia.py

@@ -29,7 +29,7 @@ def emit_block_begin(self, frame: EmitStrFrame, block: Block) -> None:
         self.newline(frame)
         self.write(f"@label {block_id};")
 
-    def emit_type_PyClass(self, attr: PyClass) -> str:
+    def emit_type_PyClass(self, frame: EmitStrFrame, attr: PyClass) -> str:
         return self.PYTYPE_MAP.get(attr.typ, "Any")
 
     def write_assign(self, frame: EmitStrFrame, result: ir.SSAValue, *args):
@@ -54,7 +54,7 @@ def emit_binaryop(
             ),
         )
 
-    def emit_type_PyAttr(self, attr: ir.PyAttr) -> str:
+    def emit_type_PyAttr(self, frame: EmitStrFrame, attr: ir.PyAttr) -> str:
         if isinstance(attr.data, (int, float)):
             return repr(attr.data)
         elif isinstance(attr.data, str):

src/kirin/interp/__init__.py

@@ -16,10 +16,9 @@
 """
 
 from .base import BaseInterpreter as BaseInterpreter
-from .impl import ImplDef as ImplDef, Signature as Signature, impl as impl
 from .frame import Frame as Frame, FrameABC as FrameABC
 from .state import InterpreterState as InterpreterState
-from .table import MethodTable as MethodTable
+from .table import Signature as Signature, MethodTable as MethodTable, impl as impl
 from .value import (
     Successor as Successor,
     YieldValue as YieldValue,
@@ -32,6 +31,11 @@
     AbstractInterpreter as AbstractInterpreter,
 )
 from .concrete import Interpreter as Interpreter
+from .undefined import (
+    Undefined as Undefined,
+    UndefinedType as UndefinedType,
+    is_undefined as is_undefined,
+)
 from .exceptions import (
     InterpreterError as InterpreterError,
     FuelExhaustedError as FuelExhaustedError,

src/kirin/interp/base.py

@@ -1,16 +1,9 @@
+from __future__ import annotations
+
 import sys
 from abc import ABC, ABCMeta, abstractmethod
 from enum import Enum
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Generic,
-    TypeVar,
-    ClassVar,
-    Optional,
-    Sequence,
-    Generator,
-)
+from typing import Any, Generic, TypeVar, ClassVar, Optional, Sequence, Generator
 from contextlib import contextmanager
 from dataclasses import field, dataclass
 
@@ -19,15 +12,12 @@
 from kirin.ir import Block, Method, Region, Statement, DialectGroup, traits
 from kirin.exception import KIRIN_INTERP_STATE
 
-from .impl import Signature
 from .frame import FrameABC
 from .state import InterpreterState
+from .table import Signature, BoundedDef
 from .value import Successor, ReturnValue, SpecialValue, StatementResult
 from .exceptions import InterpreterError
 
-if TYPE_CHECKING:
-    from kirin.registry import StatementImpl, InterpreterRegistry
-
 ValueType = TypeVar("ValueType")
 FrameType = TypeVar("FrameType", bound=FrameABC)
 
@@ -57,9 +47,6 @@ class BaseInterpreter(ABC, Generic[FrameType, ValueType], metaclass=InterpreterM
     keys: ClassVar[list[str]]
     """The name of the interpreter to select from dialects by order.
     """
-    Frame: ClassVar[type[FrameABC]] = field(init=False)
-    """The type of the frame to use for this interpreter.
-    """
     void: ValueType = field(init=False)
     """What to return when the interpreter evaluates nothing.
     """
@@ -80,7 +67,7 @@ class BaseInterpreter(ABC, Generic[FrameType, ValueType], metaclass=InterpreterM
     """
 
     # global states
-    registry: "InterpreterRegistry" = field(init=False, compare=False)
+    registry: dict[Signature, BoundedDef] = field(init=False, compare=False)
     """The interpreter registry.
     """
     symbol_table: dict[str, Statement] = field(init=False, compare=False)
@@ -433,7 +420,7 @@ def build_signature(self, frame: FrameType, stmt: Statement) -> "Signature":
 
     def lookup_registry(
         self, frame: FrameType, stmt: Statement
-    ) -> Optional["StatementImpl[Self, FrameType]"]:
+    ) -> Optional[BoundedDef]:
         """Lookup the statement implementation in the registry.
 
         Args:
@@ -444,10 +431,10 @@ def lookup_registry(
             Optional[StatementImpl]: the statement implementation if found, None otherwise.
         """
         sig = self.build_signature(frame, stmt)
-        if sig in self.registry.statements:
-            return self.registry.statements[sig]
-        elif (class_sig := Signature(stmt.__class__)) in self.registry.statements:
-            return self.registry.statements[class_sig]
+        if sig in self.registry:
+            return self.registry[sig]
+        elif (method := self.registry.get(Signature(stmt.__class__))) is not None:
+            return method
         return
 
     @abstractmethod