support inference on generics (#129)

Author: Roger-luo

src/kirin/analysis/typeinfer/__init__.py

@@ -0,0 +1,2 @@
+from .solve import TypeResolution as TypeResolution
+from .analysis import TypeInference as TypeInference

src/kirin/analysis/typeinfer.py src/kirin/analysis/typeinfer/analysis.pysrc/kirin/analysis/typeinfer.py renamed to src/kirin/analysis/typeinfer/analysis.py

@@ -2,6 +2,8 @@
 from kirin.interp.impl import Signature
 from kirin.analysis.forward import Forward, ForwardFrame
 
+from .solve import TypeResolution
+
 
 class TypeInference(Forward[types.TypeAttribute]):
     keys = ["typeinfer"]
@@ -29,7 +31,11 @@ def eval_stmt(
         method = self.lookup_registry(frame, stmt)
         if method is not None:
             return method(self, frame, stmt)
-        return tuple(result.type for result in stmt.results)
+
+        resolve = TypeResolution()
+        for arg, value in zip(stmt.args, frame.get_values(stmt.args)):
+            resolve.solve(arg.type, value)
+        return tuple(resolve.substitute(result.type) for result in stmt.results)
 
     def run_method(
         self, method: ir.Method, args: tuple[types.TypeAttribute, ...]

src/kirin/analysis/typeinfer/solve.py

@@ -0,0 +1,101 @@
+from dataclasses import field, dataclass
+
+from kirin.ir import types
+
+
+@dataclass
+class TypeResolutionResult:
+    pass
+
+
+@dataclass
+class ResolutionOk(TypeResolutionResult):
+
+    def __bool__(self):
+        return True
+
+
+Ok = ResolutionOk()
+
+
+@dataclass
+class ResolutionError(TypeResolutionResult):
+    expr: types.TypeAttribute
+    value: types.TypeAttribute
+
+    def __bool__(self):
+        return False
+
+    def __str__(self):
+        return f"expected {self.expr}, got {self.value}"
+
+
+@dataclass
+class TypeResolution:
+    vars: dict[types.TypeVar, types.TypeAttribute] = field(default_factory=dict)
+
+    def substitute(self, typ: types.TypeAttribute) -> types.TypeAttribute:
+        if isinstance(typ, types.TypeVar):
+            return self.vars.get(typ, typ)
+        elif isinstance(typ, types.Generic):
+            return types.Generic(
+                typ.body, *tuple(self.substitute(var) for var in typ.vars)
+            )
+        elif isinstance(typ, types.Union):
+            return types.Union(self.substitute(t) for t in typ.types)
+        return typ
+
+    def solve(
+        self, annot: types.TypeAttribute, value: types.TypeAttribute
+    ) -> TypeResolutionResult:
+        if isinstance(annot, types.TypeVar):
+            return self.solve_TypeVar(annot, value)
+        elif isinstance(annot, types.Generic):
+            return self.solve_Generic(annot, value)
+        elif isinstance(annot, types.Union):
+            return self.solve_Union(annot, value)
+
+        if annot.is_subseteq(value):
+            return Ok
+        else:
+            return ResolutionError(annot, value)
+
+    def solve_TypeVar(self, annot: types.TypeVar, value: types.TypeAttribute):
+        if annot in self.vars:
+            if value.is_subseteq(self.vars[annot]):
+                self.vars[annot] = value
+            elif self.vars[annot].is_subseteq(value):
+                pass
+            else:
+                return ResolutionError(annot, value)
+        else:
+            self.vars[annot] = value
+        return Ok
+
+    def solve_Generic(self, annot: types.Generic, value: types.TypeAttribute):
+        if not isinstance(value, types.Generic):
+            return ResolutionError(annot, value)
+
+        if not value.body.is_subseteq(annot.body):
+            return ResolutionError(annot.body, value.body)
+
+        for var, val in zip(annot.vars, value.vars):
+            result = self.solve(var, val)
+            if not result:
+                return result
+
+        if not annot.vararg:
+            return Ok
+
+        for val in value.vars[len(annot.vars) :]:
+            result = self.solve(annot.vararg.typ, val)
+            if not result:
+                return result
+        return Ok
+
+    def solve_Union(self, annot: types.Union, value: types.TypeAttribute):
+        for typ in annot.types:
+            result = self.solve(typ, value)
+            if result:
+                return Ok
+        return ResolutionError(annot, value)

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

@@ -21,12 +21,6 @@ def alias(
     ) -> StatementResult[types.TypeAttribute]:
         return (frame.get(stmt.value),)  # just forward the type
 
-    @impl(py.NewTuple)
-    def new_tuple(
-        self, interp, frame: Frame[types.TypeAttribute], stmt: py.NewTuple
-    ) -> StatementResult[types.TypeAttribute]:
-        return (types.Tuple.where(frame.get_values(stmt.args)),)  # make 3.10 happy
-
     @impl(py.Add, types.Float, types.Float)
     @impl(py.Add, types.Float, types.Int)
     @impl(py.Add, types.Int, types.Float)