add more docstrings for analysis (#168)

Author: Roger-luo

mkdocs.yml

@@ -41,8 +41,11 @@ nav:
       - Python Lowering: reference/kirin/lowering/
       - Analysis:
           - reference/kirin/analysis/index.md
-          - Dataflow: reference/kirin/analysis/dataflow/
           - CFG: reference/kirin/analysis/cfg.md
+          - Call Graph: reference/kirin/analysis/callgraph.md
+          - Forward Dataflow: reference/kirin/analysis/forward.md
+          - Type Inference: reference/kirin/analysis/typeinfer/
+          - Constant Propagation: reference/kirin/analysis/const/
       - Rewrite:
           - Generic: reference/kirin/rewrite/
           - Rules: reference/kirin/rules/

src/kirin/analysis/__init__.py

@@ -1,3 +1,18 @@
+"""Analysis module for kirin.
+
+This module contains the analysis framework for kirin. The analysis framework is
+built on top of the interpreter framework. This module provides a set of base classes
+and frameworks for implementing compiler analysis passes on the IR.
+
+The analysis framework contains the following modules:
+
+- [`cfg`][kirin.analysis.cfg]: Control flow graph for a given IR.
+- [`forward`][kirin.analysis.forward]: Forward dataflow analysis.
+- [`callgraph`][kirin.analysis.callgraph]: Call graph for a given IR.
+- [`typeinfer`][kirin.analysis.typeinfer]: Type inference analysis.
+- [`const`][kirin.analysis.const]: Constants used in the analysis framework.
+"""
+
 from kirin.analysis import const as const
 from kirin.analysis.cfg import CFG as CFG
 from kirin.analysis.forward import Forward as Forward, ForwardExtra as ForwardExtra

src/kirin/analysis/callgraph.py

@@ -9,8 +9,19 @@
 
 @dataclass
 class CallGraph(Printable):
+    """Call graph for a given [`ir.Method`][kirin.ir.Method].
+
+    This class implements the [`kirin.graph.Graph`][kirin.graph.Graph] protocol.
+
+    !!! note "Pretty Printing"
+        This object is pretty printable via
+        [`.print()`][kirin.print.printable.Printable.print] method.
+    """
+
     defs: dict[str, ir.Method] = field(default_factory=dict)
+    """Mapping from symbol names to methods."""
     backedges: dict[str, set[str]] = field(default_factory=dict)
+    """Mapping from symbol names to backedges."""
 
     def __init__(self, mt: ir.Method):
         self.defs = {}
@@ -26,14 +37,17 @@ def __build(self, mt: ir.Method):
                 self.__build(stmt.callee)
 
     def get_neighbors(self, node: str) -> Iterable[str]:
+        """Get the neighbors of a node in the call graph."""
         return self.backedges.get(node, ())
 
     def get_edges(self) -> Iterable[tuple[str, str]]:
+        """Get the edges of the call graph."""
         for node, neighbors in self.backedges.items():
             for neighbor in neighbors:
                 yield node, neighbor
 
     def get_nodes(self) -> Iterable[str]:
+        """Get the nodes of the call graph."""
         return self.defs.keys()
 
     def print_impl(self, printer: Printer) -> None:

src/kirin/analysis/cfg.py

@@ -9,7 +9,14 @@
 
 @dataclass
 class CFG(Printable):
-    """Control Flow Graph of a given IR statement."""
+    """Control Flow Graph of a given IR statement.
+
+    This class implements the [`kirin.graph.Graph`][kirin.graph.Graph] protocol.
+
+    !!! note "Pretty Printing"
+        This object is pretty printable via
+        [`.print()`][kirin.print.printable.Printable.print] method.
+    """
 
     parent: ir.Region
     """Parent IR statement.

src/kirin/analysis/const/__init__.py

@@ -1,3 +1,13 @@
+"""Const analysis module.
+
+This module contains the constant analysis framework for kirin. The constant
+analysis framework is built on top of the interpreter framework.
+
+This module provides a lattice for constant propagation analysis and a
+propagation algorithm for computing the constant values for each SSA value in
+the IR.
+"""
+
 from .prop import Propagate as Propagate
 from .lattice import (
     Pure as Pure,

src/kirin/analysis/const/lattice.py

@@ -1,3 +1,6 @@
+"""Lattice for constant analysis.
+"""
+
 from typing import Any, final
 from dataclasses import field, dataclass
 
@@ -22,6 +25,7 @@ class Result(
     BoundedLattice["Result"],
     _ElemVisitor,
 ):
+    """Base class for constant analysis results."""
 
     @classmethod
     def top(cls) -> "Result":
@@ -35,6 +39,7 @@ def bottom(cls) -> "Result":
 @final
 @dataclass
 class Unknown(Result, metaclass=SingletonMeta):
+    """Unknown constant value. This is the top element of the lattice."""
 
     def is_subseteq(self, other: Result) -> bool:
         return isinstance(other, Unknown)
@@ -43,6 +48,7 @@ def is_subseteq(self, other: Result) -> bool:
 @final
 @dataclass
 class Bottom(Result, metaclass=SingletonMeta):
+    """Bottom element of the lattice."""
 
     def is_subseteq(self, other: Result) -> bool:
         return True
@@ -51,6 +57,8 @@ def is_subseteq(self, other: Result) -> bool:
 @final
 @dataclass
 class Value(Result):
+    """Constant value. Wraps any Python value."""
+
     data: Any
 
     def is_subseteq_Value(self, other: "Value") -> bool:
@@ -64,11 +72,19 @@ def is_equal(self, other: Result) -> bool:
 
 @dataclass
 class PartialConst(Result):
+    """Base class for partial constant values."""
+
     pass
 
 
 @final
 class PartialTupleMeta(LatticeMeta):
+    """Metaclass for PartialTuple.
+
+    This metaclass canonicalizes PartialTuple instances with all Value elements
+    into a single Value instance.
+    """
+
     def __call__(cls, data: tuple[Result, ...]):
         if all(isinstance(x, Value) for x in data):
             return Value(tuple(x.data for x in data))  # type: ignore
@@ -78,6 +94,8 @@ def __call__(cls, data: tuple[Result, ...]):
 @final
 @dataclass
 class PartialTuple(PartialConst, metaclass=PartialTupleMeta):
+    """Partial tuple constant value."""
+
     data: tuple[Result, ...]
 
     def join(self, other: Result) -> Result:
@@ -125,6 +143,11 @@ def is_subseteq_Value(self, other: Value) -> bool:
 @final
 @dataclass
 class PartialLambda(PartialConst):
+    """Partial lambda constant value.
+
+    This represents a closure with captured variables.
+    """
+
     argnames: list[str]
     code: ir.Statement
     captured: tuple[Result, ...]
@@ -177,6 +200,7 @@ def meet(self, other: Result) -> Result:
 class Purity(
     SimpleJoinMixin["Purity"], SimpleMeetMixin["Purity"], BoundedLattice["Purity"]
 ):
+    """Base class for purity lattice."""
 
     @classmethod
     def bottom(cls) -> "Purity":
@@ -189,29 +213,43 @@ def top(cls) -> "Purity":
 
 @dataclass(frozen=True)
 class Pure(Purity, metaclass=SingletonMeta):
+    """The result is from a pure function."""
 
     def is_subseteq(self, other: Purity) -> bool:
         return isinstance(other, (NotPure, Pure))
 
 
 @dataclass(frozen=True)
 class NotPure(Purity, metaclass=SingletonMeta):
+    """The result is from an impure function."""
 
     def is_subseteq(self, other: Purity) -> bool:
         return isinstance(other, NotPure)
 
 
 @dataclass(frozen=True)
 class PurityBottom(Purity, metaclass=SingletonMeta):
+    """The bottom element of the purity lattice."""
 
     def is_subseteq(self, other: Purity) -> bool:
         return True
 
 
 @dataclass
 class JointResult(BoundedLattice["JointResult"]):
+    """Joint result of constant value and purity.
+
+    This lattice is used to join the constant value and purity of a function
+    during constant propagation analysis. This allows the analysis to track
+    both the constant value and the purity of the function, so that the analysis
+    can propagate constant values through function calls even if the function
+    is only partially pure.
+    """
+
     const: Result
+    """The constant value of the result."""
     purity: Purity = field(default_factory=Purity.top)
+    """The purity of statement that produces the result."""
 
     @classmethod
     def from_const(cls, value: Any) -> "JointResult":

src/kirin/analysis/const/prop.py

@@ -13,6 +13,22 @@ class ExtraFrameInfo:
 
 @dataclass
 class Propagate(ForwardExtra[JointResult, ExtraFrameInfo]):
+    """Forward dataflow analysis for constant propagation.
+
+    This analysis is a forward dataflow analysis that propagates constant values
+    through the program. It uses the `JointResult` lattice to track the constant
+    values and purity of the values.
+
+    The analysis is implemented as a forward dataflow analysis, where the
+    `eval_stmt` method is overridden to handle the different types of statements
+    in the IR. The analysis uses the `interp.Interpreter` to evaluate the
+    statements and propagate the constant values.
+
+    When a statement is registered under the "constprop" key in the method table,
+    the analysis will call the method to evaluate the statement instead of using
+    the interpreter. This allows for custom handling of statements.
+    """
+
     keys = ["constprop"]
     lattice = JointResult
 

src/kirin/analysis/forward.py

@@ -66,6 +66,13 @@ def set_values(
         ssa: Iterable[ir.SSAValue],
         results: Iterable[LatticeElemType],
     ):
+        """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. This method is overridden to join the results if the SSA value
+        already exists in the frame.
+        """
         for ssa_value, result in zip(ssa, results):
             if ssa_value in frame.entries:
                 frame.entries[ssa_value] = frame.entries[ssa_value].join(result)
@@ -88,4 +95,11 @@ def new_frame(self, code: ir.Statement) -> ForwardFrame[LatticeElemType, ExtraTy
 
 
 class Forward(ForwardExtra[LatticeElemType, None], ABC):
+    """Forward dataflow analysis.
+
+    This is the base class for forward dataflow analysis. If your analysis
+    requires extra information per frame, you should subclass
+    [`ForwardExtra`][kirin.analysis.forward.ForwardExtra] instead.
+    """
+
     pass

src/kirin/analysis/typeinfer/__init__.py

@@ -1,2 +1,5 @@
+"""Type inference analysis for kirin.
+"""
+
 from .solve import TypeResolution as TypeResolution
 from .analysis import TypeInference as TypeInference

src/kirin/analysis/typeinfer/analysis.py

@@ -10,6 +10,17 @@
 
 @final
 class TypeInference(Forward[types.TypeAttribute]):
+    """Type inference analysis for kirin.
+
+    This analysis uses the forward dataflow analysis framework to infer the types of
+    the IR. The analysis uses the type information within the IR to determine the
+    method dispatch.
+
+    The analysis will fallback to a type resolution algorithm if the type information
+    is not available in the IR but the type information is available in the abstract
+    values.
+    """
+
     keys = ["typeinfer"]
     lattice = types.TypeAttribute