Propagate foast & past typing information to ITIR

src/gt4py/next/ffront/foast_to_itir.py

@@ -249,6 +249,11 @@ def _visit_assign(self, node: foast.Assign, **kwargs) -> tuple[itir.Sym, itir.Ex
         return sym, expr
 
     def visit_Symbol(self, node: foast.Symbol, **kwargs) -> itir.Sym:
+        # TODO(tehrengruber): extend to more types
+        if isinstance(node.type, ts.FieldType):
+            kind = "Iterator"
+            dtype = node.type.dtype.kind.name.lower()
+            return itir.Sym(id=node.id, kind=kind, dtype=dtype)
         return im.sym(node.id)
 
     def visit_Name(self, node: foast.Name, **kwargs) -> itir.SymRef:

src/gt4py/next/ffront/past_to_itir.py

@@ -348,6 +348,11 @@ def visit_Name(self, node: past.Name, **kwargs) -> itir.SymRef:
         return itir.SymRef(id=node.id)
 
     def visit_Symbol(self, node: past.Symbol, **kwargs) -> itir.Sym:
+        # TODO(tehrengruber): extend to more types
+        if isinstance(node.type, ts.FieldType):
+            kind = "Iterator"
+            dtype = node.type.dtype.kind.name.lower()
+            return itir.Sym(id=node.id, kind=kind, dtype=dtype)
         return itir.Sym(id=node.id)
 
     def visit_BinOp(self, node: past.BinOp, **kwargs) -> itir.FunCall:

src/gt4py/next/iterator/ir.py

@@ -12,7 +12,8 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-from typing import ClassVar, List, Union
+import typing
+from typing import ClassVar, List, Optional, Union
 
 import gt4py.eve as eve
 from gt4py.eve import Coerced, SymbolName, SymbolRef, datamodels
@@ -38,6 +39,22 @@ def __hash__(self) -> int:
 
 class Sym(Node):  # helper
     id: Coerced[SymbolName]  # noqa: A003
+    # TODO(tehrengruber): Revisit. Using strings is a workaround to avoid coupling with the
+    #   type inference.
+    kind: Optional[typing.Literal["Iterator", "Value"]] = None
+    dtype: Optional[str] = None
+
+    @datamodels.validator("dtype")
+    def _kind_validator(self: datamodels.DataModelTP, attribute: datamodels.Attribute, value: str):
+        if value and value not in ["Iterator", "Value"]:
+            raise ValueError(f"Invalid kind `{value}`, must be one of `Iterator`, `Value`.")
+
+    @datamodels.validator("dtype")
+    def _dtype_validator(self: datamodels.DataModelTP, attribute: datamodels.Attribute, value: str):
+        if value and value not in TYPEBUILTINS:
+            raise ValueError(
+                f"Invalid dtype `{value}`, must be one of `{'`, `'.join(TYPEBUILTINS)}`."
+            )
 
 
 @noninstantiable

src/gt4py/next/iterator/type_inference.py

@@ -417,7 +417,25 @@ def visit(self, node, **kwargs) -> typing.Any:
         return result
 
     def visit_Sym(self, node: ir.Sym, **kwargs) -> Type:
-        return TypeVar.fresh()
+        result = TypeVar.fresh()
+        if node.kind:
+            kind = {"Iterator": Iterator(), "Value": Value()}[node.kind]
+            self.constraints.add(
+                (Val(kind=kind, current_loc=TypeVar.fresh(), defined_loc=TypeVar.fresh()), result)
+            )
+        if node.dtype:
+            assert node.dtype in ir.TYPEBUILTINS
+            self.constraints.add(
+                (
+                    Val(
+                        dtype=Primitive(name=node.dtype),
+                        current_loc=TypeVar.fresh(),
+                        defined_loc=TypeVar.fresh(),
+                    ),
+                    result,
+                )
+            )
+        return result
 
     def visit_SymRef(self, node: ir.SymRef, *, symtable, **kwargs) -> Type:
         if node.id in ir.BUILTINS:

tests/next_tests/iterator_tests/test_type_inference.py

@@ -93,6 +93,26 @@ def test_lambda():
     assert ti.pformat(inferred) == "(T₀) → T₀"
 
 
+def test_typed_lambda():
+    testee = ir.Lambda(
+        params=[ir.Sym(id="x", kind="Iterator", dtype="float")], expr=ir.SymRef(id="x")
+    )
+    expected_val = ti.Val(
+        kind=ti.Iterator(),
+        dtype=ti.Primitive(name="float"),
+        size=ti.TypeVar(idx=0),
+        current_loc=ti.TypeVar(idx=1),
+        defined_loc=ti.TypeVar(idx=2),
+    )
+    expected = ti.FunctionType(
+        args=ti.Tuple.from_elems(expected_val),
+        ret=expected_val,
+    )
+    inferred = ti.infer(testee)
+    assert inferred == expected
+    assert ti.pformat(inferred) == "(It[T₁, T₂, float⁰]) → It[T₁, T₂, float⁰]"
+
+
 def test_plus():
     testee = ir.SymRef(id="plus")
     t = ti.Val(kind=ti.Value(), dtype=ti.TypeVar(idx=0), size=ti.TypeVar(idx=1))