Fix SCF loop type inference with bounded widening

src/kirin/analysis/typeinfer/widen.py

@@ -0,0 +1,26 @@
+from kirin import types
+
+
+def widen(
+    previous: types.TypeAttribute, current: types.TypeAttribute
+) -> types.TypeAttribute:
+    """Include both types, forgetting changing literals within matching generics."""
+    if current.is_subseteq(previous):
+        return previous
+    if previous.is_subseteq(types.Bottom):
+        return current
+    if (
+        isinstance(previous, types.Generic)
+        and isinstance(current, types.Generic)
+        and previous.body == current.body
+        and len(previous.vars) == len(current.vars)
+        and previous.vararg is None
+        and current.vararg is None
+    ):
+        return types.Generic(
+            previous.body,
+            *(widen(old, new) for old, new in zip(previous.vars, current.vars)),
+        )
+    if isinstance(previous, types.Literal) and isinstance(current, types.Literal):
+        return types.Any
+    return previous.join(current)

src/kirin/dialects/ilist/typeinfer.py

@@ -43,7 +43,7 @@ def eltype_list(
     def new(self, interp: TypeInference, frame: Frame[types.TypeAttribute], stmt: New):
         values = frame.get_values(stmt.values)
         if not values:
-            return (IListType[types.Any, types.Literal(0)],)
+            return (IListType[stmt.elem_type, types.Literal(0)],)
 
         elem_type = values[0]
         for v in values:

src/kirin/dialects/scf/typeinfer.py

@@ -2,11 +2,16 @@
 from kirin.analysis import ForwardFrame, TypeInference
 from kirin.dialects import func
 from kirin.dialects.eltype import ElType
+from kirin.analysis.typeinfer.widen import widen
 
 from . import absint
 from .stmts import For, IfElse
 from ._dialect import dialect
 
+# Precision/performance heuristic: allow structural widening before falling back
+# to top. Correctness requires a verified invariant regardless of this budget.
+_LOOP_WIDENING_BUDGET = 16
+
 
 @dialect.register(key="typeinfer")
 class TypeInfer(absint.Methods):
@@ -40,16 +45,48 @@ def for_loop(
         if not isinstance(eltype, tuple):  # error
             return
         item = eltype[0]
-        frame.set_values(block_args, (item,) + loop_vars)
 
         if isinstance(body_block.last_stmt, func.Return):
+            frame.set_values(block_args, (item,) + loop_vars)
             frame.worklist.append(interp.Successor(body_block, item, *loop_vars))
             return  # if terminate is Return, there is no result
 
-        loop_vars_ = interp_.frame_call_region(frame, stmt, stmt.body, item, *loop_vars)
-        if isinstance(loop_vars_, interp.ReturnValue):
-            return loop_vars_
-        elif isinstance(loop_vars_, tuple):
-            return interp_.join_results(loop_vars, loop_vars_)
-        else:  # None, loop has no result
-            return
+        candidate = loop_vars
+        # Each fallback update promotes at least one changing component to top.
+        # Allow one update per component and a final invariant check.
+        max_iterations = _LOOP_WIDENING_BUDGET + len(candidate) + 1
+        for iteration in range(max_iterations):
+            # A fresh frame prevents earlier visits and intermediate types from
+            # contaminating inference under the current loop invariant.
+            with interp_.new_frame(stmt, has_parent_access=True) as body_frame:
+                yielded = interp_.frame_call_region(
+                    body_frame, stmt, stmt.body, item, *candidate
+                )
+
+            if not isinstance(yielded, tuple):
+                frame.set_values(body_frame.entries.keys(), body_frame.entries.values())
+                return yielded
+
+            required = tuple(
+                initial.join(value) for initial, value in zip(loop_vars, yielded)
+            )
+            if all(new.is_subseteq(old) for old, new in zip(candidate, required)):
+                frame.set_values(body_frame.entries.keys(), body_frame.entries.values())
+                frame.set_values(block_args, (item,) + candidate)
+                return candidate
+
+            # Structural growth can also produce unbounded unions or nesting.
+            # Fall back to top for changing components, then verify the invariant
+            # with another body evaluation rather than returning a partial result.
+            candidate = tuple(
+                (
+                    widen(old, new)
+                    if iteration < _LOOP_WIDENING_BUDGET
+                    else old if new.is_subseteq(old) else types.Any
+                )
+                for old, new in zip(candidate, yielded)
+            )
+
+        raise interp.InterpreterError(
+            f"scf.For type inference did not converge after {max_iterations} iterations"
+        )

test/analysis/dataflow/typeinfer/test_widen.py

@@ -0,0 +1,36 @@
+from pytest import mark
+
+from kirin import types
+from kirin.dialects.ilist import IListType
+from kirin.analysis.typeinfer.widen import widen
+
+
+@mark.parametrize(
+    "previous,current,expected",
+    [
+        (types.Bottom, types.Int, types.Int),
+        (types.Any, types.Int, types.Any),
+        (types.Int, types.Float, types.Int | types.Float),
+        (
+            IListType[types.Int, types.Literal(1)],
+            IListType[types.Float, types.Literal(1)],
+            IListType[types.Int | types.Float, types.Literal(1)],
+        ),
+        (
+            IListType[types.Int, types.Literal(0)],
+            IListType[types.Int, types.Literal(1)],
+            IListType[types.Int, types.Any],
+        ),
+        (
+            types.Tuple[IListType[types.Int, types.Literal(0)], types.Bool],
+            types.Tuple[IListType[types.Int, types.Literal(1)], types.Bool],
+            types.Tuple[IListType[types.Int, types.Any], types.Bool],
+        ),
+    ],
+)
+def test_widen_preserves_upper_bounds(previous, current, expected):
+    result = widen(previous, current)
+    assert result == expected
+    assert previous.is_subseteq(result)
+    assert current.is_subseteq(result)
+    assert widen(result, current) == result

test/dialects/scf/test_loop_inference.py

@@ -0,0 +1,249 @@
+from typing import Any
+
+from pytest import mark
+
+from kirin import types
+from kirin.passes import TypeInfer
+from kirin.prelude import structural_no_opt
+from kirin.rewrite import Walk
+from kirin.analysis import TypeInference
+from kirin.dialects import scf, ilist
+
+
+@structural_no_opt
+def identity(x: int) -> int:
+    return x
+
+
+@mark.parametrize("n", [0, 1, 3])
+def test_unroll_map_in_growing_loop(n):
+    @structural_no_opt
+    def grow(n: int) -> ilist.IList[int, Any]:
+        xs = ilist.IList([])
+        for _ in range(n):
+            xs = xs + [4]
+            xs = ilist.map(identity, xs)
+        return xs
+
+    expected = list(grow(n))
+    TypeInfer(grow.dialects, no_raise=False)(grow)
+    Walk(ilist.rewrite.Unroll()).rewrite(grow.code)
+    assert list(grow(n)) == expected
+
+
+def test_growing_list_invariant():
+    @structural_no_opt
+    def grow(n: int):
+        xs = [1]
+        for i in range(n):
+            xs = xs + [i]
+        return xs
+
+    frame, result = TypeInference(grow.dialects).run(grow)
+    expected = ilist.IListType[types.Int, types.Any]
+    assert result == expected
+    loop = next(stmt for stmt in grow.code.walk() if isinstance(stmt, scf.For))
+    assert all(frame.get(arg) == expected for arg in loop.body.blocks[0].args[1:])
+
+
+def test_stable_length_and_multiple_carried_values():
+    @structural_no_opt
+    def loop(n: int):
+        fixed = [1]
+        growing = [2]
+        for i in range(n):
+            fixed = [i]
+            growing = growing + fixed
+        return fixed, growing
+
+    _, result = TypeInference(loop.dialects).run(loop)
+    assert (
+        result
+        == types.Tuple[
+            ilist.IListType[types.Int, types.Literal(1)],
+            ilist.IListType[types.Int, types.Any],
+        ]
+    )
+
+
+def test_nested_loops_and_repeated_inference():
+    @structural_no_opt
+    def grow(n: int):
+        xs = [1]
+        for i in range(n):
+            for j in range(n):
+                xs = xs + [i + j]
+        return xs
+
+    expected = list(grow(3))
+    inference = TypeInfer(grow.dialects, no_raise=False)
+    inference(grow)
+    assert grow.return_type == ilist.IListType[types.Int, types.Any]
+    first = [value.type for stmt in grow.code.walk() for value in stmt.results]
+    inference(grow)
+    assert [value.type for stmt in grow.code.walk() for value in stmt.results] == first
+    assert list(grow(3)) == expected
+
+
+def test_changing_element_type():
+    @structural_no_opt
+    def grow(n: int):
+        xs = [1]
+        for _ in range(n):
+            xs = xs + [1.5]
+        return xs
+
+    _, result = TypeInference(grow.dialects).run(grow)
+    assert result == ilist.IListType[types.Int | types.Float, types.Any]
+
+
+def test_nested_tuple_growth_terminates():
+    @structural_no_opt
+    def nest(n: int):
+        value = ()
+        for _ in range(n):
+            value = (value,)
+        return value
+
+    _, result = TypeInference(nest.dialects).run(nest)
+    assert result == types.Any
+
+
+def test_typed_empty_list():
+    from kirin.interp import Frame
+    from kirin.dialects.ilist.typeinfer import TypeInfer as IListTypeInfer
+
+    stmt = ilist.New((), elem_type=types.Int)
+    result = IListTypeInfer().new(TypeInference(structural_no_opt), Frame(stmt), stmt)
+    assert result == (ilist.IListType[types.Int, types.Literal(0)],)
+
+
+def test_inconsistent_lattice_fails_with_bounded_iterations():
+    from contextlib import nullcontext
+    from unittest.mock import Mock
+
+    from pytest import raises
+
+    from kirin import interp
+    from kirin.analysis import ForwardFrame
+    from kirin.dialects.scf.typeinfer import TypeInfer as SCFTypeInfer
+
+    @structural_no_opt
+    def grow(n: int):
+        xs = [1]
+        for i in range(n):
+            xs = xs + [i]
+        return xs
+
+    loop = next(stmt for stmt in grow.code.walk() if isinstance(stmt, scf.For))
+    # Model a broken lattice that rejects containment even after widening to top.
+    value = Mock(spec=types.TypeAttribute)
+    value.is_subseteq.return_value = False
+    value.join.return_value = value
+    values = (value,) * len(loop.initializers)
+    frame = Mock(spec=ForwardFrame)
+    frame.get_values.return_value = values
+    interpreter = Mock(spec=TypeInference)
+    interpreter.frame_eval.return_value = (types.Int,)
+    interpreter.frame_call_region.return_value = values
+    interpreter.new_frame.side_effect = lambda *args, **kwargs: nullcontext(
+        ForwardFrame(loop)
+    )
+
+    with raises(
+        interp.InterpreterError, match="scf.For type inference did not converge"
+    ):
+        SCFTypeInfer().for_loop(interpreter, frame, loop)
+
+    assert interpreter.frame_call_region.call_count <= 32
+    frame.set_values.assert_not_called()
+
+
+@mark.parametrize("n", [0, 1, 2, 5])
+def test_branch_dependent_growth_preserves_execution(n):
+    @structural_no_opt
+    def grow(n: int):
+        xs = [1]
+        for i in range(n):
+            if i % 2 == 0:
+                xs = xs + [i]
+            else:
+                xs = xs + [i, i]
+            xs = ilist.map(identity, xs)
+        return xs
+
+    expected = list(grow(n))
+    TypeInfer(grow.dialects, no_raise=False)(grow)
+    maps = [stmt for stmt in grow.code.walk() if isinstance(stmt, ilist.Map)]
+    assert len(maps) == 1
+    assert ilist.IListType[types.Int, types.Any].is_subseteq(maps[0].collection.type)
+    Walk(ilist.rewrite.Unroll()).rewrite(grow.code)
+    assert list(grow(n)) == expected
+
+
+@mark.parametrize("n", [0, 1, 2, 5])
+def test_dependent_carried_lists_preserve_execution(n):
+    @structural_no_opt
+    def grow(n: int):
+        xs = [1]
+        ys = [2]
+        for i in range(n):
+            xs = ilist.map(identity, ys)
+            ys = xs + [i]
+        return xs, ys
+
+    before_xs, before_ys = grow(n)
+    TypeInfer(grow.dialects, no_raise=False)(grow)
+    mapped = next(stmt for stmt in grow.code.walk() if isinstance(stmt, ilist.Map))
+    assert mapped.collection.type == ilist.IListType[types.Int, types.Any]
+    Walk(ilist.rewrite.Unroll()).rewrite(grow.code)
+    after_xs, after_ys = grow(n)
+    assert list(after_xs) == list(before_xs)
+    assert list(after_ys) == list(before_ys)
+
+
+@mark.parametrize("budget", [0, 1, 16])
+def test_budget_fallback_preserves_stable_components(monkeypatch, budget):
+    monkeypatch.setattr("kirin.dialects.scf.typeinfer._LOOP_WIDENING_BUDGET", budget)
+
+    @structural_no_opt
+    def nest(n: int):
+        value = ()
+        fixed = [1]
+        for i in range(n):
+            value = (value,)
+            fixed = [i]
+        return value, fixed
+
+    before_value, before_fixed = nest(3)
+    TypeInfer(nest.dialects, no_raise=False)(nest)
+    assert (
+        nest.return_type
+        == types.Tuple[types.Any, ilist.IListType[types.Int, types.Literal(1)]]
+    )
+    Walk(ilist.rewrite.Unroll()).rewrite(nest.code)
+    after_value, after_fixed = nest(3)
+    assert after_value == before_value
+    assert list(after_fixed) == list(before_fixed)
+
+
+def test_fallback_propagates_through_dependent_values(monkeypatch):
+    monkeypatch.setattr("kirin.dialects.scf.typeinfer._LOOP_WIDENING_BUDGET", 0)
+
+    @structural_no_opt
+    def nest(n: int):
+        first = ()
+        second = ()
+        third = ()
+        for _ in range(n):
+            first = second
+            second = third
+            third = (third,)
+        return first, second, third
+
+    expected = [nest(n) for n in (0, 1, 4)]
+    TypeInfer(nest.dialects, no_raise=False)(nest)
+    assert nest.return_type == types.Tuple[types.Any, types.Any, types.Any]
+    loop = next(stmt for stmt in nest.code.walk() if isinstance(stmt, scf.For))
+    assert all(arg.type == types.Any for arg in loop.body.blocks[0].args[1:])
+    assert [nest(n) for n in (0, 1, 4)] == expected