remove return_type make return_type return signature output (#152)

Author: Roger-luo

src/kirin/ir/method.py

@@ -2,11 +2,10 @@
 from typing import TYPE_CHECKING, Generic, TypeVar, Callable, ParamSpec
 from dataclasses import field, dataclass
 
-from kirin.ir.traits import CallableStmtInterface
+from kirin.ir.traits import HasSignature, CallableStmtInterface
 from kirin.exceptions import InterpreterError, VerificationError
 from kirin.ir.nodes.stmt import Statement
 from kirin.print.printer import Printer
-from kirin.ir.attrs.types import TypeAttribute
 from kirin.print.printable import Printable
 
 if TYPE_CHECKING:
@@ -27,7 +26,6 @@ class Method(Printable, Generic[Param, RetType]):
     # values contained if closure
     fields: tuple = field(default_factory=tuple)  # own
     file: str = ""
-    return_type: TypeAttribute | None = None
     inferred: bool = False
     """if typeinfer has been run on this method
     """
@@ -58,6 +56,13 @@ def callable_region(self):
             raise ValueError("Method body must implement CallableStmtInterface")
         return trait.get_callable_region(self.code)
 
+    @property
+    def return_type(self):
+        trait = self.code.get_trait(HasSignature)
+        if trait is None:
+            raise ValueError("Method body must implement HasSignature")
+        return trait.get_signature(self.code).output
+
     def __repr__(self) -> str:
         return f'Method("{self.sym_name}")'
 

src/kirin/passes/typeinfer.py

@@ -1,9 +1,10 @@
 from dataclasses import dataclass
 
-from kirin.ir import Method
+from kirin.ir import Method, HasSignature
 from kirin.rewrite import Walk
 from kirin.passes.abc import Pass
 from kirin.rewrite.abc import RewriteResult
+from kirin.dialects.func import Signature
 from kirin.analysis.typeinfer import TypeInference
 from kirin.rewrite.apply_type import ApplyType
 
@@ -16,8 +17,10 @@ def __post_init__(self):
 
     def unsafe_run(self, mt: Method) -> RewriteResult:
         return_type = self.infer.eval(mt, mt.arg_types).expect()
-        mt.return_type = return_type
+        if trait := mt.code.get_trait(HasSignature):
+            trait.set_signature(mt.code, Signature(mt.arg_types, return_type))
+
+        result = Walk(ApplyType(self.infer.results)).rewrite(mt.code)
         mt.inferred = True
-        result = Walk(ApplyType(return_type, self.infer.results)).rewrite(mt.code)
         self.infer.results.clear()
         return result

src/kirin/rewrite/apply_type.py

@@ -1,53 +1,27 @@
 from dataclasses import dataclass
 
-from kirin.ir import (
-    Block,
-    SSAValue,
-    Statement,
-    HasSignature,
-    CallableStmtInterface,
-    types,
-)
+from kirin.ir import Block, SSAValue, Statement, types
 from kirin.rewrite.abc import RewriteRule, RewriteResult
-from kirin.dialects.func import Signature
 
 
 @dataclass
 class ApplyType(RewriteRule):
-    ret_type: types.TypeAttribute
     results: dict[SSAValue, types.TypeAttribute]
 
-    def get_type(self, value: SSAValue) -> types.TypeAttribute:
-        return self.results[value]
-        # if isinstance(typ, types.PyConst):
-        #     return typ.typ
-        # return typ
-
     def rewrite_Block(self, node: Block) -> RewriteResult:
         has_done_something = False
         for arg in node.args:
             if arg in self.results:
-                arg.type = self.get_type(arg)
+                arg.type = self.results[arg]
                 has_done_something = True
 
         return RewriteResult(has_done_something=has_done_something)
 
     def rewrite_Statement(self, node: Statement) -> RewriteResult:
-        if (fn_trait := node.get_trait(HasSignature)) and (
-            call_trait := node.get_trait(CallableStmtInterface)
-        ):
-            signature = fn_trait.get_signature(node)
-            body = call_trait.get_callable_region(node)
-            inputs = [
-                self.get_type(arg) if arg in self.results else signature.inputs[idx]
-                for idx, arg in enumerate(body.blocks[0].args[1:])
-            ]
-            fn_trait.set_signature(node, Signature(tuple(inputs), self.ret_type))
-
         has_done_something = False
         for result in node._results:
             if result in self.results:
-                result.type = self.get_type(result)
+                result.type = self.results[result]
                 has_done_something = True
 
         return RewriteResult(has_done_something=has_done_something)

test/analysis/dataflow/typeinfer/test_inter_method.py

@@ -27,4 +27,4 @@ def test_inter_method_infer():
     # NOTE: inference of moo should not update foo
     assert foo.arg_types[0] == types.Int
     assert foo.inferred is False
-    assert foo.return_type is None
+    assert foo.return_type is types.Any