generalize Const -> Hinted (#136)

This PR generalize `ir.types.Const` to `ir.types.Hinted` allow more
flexible data being stored in the type annotation to assist type
inference, e.g inference results from const propagation.

Author: Roger-luo

src/kirin/analysis/typeinfer/analysis.py

@@ -17,8 +17,8 @@ def build_signature(
     ) -> Signature:
         _args = ()
         for x in frame.get_values(stmt.args):
-            if isinstance(x, types.Const):
-                _args += (x.typ,)
+            if isinstance(x, types.Hinted):
+                _args += (x.type,)
             elif isinstance(x, types.Generic):
                 _args += (x.body,)
             else:
@@ -43,6 +43,6 @@ def run_method(
         if len(self.state.frames) < self.max_depth:
             # NOTE: widen method type here
             return self.run_callable(
-                method.code, (types.Const(method, types.PyClass(ir.Method)),) + args
+                method.code, (types.Hinted(types.PyClass(ir.Method), method),) + args
             )
         return types.Bottom

src/kirin/dialects/fcf/typeinfer.py

@@ -37,7 +37,7 @@ def fold(
         stmt: Foldl | Foldr,
         values: tuple[ir.types.TypeAttribute, ...],
     ):
-        if not isinstance(values[0], ir.types.Const):
+        if not isinstance(values[0], ir.types.Hinted):
             return (stmt.result.type,)  # give up on dynamic calls
 
         fn: ir.Method = values[0].data
@@ -71,7 +71,7 @@ def map_list(
         stmt: Map,
     ):
         fn_value = frame.get(stmt.fn)
-        if not isinstance(fn_value, ir.types.Const):
+        if not isinstance(fn_value, ir.types.Hinted):
             return (ir.types.List[ir.types.Any],)  # give up on dynamic calls
 
         fn: ir.Method = fn_value.data
@@ -92,7 +92,7 @@ def map_range(
         stmt: Map,
     ):
         fn_value = frame.get(stmt.fn)
-        if not isinstance(fn_value, ir.types.Const):
+        if not isinstance(fn_value, ir.types.Hinted):
             return (ir.types.List,)  # give up on dynamic calls
 
         fn: ir.Method = fn_value.data
@@ -114,7 +114,7 @@ def scan(
         init = frame.get(stmt.init)
         coll = frame.get(stmt.coll)
 
-        if not isinstance(fn_value, ir.types.Const):
+        if not isinstance(fn_value, ir.types.Hinted):
             return (ir.types.Tuple[init, ir.types.List[ir.types.Any]],)
 
         fn: ir.Method = fn_value.data

src/kirin/dialects/func/typeinfer.py

@@ -30,7 +30,7 @@ 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.Const):
+        if not isinstance(callee, ir.types.Hinted):
             return (stmt.result.type,)
 
         mt: ir.Method = callee.data

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.Const(value.data, value.type),),
+            result_types=(types.Hinted(value.type, value.data),),
         )
 
     def print_impl(self, printer: Printer) -> None:

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

@@ -26,9 +26,9 @@ def __init__(self, start: SSAValue, stop: SSAValue, step: SSAValue) -> None:
             if stop.type.is_subseteq(types.NoneType):
                 result_type = types.Bottom
             else:
-                result_type = types.Slice[types.widen_const(stop.type)]
+                result_type = types.Slice[types.unwrap_hinted(stop.type)]
         else:
-            result_type = types.Slice[types.widen_const(start.type)]
+            result_type = types.Slice[types.unwrap_hinted(start.type)]
 
         super().__init__(
             args=(start, stop, step),

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

@@ -180,8 +180,8 @@ def getitem(
         stmt: py.GetItem,
     ) -> StatementResult[types.TypeAttribute]:
         obj = frame.get(stmt.obj)
-        if isinstance(obj, types.Const):  # unwrap const
-            obj = obj.typ
+        if isinstance(obj, types.Hinted):  # unwrap const
+            obj = obj.type
         index: types.TypeAttribute = frame.get(stmt.index)
         # TODO: replace this when we can multiple dispatch
         if obj.is_subseteq(types.Tuple):
@@ -226,7 +226,7 @@ def getitem_tuple_index(
         obj: types.Generic,
         index: types.TypeAttribute,
     ):
-        if isinstance(index, types.Const):  # const
+        if isinstance(index, types.Hinted):  # const
             if obj.vararg and index.data >= len(obj.vars):
                 return (obj.vararg.typ,)
             elif index.data < len(obj.vars):
@@ -243,7 +243,7 @@ def getitem_tuple_slice(
         obj: types.Generic,
         index: types.TypeAttribute,
     ):
-        if isinstance(index, types.Const):
+        if isinstance(index, types.Hinted):
             data: slice = index.data
             if obj.vararg and data.stop >= len(obj.vars):
                 return (
@@ -319,13 +319,13 @@ def slice(
     ) -> StatementResult[types.TypeAttribute]:
         start, stop, step = frame.get_values(stmt.args)
         if (
-            isinstance(start, types.Const)
-            and isinstance(stop, types.Const)
-            and isinstance(step, types.Const)
+            isinstance(start, types.Hinted)
+            and isinstance(stop, types.Hinted)
+            and isinstance(step, types.Hinted)
             and isinstance(stmt.result.type, types.TypeAttribute)
         ):
             return (
-                types.Const(slice(start.data, stop.data, step.data), stmt.result.type),
+                types.Hinted(stmt.result.type, slice(start.data, stop.data, step.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.Const) -> ValueType: ...
+    def emit_type_Const(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

@@ -2,8 +2,8 @@ from dataclasses import dataclass
 
 from kirin.ir.attrs import Attribute
 from kirin.ir.types import (
-    Const,
     Union,
+    Hinted,
     Generic,
     Literal,
     PyClass,
@@ -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: Const) -> bool: ...
+    def is_subseteq_Const(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

@@ -162,8 +162,8 @@ 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: "Const") -> bool:
-        return self.is_subseteq(other.typ)
+    def is_subseteq_Const(self, other: "Hinted") -> bool:
+        return self.is_subseteq(other.type)
 
     def __hash__(self) -> int:
         return hash((PyClass, self.typ))
@@ -349,7 +349,7 @@ def __init__(
             self.body = body
         else:
             self.body = PyClass(body)
-        self.vars, self.vararg = split_type_args(vars)
+        self.vars, self.vararg = _split_type_args(vars)
 
     def is_subseteq_Literal(self, other: Literal) -> bool:
         return False
@@ -408,7 +408,7 @@ def where(self, typ: TypeVarValue | tuple[TypeVarValue, ...]) -> "Generic":
         else:
             typs = (typ,)
 
-        args, vararg = split_type_args(typs)
+        args, vararg = _split_type_args(typs)
         if self.vararg is None and vararg is None:
             assert len(args) <= len(
                 self.vars
@@ -446,45 +446,52 @@ def where(self, typ: TypeVarValue | tuple[TypeVarValue, ...]) -> "Generic":
         raise TypeError("Type arguments do not match")
 
 
-ConstType = typing.TypeVar("ConstType")
+HintedData = typing.TypeVar("HintedData")
 
 
 @typing.final
 @dataclass
-class Const(TypeAttribute, typing.Generic[ConstType]):
-    name = "Const"
-    data: ConstType
-    typ: TypeAttribute
+class Hinted(TypeAttribute, typing.Generic[HintedData]):
+    """Type wrapped with a hint.
+
+    `Hinted` is used to represent a type with additional data that can be used as
+    a hint for type inference. The additional data is only used for specific type
+    inference purposes, or improve certain type inference precision, it does not affect the
+    order of types in the lattice.
+    """
 
-    def __init__(self, data: ConstType, typ: TypeAttribute | None = None):
+    name = "Hinted"
+    type: TypeAttribute
+    data: HintedData
+
+    def __init__(self, type: TypeAttribute, data: HintedData):
         self.data = data
-        if isinstance(typ, Const):
-            typ = widen_const(typ)
-        elif typ is None:
-            typ = PyClass(type(data))
-        self.typ = typ
+        if isinstance(type, Hinted):
+            type = type.type
+        self.type = type
 
     def is_equal(self, other: TypeAttribute) -> bool:
         return (
-            isinstance(other, Const)
+            isinstance(other, Hinted)
             and self.data == other.data
-            and self.typ.is_equal(other.typ)
+            and self.type.is_equal(other.type)
         )
 
     def is_subseteq_fallback(self, other: TypeAttribute) -> bool:
-        return self.typ.is_subseteq(other)
+        return self.type.is_subseteq(other)
 
     def __hash__(self) -> int:
-        return hash(self.typ)
+        return hash(self.type)
 
     def print_impl(self, printer: Printer) -> None:
         printer.print_name(self, prefix="!")
-        printer.plain_print("(", self.data, ", ")
-        printer.print(self.typ)
+        printer.plain_print("(")
+        printer.print(self.type)
+        printer.plain_print(", ", self.data)
         printer.plain_print(")")
 
 
-def split_type_args(
+def _split_type_args(
     args: tuple[TypeVarValue, ...]
 ) -> tuple[tuple[TypeAttribute, ...], Vararg | None]:
     if args is None or len(args) == 0:
@@ -539,10 +546,10 @@ def hint2type(hint) -> TypeAttribute:
     return Generic(body, *params)
 
 
-def widen_const(typ: TypeAttribute) -> TypeAttribute:
-    if isinstance(typ, Const):
-        return typ.typ
-    return typ
+def unwrap_hinted(hint: TypeAttribute) -> TypeAttribute:
+    if isinstance(hint, Hinted):
+        return hint.type
+    return hint
 
 
 Any = AnyType()

test/analysis/dataflow/typeinfer/test_unstable.py

@@ -25,17 +25,17 @@ def results_at(block_id, stmt_id):
         return stmt_at(block_id, stmt_id).results
 
     assert [infer.results[result] for result in results_at(0, 0)] == [
-        types.Const(1, types.Int)
+        types.Hinted(types.Int, 1)
     ]
     assert [infer.results[result] for result in results_at(0, 1)] == [types.Int]
     assert [infer.results[result] for result in results_at(0, 2)] == [
-        types.Const(10, types.Int)
+        types.Hinted(types.Int, 10)
     ]
     assert [infer.results[result] for result in results_at(0, 3)] == [types.Bool]
 
     assert [infer.results[result] for result in results_at(1, 0)] == [types.Int]
     assert [infer.results[result] for result in results_at(2, 0)] == [
-        types.Const(1.2, types.Float)
+        types.Hinted(types.Float, 1.2)
     ]
     assert [infer.results[result] for result in results_at(2, 1)] == [types.Float]