wrap constants after const propagation (#137)

this PR insert a rewrite after const folding, this new rule will insert
the inference result of constant propagation into the IR type annotation
if the value is a constant or partial constant.

Author: Roger-luo

src/kirin/analysis/const/__init__.py

@@ -7,6 +7,7 @@
     NotPure as NotPure,
     Unknown as Unknown,
     JointResult as JointResult,
+    PartialConst as PartialConst,
     PartialTuple as PartialTuple,
     PurityBottom as PurityBottom,
     PartialLambda as PartialLambda,

src/kirin/analysis/const/lattice.py

@@ -62,6 +62,11 @@ def is_equal(self, other: Result) -> bool:
         return False
 
 
+@dataclass
+class PartialConst(Result):
+    pass
+
+
 @final
 class PartialTupleMeta(LatticeMeta):
     def __call__(cls, data: tuple[Result, ...]):
@@ -72,7 +77,7 @@ def __call__(cls, data: tuple[Result, ...]):
 
 @final
 @dataclass
-class PartialTuple(Result, metaclass=PartialTupleMeta):
+class PartialTuple(PartialConst, metaclass=PartialTupleMeta):
     data: tuple[Result, ...]
 
     def join(self, other: Result) -> Result:
@@ -119,7 +124,7 @@ def is_subseteq_Value(self, other: Value) -> bool:
 
 @final
 @dataclass
-class PartialLambda(Result):
+class PartialLambda(PartialConst):
     argnames: list[str]
     code: ir.Statement
     captured: tuple[Result, ...]

src/kirin/analysis/typeinfer/analysis.py

@@ -1,4 +1,7 @@
+from typing import TypeGuard
+
 from kirin import ir, types, interp
+from kirin.analysis import const
 from kirin.interp.impl import Signature
 from kirin.analysis.forward import Forward, ForwardFrame
 
@@ -17,14 +20,16 @@ def build_signature(
     ) -> Signature:
         _args = ()
         for x in frame.get_values(stmt.args):
-            if isinstance(x, types.Hinted):
-                _args += (x.type,)
-            elif isinstance(x, types.Generic):
-                _args += (x.body,)
-            else:
-                _args += (x,)
+            _args += (self._unwrap(x),)
         return Signature(stmt.__class__, _args)
 
+    def _unwrap(self, value: types.TypeAttribute) -> types.TypeAttribute:
+        if isinstance(value, types.Hinted):
+            return self._unwrap(value.type)
+        elif isinstance(value, types.Generic):
+            return value.body
+        return value
+
     def eval_stmt(
         self, frame: ForwardFrame[types.TypeAttribute, None], stmt: ir.Statement
     ) -> interp.StatementResult[types.TypeAttribute]:
@@ -46,3 +51,19 @@ def run_method(
                 method.code, (types.Hinted(types.PyClass(ir.Method), method),) + args
             )
         return types.Bottom
+
+    def is_const(
+        self, value: types.TypeAttribute
+    ) -> TypeGuard[types.Hinted[const.Value]]:
+        return isinstance(value, types.Hinted) and isinstance(value.data, const.Value)
+
+    def is_partial_const(
+        self, value: types.TypeAttribute
+    ) -> TypeGuard[
+        types.Hinted[const.Value]
+        | types.Hinted[const.PartialTuple]
+        | types.Hinted[const.PartialLambda]
+    ]:
+        return isinstance(value, types.Hinted) and isinstance(
+            value.data, (const.Value, const.PartialConst)
+        )

src/kirin/dialects/fcf/typeinfer.py

@@ -37,10 +37,10 @@ def fold(
         stmt: Foldl | Foldr,
         values: tuple[ir.types.TypeAttribute, ...],
     ):
-        if not isinstance(values[0], ir.types.Hinted):
+        if not interp.is_const(values[0]):
             return (stmt.result.type,)  # give up on dynamic calls
 
-        fn: ir.Method = values[0].data
+        fn: ir.Method = values[0].data.data
         coll: ir.types.TypeAttribute = values[1]
         init: ir.types.TypeAttribute = values[2]
 
@@ -71,10 +71,10 @@ def map_list(
         stmt: Map,
     ):
         fn_value = frame.get(stmt.fn)
-        if not isinstance(fn_value, ir.types.Hinted):
+        if not interp.is_const(fn_value):
             return (ir.types.List[ir.types.Any],)  # give up on dynamic calls
 
-        fn: ir.Method = fn_value.data
+        fn: ir.Method = fn_value.data.data
         coll: ir.types.TypeAttribute = frame.get(stmt.coll)
         if isinstance(coll, ir.types.Generic) and coll.is_subseteq(ir.types.List):
             elem = interp.eval(fn, (coll.vars[0],)).value
@@ -92,10 +92,10 @@ def map_range(
         stmt: Map,
     ):
         fn_value = frame.get(stmt.fn)
-        if not isinstance(fn_value, ir.types.Hinted):
+        if not interp.is_const(fn_value):
             return (ir.types.List,)  # give up on dynamic calls
 
-        fn: ir.Method = fn_value.data
+        fn: ir.Method = fn_value.data.data
         elem = interp.eval(fn, (ir.types.Int,)).value
         # fn errors forward the error
         if isinstance(elem, Err):
@@ -114,10 +114,10 @@ def scan(
         init = frame.get(stmt.init)
         coll = frame.get(stmt.coll)
 
-        if not isinstance(fn_value, ir.types.Hinted):
+        if not interp.is_const(fn_value):
             return (ir.types.Tuple[init, ir.types.List[ir.types.Any]],)
 
-        fn: ir.Method = fn_value.data
+        fn: ir.Method = fn_value.data.data
         if isinstance(coll, ir.types.Generic) and coll.is_subseteq(ir.types.List):
             ret = interp.eval(fn, (init, coll.vars[0])).value
             if isinstance(ret, Err):

src/kirin/dialects/func/typeinfer.py

@@ -30,10 +30,10 @@ def return_(self, interp: TypeInference, frame: Frame, stmt: Return) -> ReturnVa
     def call(self, interp: TypeInference, frame: Frame, stmt: Call):
         # give up on dynamic method calls
         callee = frame.get(stmt.callee)
-        if not isinstance(callee, ir.types.Hinted):
+        if not interp.is_const(callee):
             return (stmt.result.type,)
 
-        mt: ir.Method = callee.data
+        mt: ir.Method = callee.data.data
         return self._invoke_method(
             interp,
             mt,

src/kirin/dialects/py/stmts/_stmts/constant.py

@@ -23,7 +23,7 @@ def __init__(self, value: T | data.PyAttr[T]) -> None:
             value = data.PyAttr(value)
         super().__init__(
             properties={"value": value},
-            result_types=(types.Hinted(value.type, value.data),),
+            result_types=(value.type,),
         )
 
     def print_impl(self, printer: Printer) -> None:

src/kirin/dialects/py/stmts/typeinfer.py

@@ -1,5 +1,6 @@
 from kirin.ir import types
 from kirin.interp import Frame, MethodTable, StatementResult, impl
+from kirin.analysis.typeinfer import TypeInference
 
 from . import _stmts as py
 from .dialect import dialect
@@ -175,12 +176,12 @@ def mat_mult(self, interp, frame, stmt) -> StatementResult[types.TypeAttribute]:
     @impl(py.GetItem)
     def getitem(
         self,
-        interp,
+        interp: TypeInference,
         frame: Frame[types.TypeAttribute],
         stmt: py.GetItem,
     ) -> StatementResult[types.TypeAttribute]:
         obj = frame.get(stmt.obj)
-        if isinstance(obj, types.Hinted):  # unwrap const
+        if interp.is_const(obj):  # unwrap const
             obj = obj.type
         index: types.TypeAttribute = frame.get(stmt.index)
         # TODO: replace this when we can multiple dispatch
@@ -221,30 +222,31 @@ def getitem_tuple(
 
     def getitem_tuple_index(
         self,
-        interp,
+        interp: TypeInference,
         stmt: py.GetItem,
         obj: types.Generic,
         index: types.TypeAttribute,
     ):
-        if isinstance(index, types.Hinted):  # const
-            if obj.vararg and index.data >= len(obj.vars):
+        if interp.is_const(index) and index.type.is_subseteq(types.Int):
+            index_: int = index.data.data
+            if obj.vararg and index_ >= len(obj.vars):
                 return (obj.vararg.typ,)
-            elif index.data < len(obj.vars):
-                return (obj.vars[index.data],)
+            elif index_ < len(obj.vars):
+                return (obj.vars[index_],)
             else:
                 return (types.Bottom,)
         else:
             return (self.getitem_tuple_union(obj),)
 
     def getitem_tuple_slice(
         self,
-        interp,
+        interp: TypeInference,
         stmt: py.GetItem,
         obj: types.Generic,
         index: types.TypeAttribute,
     ):
-        if isinstance(index, types.Hinted):
-            data: slice = index.data
+        if interp.is_const(index):
+            data: slice = index.data.data
             if obj.vararg and data.stop >= len(obj.vars):
                 return (
                     types.Union(
@@ -315,17 +317,15 @@ def new_list(
 
     @impl(py.Slice)
     def slice(
-        self, interp, frame: Frame[types.TypeAttribute], stmt: py.Slice
+        self, interp: TypeInference, frame: Frame[types.TypeAttribute], stmt: py.Slice
     ) -> StatementResult[types.TypeAttribute]:
         start, stop, step = frame.get_values(stmt.args)
-        if (
-            isinstance(start, types.Hinted)
-            and isinstance(stop, types.Hinted)
-            and isinstance(step, types.Hinted)
-            and isinstance(stmt.result.type, types.TypeAttribute)
-        ):
+        if interp.is_const(start) and interp.is_const(stop) and interp.is_const(step):
             return (
-                types.Hinted(stmt.result.type, slice(start.data, stop.data, step.data)),
+                types.Hinted(
+                    stmt.result.type,
+                    slice(start.data.data, stop.data.data, step.data.data),
+                ),
             )
 
         return (stmt.result.type,)

src/kirin/emit/abc.pyi

@@ -37,7 +37,7 @@ class EmitABC(interp.BaseInterpreter[FrameType, ValueType]):
     def emit_type_TypeVar(self, attr: ir.types.TypeVar) -> ValueType: ...
     def emit_type_Vararg(self, attr: ir.types.Vararg) -> ValueType: ...
     def emit_type_Generic(self, attr: ir.types.Generic) -> ValueType: ...
-    def emit_type_Const(self, attr: ir.types.Hinted) -> ValueType: ...
+    def emit_type_Hinted(self, attr: ir.types.Hinted) -> ValueType: ...
     def emit_type_PyClass(self, attr: ir.types.PyClass) -> ValueType: ...
     def emit_attribute_fallback(self, attr: ir.Attribute) -> ValueType: ...
     def emit_stmt_begin(self, frame: FrameType, stmt: ir.Statement) -> None: ...

src/kirin/ir/_types.pyi

@@ -16,7 +16,7 @@ class _TypeAttribute(Attribute):
     def is_subseteq_Union(self, other: Union) -> bool: ...
     def is_subseteq_Literal(self, other: Literal) -> bool: ...
     def is_subseteq_TypeVar(self, other: TypeVar) -> bool: ...
-    def is_subseteq_Const(self, other: Hinted) -> bool: ...
+    def is_subseteq_Hinted(self, other: Hinted) -> bool: ...
     def is_subseteq_PyClass(self, other: PyClass) -> bool: ...
     def is_subseteq_Generic(self, other: Generic) -> bool: ...
     def is_subseteq_fallback(self, other: TypeAttribute) -> bool: ...

src/kirin/ir/types.py

@@ -68,6 +68,9 @@ def join(self, other: "TypeAttribute") -> "TypeAttribute":
             return Union(self, other)
         return AnyType()  # don't know how to join
 
+    def is_subseteq_Hinted(self, other: "Hinted") -> bool:
+        return self.is_subseteq(other.type)
+
     def print_impl(self, printer: Printer) -> None:
         printer.print_name(self, prefix="!")
 
@@ -162,9 +165,6 @@ def is_subseteq_Generic(self, other: "Generic") -> bool:
     def is_subseteq_TypeVar(self, other: "TypeVar") -> bool:
         return self.is_subseteq(other.bound)
 
-    def is_subseteq_Const(self, other: "Hinted") -> bool:
-        return self.is_subseteq(other.type)
-
     def __hash__(self) -> int:
         return hash((PyClass, self.typ))