add some docstrings for interp interface (#167)

Author: Roger-luo

src/kirin/interp/__init__.py

@@ -1,3 +1,20 @@
+"""Interpreter module for Kirin.
+
+This module contains the interpreter framework for Kirin. The interpreter
+framework is used to implement concrete and abstract interpreters for the
+IR. The interpreter framework provides a set of classes and interfaces to
+implement interpreters for the IR.
+
+The interpreter framework is designed to be extensible and customizable. It
+provides a set of base classes and interfaces for implementing concrete and
+abstract interpreters:
+
+- [`BaseInterpreter`][kirin.interp.BaseInterpreter]: Base class for implementing concrete interpreters.
+- [`AbstractInterpreter`][kirin.interp.AbstractInterpreter]: Base class for implementing abstract interpreters.
+- [`Frame`][kirin.interp.Frame]: Base class for interpreter frame.
+- [`MethodTable`][kirin.interp.MethodTable]: Method table for registering implementations of statements.
+"""
+
 from . import result as result
 from .base import BaseInterpreter as BaseInterpreter
 from .impl import ImplDef as ImplDef, Signature as Signature, impl as impl

src/kirin/interp/abstract.py

@@ -15,6 +15,12 @@
 
 @dataclass
 class AbstractFrame(Frame[ResultType]):
+    """Interpreter frame for abstract interpreter.
+
+    This frame is used to store the state of the abstract interpreter.
+    It contains the worklist of successors to be processed.
+    """
+
     worklist: WorkList[Successor[ResultType]] = field(default_factory=WorkList)
 
 
@@ -34,6 +40,29 @@ class AbstractInterpreter(
     ABC,
     metaclass=AbstractInterpreterMeta,
 ):
+    """Abstract interpreter for the IR.
+
+    This is a base class for implementing abstract interpreters for the IR.
+    It provides a framework for implementing abstract interpreters given a
+    bounded lattice type.
+
+    The abstract interpreter is a forward dataflow analysis that computes
+    the abstract values for each SSA value in the IR. The abstract values
+    are computed by evaluating the statements in the IR using the abstract
+    lattice operations.
+
+    The abstract interpreter is implemented as a worklist algorithm. The
+    worklist contains the successors of the current block to be processed.
+    The abstract interpreter processes each successor by evaluating the
+    statements in the block and updating the abstract values in the frame.
+
+    The abstract interpreter provides hooks for customizing the behavior of
+    the interpreter.
+    The [`prehook_succ`][kirin.interp.abstract.AbstractInterpreter.prehook_succ] and
+    [`posthook_succ`][kirin.interp.abstract.AbstractInterpreter.posthook_succ] methods
+    can be used to perform custom actions before and after processing a successor.
+    """
+
     lattice: type[BoundedLattice[ResultType]] = field(init=False)
     """lattice type for the abstract interpreter.
     """
@@ -50,12 +79,46 @@ def __init_subclass__(cls) -> None:
         super().__init_subclass__()
 
     def prehook_succ(self, frame: AbstractFrameType, succ: Successor):
+        """Hook called before processing a successor.
+
+        This method can be used to perform custom actions before processing
+        a successor. It is called before evaluating the statements in the block.
+
+        Args:
+            frame: The current frame of the interpreter.
+            succ: The successor to be processed.
+        """
         return
 
     def posthook_succ(self, frame: AbstractFrameType, succ: Successor):
+        """Hook called after processing a successor.
+
+        This method can be used to perform custom actions after processing
+        a successor. It is called after evaluating the statements in the block.
+
+        Args:
+            frame: The current frame of the interpreter.
+            succ: The successor that was processed.
+        """
         return
 
-    def should_exec_stmt(self, stmt: Statement):
+    def should_exec_stmt(self, stmt: Statement) -> bool:
+        """This method can be used to control which statements are executed
+        during the abstract interpretation. By default, all statements are
+        executed.
+
+        This method is useful when one wants to skip certain statements
+        during the abstract interpretation and is certain that the skipped
+        statements do not affect the final result. This would allow saving
+        computation time and memory by not evaluating the skipped statements
+        and their results.
+
+        Args:
+            stmt: The statement to be executed.
+
+        Returns:
+            True if the statement should be executed, False otherwise.
+        """
         return True
 
     def set_values(
@@ -64,6 +127,12 @@ def set_values(
         ssa: Iterable[SSAValue],
         results: Iterable[ResultType],
     ):
+        """Set the abstract values for the given SSA values in the frame.
+
+        This method is used to customize how the abstract values are set in
+        the frame. By default, the abstract values are set directly in the
+        frame.
+        """
         frame.set_values(ssa, results)
 
     def eval_recursion_limit(self, frame: AbstractFrameType) -> ResultType:

src/kirin/interp/base.py

@@ -24,6 +24,8 @@
 
 
 class InterpreterMeta(ABCMeta):
+    """A metaclass for interpreters."""
+
     pass
 
 
@@ -35,8 +37,8 @@ class BaseInterpreter(ABC, Generic[FrameType, ValueType], metaclass=InterpreterM
     designed to be subclassed to provide the actual implementation of
     the interpreter.
 
-    When subclassing, if the bases contains `ABC` no checks will be
-    performed on the subclass. If the subclass does not contain `ABC`,
+    ### Required Overrides
+    When subclassing, if the subclass does not contain `ABC`,
     the subclass must define the following attributes:
 
     - `keys`: a list of strings that defines the order of dialects to select from.
@@ -83,12 +85,13 @@ def __post_init__(self) -> None:
         self.registry = self.dialects.registry.interpreter(keys=self.keys)
 
     def initialize(self) -> Self:
-        """Initialize the interpreter global states.
-
-        This method is called before calling `eval` to initialize the
+        """Initialize the interpreter global states. This method is called before
+        calling [`run`][kirin.interp.base.BaseInterpreter.run] to initialize the
         interpreter global states.
 
-        Override this method to add custom global states.
+        !!! note "Default Implementation"
+            This method provides default behavior but may be overridden by subclasses
+            to customize or extend functionality.
         """
         self.symbol_table: dict[str, Statement] = {}
         self.state: InterpreterState[FrameType] = InterpreterState()
@@ -119,7 +122,16 @@ def run(
         args: tuple[ValueType, ...],
         kwargs: dict[str, ValueType] | None = None,
     ) -> Result[ValueType]:
-        """Run a method."""
+        """Run a method. This is the main entry point of the interpreter.
+
+        Args:
+            mt (Method): the method to run.
+            args (tuple[ValueType]): the arguments to the method, does not include self.
+            kwargs (dict[str, ValueType], optional): the keyword arguments to the method.
+
+        Returns:
+            Result[ValueType]: the result of the method.
+        """
         if self._eval_lock:
             raise InterpreterError(
                 "recursive eval is not allowed, use run_method instead"
@@ -340,6 +352,15 @@ def build_signature(self, frame: FrameType, stmt: Statement) -> "Signature":
     def lookup_registry(
         self, frame: FrameType, stmt: Statement
     ) -> Optional["StatementImpl[Self, FrameType]"]:
+        """Lookup the statement implementation in the registry.
+
+        Args:
+            frame: the current frame
+            stmt: the statement to run
+
+        Returns:
+            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]
@@ -348,7 +369,17 @@ def lookup_registry(
         return
 
     @abstractmethod
-    def run_ssacfg_region(self, frame: FrameType, region: Region) -> ValueType: ...
+    def run_ssacfg_region(self, frame: FrameType, region: Region) -> ValueType:
+        """This implements how to run a region with MLIR SSA CFG convention.
+
+        Args:
+            frame: the current frame.
+            region: the region to run.
+
+        Returns:
+            ValueType: the result of running the region.
+        """
+        ...
 
     class FuelResult(Enum):
         Stop = 0

src/kirin/interp/concrete.py

@@ -11,6 +11,13 @@
 
 
 class Interpreter(BaseInterpreter[Frame[Any], Any]):
+    """Concrete interpreter for the IR.
+
+    This is a concrete interpreter for the IR. It evaluates the IR by
+    executing the statements in the IR using a simple stack-based
+    interpreter.
+    """
+
     keys = ["main"]
     void = None
 

src/kirin/interp/exceptions.py

@@ -3,13 +3,24 @@
 
 # errors
 class InterpreterError(Exception):
+    """Generic interpreter error.
+
+    This is the base class for all interpreter errors. Interpreter
+    errors will be catched by the interpreter and handled appropriately
+    as an error with stack trace (of Kirin, not Python) from the interpreter.
+    """
+
     pass
 
 
 @dataclass
 class WrapException(InterpreterError):
+    """A special interpreter error that wraps a Python exception."""
+
     exception: Exception
 
 
 class FuelExhaustedError(InterpreterError):
+    """An error raised when the interpreter runs out of fuel."""
+
     pass

src/kirin/interp/frame.py

@@ -13,6 +13,8 @@
 
 @dataclass
 class FrameABC(ABC, Generic[ValueType]):
+    """Abstract base class for interpreter frame."""
+
     code: Statement
     """func statement being interpreted.
     """
@@ -24,17 +26,49 @@ def from_func_like(cls, code: Statement) -> Self:
         ...
 
     @abstractmethod
-    def get(self, key: SSAValue) -> ValueType: ...
+    def get(self, key: SSAValue) -> ValueType:
+        """Get the value for the given [`SSAValue`][kirin.ir.SSAValue] key.
+        See also [`get_values`][kirin.interp.frame.Frame.get_values].
+
+        Args:
+            key(SSAValue): The key to get the value for.
+
+        Returns:
+            ValueType: The value.
+        """
+        ...
 
     @abstractmethod
-    def set(self, key: SSAValue, value: ValueType) -> None: ...
+    def set(self, key: SSAValue, value: ValueType) -> None:
+        """Set the value for the given [`SSAValue`][kirin.ir.SSAValue] key.
+        See also [`set_values`][kirin.interp.frame.Frame.set_values].
+
+        Args:
+            key(SSAValue): The key to set the value for.
+            value(ValueType): The value.
+        """
+        ...
 
     def get_values(self, keys: Iterable[SSAValue]) -> tuple[ValueType, ...]:
-        """Get the values of the given `SSAValue` keys."""
+        """Get the values of the given [`SSAValue`][kirin.ir.SSAValue] keys.
+        See also [`get`][kirin.interp.frame.Frame.get].
+
+        Args:
+            keys(Iterable[SSAValue]): The keys to get the values for.
+
+        Returns:
+            tuple[ValueType, ...]: The values.
+        """
         return tuple(self.get(key) for key in keys)
 
     def set_values(self, keys: Iterable[SSAValue], values: Iterable[ValueType]) -> None:
-        """Set the values of the given `SSAValue` keys."""
+        """Set the values of the given [`SSAValue`][kirin.ir.SSAValue] keys.
+        This is a convenience method to set multiple values at once.
+
+        Args:
+            keys(Iterable[SSAValue]): The keys to set the values for.
+            values(Iterable[ValueType]): The values.
+        """
         for key, value in zip(keys, values):
             self.set(key, value)
 
@@ -46,6 +80,8 @@ def set_stmt(self, stmt: Statement) -> Self:
 
 @dataclass
 class Frame(FrameABC[ValueType]):
+    """Interpreter frame."""
+
     lino: int = 0
     stmt: Statement | None = None
     """statement being interpreted.
@@ -65,6 +101,7 @@ class Frame(FrameABC[ValueType]):
 
     @classmethod
     def from_func_like(cls, code: Statement) -> Self:
+        """Create a new frame for the given statement."""
         return cls(code=code)
 
     def get(self, key: SSAValue) -> ValueType: