fix inference with scf (#241)

this PR fixes inference when there are `func.Return` inside the `scf`
region.

Author: Roger-luo

src/kirin/analysis/typeinfer/analysis.py

@@ -24,6 +24,13 @@ class TypeInference(Forward[types.TypeAttribute]):
     keys = ["typeinfer"]
     lattice = types.TypeAttribute
 
+    def run_analysis(
+        self, method: ir.Method, args: tuple[types.TypeAttribute, ...] | None = None
+    ) -> tuple[ForwardFrame[types.TypeAttribute], types.TypeAttribute]:
+        if args is None:
+            args = method.arg_types
+        return super().run_analysis(method, args)
+
     # NOTE: unlike concrete interpreter, instead of using type information
     # within the IR. Type inference will use the interpreted
     # value (which is a type) to determine the method dispatch.

src/kirin/dialects/eltype.py

@@ -17,7 +17,7 @@ class ElType(ir.Statement):
     This statement is used by other dialects to query the element type of a value.
     """
 
-    container: ir.SSAValue = info.argument(types.PyClass(types.TypeAttribute))
+    container: ir.SSAValue = info.argument(types.Any)
     """The value to query the element type of."""
     elem: ir.ResultValue = info.result(types.PyClass(types.TypeAttribute))
     """The element type of the value."""

src/kirin/dialects/func/typeinfer.py

@@ -2,6 +2,7 @@
 
 from kirin import ir, types
 from kirin.interp import Frame, MethodTable, ReturnValue, impl
+from kirin.analysis import const
 from kirin.analysis.typeinfer import TypeInference, TypeResolution
 from kirin.dialects.func.stmts import (
     Call,
@@ -24,6 +25,8 @@ def const_none(self, interp: TypeInference, frame: Frame, stmt: ConstantNone):
 
     @impl(Return)
     def return_(self, interp: TypeInference, frame: Frame, stmt: Return) -> ReturnValue:
+        if isinstance(hint := stmt.value.hints.get("const"), const.Value):
+            return ReturnValue(types.Literal(hint.data))
         return ReturnValue(frame.get(stmt.value))
 
     @impl(Call)

src/kirin/dialects/ilist/typeinfer.py

@@ -1,5 +1,6 @@
 from kirin import types
 from kirin.interp import Frame, MethodTable, impl
+from kirin.dialects.eltype import ElType
 from kirin.dialects.py.binop import Add
 from kirin.analysis.typeinfer import TypeInference
 from kirin.dialects.py.indexing import GetItem
@@ -19,6 +20,16 @@ def _get_list_len(typ: types.Generic):
         else:
             return types.Any
 
+    @impl(ElType, types.PyClass(IList))
+    def eltype_list(
+        self, interp: TypeInference, frame: Frame[types.TypeAttribute], stmt: ElType
+    ):
+        list_type = frame.get(stmt.container)
+        if isinstance(list_type, types.Generic):
+            return (list_type.vars[0],)
+        else:
+            return (types.Any,)
+
     @impl(New)
     def new(self, interp: TypeInference, frame: Frame[types.TypeAttribute], stmt: New):
         values = frame.get_values(stmt.values)

src/kirin/dialects/py/list/typeinfer.py

@@ -1,4 +1,5 @@
 from kirin import types, interp
+from kirin.dialects.eltype import ElType
 from kirin.dialects.py.binop import Add
 from kirin.dialects.py.indexing import GetItem
 
@@ -8,6 +9,14 @@
 @dialect.register(key="typeinfer")
 class TypeInfer(interp.MethodTable):
 
+    @interp.impl(ElType, types.PyClass(list))
+    def eltype_list(self, interp, frame: interp.Frame, stmt: ElType):
+        list_type = frame.get(stmt.container)
+        if isinstance(list_type, types.Generic):
+            return (list_type.vars[0],)
+        else:
+            return (types.Any,)
+
     @interp.impl(Add, types.PyClass(list), types.PyClass(list))
     def add(self, interp, frame: interp.Frame, stmt: Add):
         lhs_type = frame.get(stmt.lhs)

src/kirin/dialects/py/range.py

@@ -14,8 +14,9 @@
 import ast
 from dataclasses import dataclass
 
-from kirin import ir, types, lowering, exceptions
+from kirin import ir, types, interp, lowering, exceptions
 from kirin.decl import info, statement
+from kirin.dialects import eltype
 
 dialect = ir.Dialect("py.range")
 
@@ -48,6 +49,14 @@ def lower_Call_range(
         return _lower_range(state, node)
 
 
+@dialect.register(key="typeinfer")
+class TypeInfer(interp.MethodTable):
+
+    @interp.impl(eltype.ElType, types.PyClass(range))
+    def eltype_range(self, interp_, frame: interp.Frame, stmt: eltype.ElType):
+        return (types.Int,)
+
+
 def _lower_range(state: lowering.LoweringState, node: ast.Call) -> lowering.Result:
     if len(node.args) == 1:
         start = state.visit(ast.Constant(0)).expect_one()

src/kirin/dialects/py/tuple.py

@@ -18,6 +18,7 @@
 from kirin.decl import info, statement
 from kirin.analysis import const
 from kirin.emit.julia import EmitJulia, EmitStrFrame
+from kirin.dialects.eltype import ElType
 from kirin.dialects.py.binop import Add
 
 dialect = ir.Dialect("py.tuple")
@@ -49,6 +50,17 @@ def new(self, interp: interp.Interpreter, frame: interp.Frame, stmt: New):
 @dialect.register(key="typeinfer")
 class TypeInfer(interp.MethodTable):
 
+    @interp.impl(ElType, types.PyClass(tuple))
+    def eltype_tuple(self, interp, frame: interp.Frame, stmt: ElType):
+        tuple_type = frame.get(stmt.container)
+        if isinstance(tuple_type, types.Generic):
+            ret = tuple_type.vars[0]
+            for var in tuple_type.vars[1:]:
+                ret = ret.join(var)
+            return (ret,)
+        else:
+            return (types.Any,)
+
     @interp.impl(Add, types.PyClass(tuple), types.PyClass(tuple))
     def add(self, interp, frame: interp.Frame[types.TypeAttribute], stmt):
         lhs = frame.get(stmt.lhs)

src/kirin/dialects/scf/constprop.py

@@ -1,6 +1,6 @@
 from collections.abc import Iterable
 
-from kirin import interp
+from kirin import ir, interp
 from kirin.analysis import const
 
 from .stmts import For, Yield, IfElse
@@ -32,35 +32,50 @@ def if_else(
     ):
         cond = frame.get(stmt.cond)
         if isinstance(cond, const.Value):
-            with interp_.state.new_frame(interp_.new_frame(stmt)) as body_frame:
-                body_frame.entries.update(frame.entries)
-                if cond.data:
-                    results = interp_.run_ssacfg_region(body_frame, stmt.then_body)
-                else:
-                    results = interp_.run_ssacfg_region(body_frame, stmt.else_body)
+            if cond.data:
+                body = stmt.then_body
+            else:
+                body = stmt.else_body
+            body_frame, ret = self._prop_const_cond_ifelse(
+                interp_, frame, stmt, cond, body
+            )
+            frame.entries.update(body_frame.entries)
+            return ret
+        else:
+            then_frame, then_results = self._prop_const_cond_ifelse(
+                interp_, frame, stmt, const.Value(True), stmt.then_body
+            )
+            else_frame, else_results = self._prop_const_cond_ifelse(
+                interp_, frame, stmt, const.Value(False), stmt.else_body
+            )
+            ret = interp_.join_results(then_results, else_results)
 
-            if not body_frame.frame_is_not_pure:
+            if not then_frame.frame_is_not_pure or not else_frame.frame_is_not_pure:
                 frame.should_be_pure.add(stmt)
-        else:
-            with interp_.state.new_frame(interp_.new_frame(stmt)) as then_body_frame:
-                then_body_frame.entries.update(frame.entries)
-                then_results = interp_.run_ssacfg_region(
-                    then_body_frame, stmt.then_body
-                )
 
-            with interp_.state.new_frame(interp_.new_frame(stmt)) as else_body_frame:
-                else_body_frame.entries.update(frame.entries)
-                else_results = interp_.run_ssacfg_region(
-                    else_body_frame, stmt.else_body
-                )
-            results = interp_.join_results(then_results, else_results)
+            # NOTE: then_frame and else_frame do not change
+            # parent frame variables value except cond
+            frame.entries.update(then_frame.entries)
+            frame.entries.update(else_frame.entries)
+            frame.set(stmt.cond, cond)
+        return ret
 
-            if (
-                not then_body_frame.frame_is_not_pure
-                or not else_body_frame.frame_is_not_pure
-            ):
-                frame.should_be_pure.add(stmt)
-        return results
+    def _prop_const_cond_ifelse(
+        self,
+        interp_: const.Propagate,
+        frame: const.Frame,
+        stmt: IfElse,
+        cond: const.Value,
+        body: ir.Region,
+    ):
+        with interp_.state.new_frame(interp_.new_frame(stmt)) as body_frame:
+            body_frame.entries.update(frame.entries)
+            body_frame.set(body.blocks[0].args[0], cond)
+            results = interp_.run_ssacfg_region(body_frame, body)
+
+        if not body_frame.frame_is_not_pure:
+            frame.should_be_pure.add(stmt)
+        return body_frame, results
 
     @interp.impl(For)
     def for_loop(
@@ -70,33 +85,44 @@ def for_loop(
         stmt: For,
     ):
         iterable = frame.get(stmt.iterable)
-        loop_vars = frame.get_values(stmt.initializers)
-        block_args = stmt.body.blocks[0].args
-
         if isinstance(iterable, const.Value):
-            frame_is_not_pure = False
-            if not isinstance(iterable.data, Iterable):
-                raise interp.InterpreterError(
-                    f"Expected iterable, got {type(iterable.data)}"
-                )
-            for value in iterable.data:
-                with interp_.state.new_frame(interp_.new_frame(stmt)) as body_frame:
-                    body_frame.entries.update(frame.entries)
-                    body_frame.set_values(
-                        block_args,
-                        (const.Value(value),) + loop_vars,
-                    )
-                    loop_vars = interp_.run_ssacfg_region(body_frame, stmt.body)
-
-                if body_frame.frame_is_not_pure:
-                    frame_is_not_pure = True
-                if loop_vars is None:
-                    loop_vars = ()
-                elif isinstance(loop_vars, interp.ReturnValue):
-                    return loop_vars
-
-            if not frame_is_not_pure:
-                frame.should_be_pure.add(stmt)
-            return loop_vars
+            return self._prop_const_iterable_forloop(interp_, frame, stmt, iterable)
         else:  # TODO: support other iteration
             return tuple(interp_.lattice.top() for _ in stmt.results)
+
+    def _prop_const_iterable_forloop(
+        self,
+        interp_: const.Propagate,
+        frame: const.Frame,
+        stmt: For,
+        iterable: const.Value,
+    ):
+        frame_is_not_pure = False
+        if not isinstance(iterable.data, Iterable):
+            raise interp.InterpreterError(
+                f"Expected iterable, got {type(iterable.data)}"
+            )
+
+        loop_vars = frame.get_values(stmt.initializers)
+        body_block = stmt.body.blocks[0]
+        block_args = body_block.args
+
+        for value in iterable.data:
+            with interp_.state.new_frame(interp_.new_frame(stmt)) as body_frame:
+                body_frame.entries.update(frame.entries)
+                body_frame.set_values(
+                    block_args,
+                    (const.Value(value),) + loop_vars,
+                )
+                loop_vars = interp_.run_ssacfg_region(body_frame, stmt.body)
+
+            if body_frame.frame_is_not_pure:
+                frame_is_not_pure = True
+            if loop_vars is None:
+                loop_vars = ()
+            elif isinstance(loop_vars, interp.ReturnValue):
+                return loop_vars
+
+        if not frame_is_not_pure:
+            frame.should_be_pure.add(stmt)
+        return loop_vars

src/kirin/dialects/scf/lowering.py

@@ -15,13 +15,19 @@ def lower_If(self, state: lowering.LoweringState, node: ast.If) -> lowering.Resu
         cond = state.visit(node.test).expect_one()
         frame = state.current_frame
         body_frame = lowering.Frame.from_stmts(node.body, state, globals=frame.globals)
+        then_cond = body_frame.curr_block.args.append_from(types.Bool, cond.name)
+        if cond.name:
+            body_frame.defs[cond.name] = then_cond
         state.push_frame(body_frame)
         state.exhaust(body_frame)
         state.pop_frame(finalize_next=False)  # NOTE: scf does not have multiple blocks
 
         else_frame = lowering.Frame.from_stmts(
             node.orelse, state, globals=frame.globals
         )
+        else_cond = else_frame.curr_block.args.append_from(types.Bool, cond.name)
+        if cond.name:
+            else_frame.defs[cond.name] = else_cond
         state.push_frame(else_frame)
         state.exhaust(else_frame)
         state.pop_frame(finalize_next=False)  # NOTE: scf does not have multiple blocks
@@ -96,7 +102,11 @@ def new_block_arg_if_inside_loop(frame: lowering.Frame, capture: ir.SSAValue):
         unpacking(state, node.target, loop_var)
         state.exhaust(body_frame)
         # NOTE: this frame won't have phi nodes
-        body_frame.append_stmt(Yield(*[body_frame.defs[name] for name in yields]))  # type: ignore
+        if yields and (
+            body_frame.curr_block.last_stmt is None
+            or not body_frame.curr_block.last_stmt.has_trait(ir.IsTerminator)
+        ):
+            body_frame.append_stmt(Yield(*[body_frame.defs[name] for name in yields]))  # type: ignore
         state.pop_frame(finalize_next=False)  # NOTE: scf does not have multiple blocks
 
         initializers: list[ir.SSAValue] = []