feat[next]: Add support for tuple comprehensions (not stacked)

src/gt4py/next/ffront/field_operator_ast.py

@@ -98,6 +98,35 @@ class TupleExpr(Expr):
     elts: list[Expr]
 
 
+# TODO(tehrengruber): extend this to supported nested tuple comprehension.
+#  e.g. `tuple(element_expr for child in nested_tuple for grand_child in child)`
+#  would be represented by:
+#  ```
+#  class TupleComprehension(Expr):                              # ruff: noqa: ERA001
+#    inner: TupleComprehensionMapper | NestedTupleCompr         # ruff: noqa: ERA001
+#  class NestedTupleCompr(Expr, SymbolTableTrait):              # ruff: noqa: ERA001
+#    params: tuple[DataSymbol]                                  # ruff: noqa: ERA001
+#    body: TupleComprehension                                   # ruff: noqa: ERA001
+#  ```
+class TupleComprehension(Expr):
+    """
+    tuple(element_expr for target in iterable)
+    Note: The structure here differs from the one in the Python AST. Here we group target and
+    element expression in order to cleanly nest by the symbols being introduced, whereas in
+    the Python AST target and iterable are grouped into generator nodes.
+    """
+
+    inner: TupleComprehensionMapper
+    iterable: Expr
+
+
+# This is essentially a lambda. The difference is that for a lambda we might not know the type of
+# the args; therefore this is named differently at the moment.
+class TupleComprehensionMapper(LocatedNode, SymbolTableTrait):
+    target: Any  # TODO(tehrengruber): should be NestedTuple[DataSymbol], but this breaks in eve
+    element_expr: Expr
+
+
 class UnaryOp(Expr):
     op: dialect_ast_enums.UnaryOperator
     operand: Expr

src/gt4py/next/ffront/foast_passes/type_deduction.py

@@ -5,13 +5,13 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
-
 import textwrap
 from typing import Any, Optional, Sequence, TypeAlias, TypeVar, cast
 
 import gt4py.next.ffront.field_operator_ast as foast
 from gt4py import eve
 from gt4py.eve import NodeTranslator, NodeVisitor, traits
+from gt4py.eve.extended_typing import NestedTuple
 from gt4py.next import errors
 from gt4py.next.common import Dimension, DimensionKind, promote_dims
 from gt4py.next.ffront import (
@@ -24,6 +24,7 @@
 from gt4py.next.ffront.foast_passes import utils as foast_utils
 from gt4py.next.iterator import builtins
 from gt4py.next.type_system import type_info, type_specifications as ts, type_translation
+from gt4py.next.utils import tree_map
 
 
 OperatorNodeT = TypeVar("OperatorNodeT", bound=foast.LocatedNode)
@@ -453,6 +454,10 @@ def visit_Subscript(self, node: foast.Subscript, **kwargs: Any) -> foast.Subscri
                         f"Tuples need to be indexed with literal integers, got '{node.index}'.",
                     ) from ex
                 new_type = types[index]
+            case ts.VarArgType(element_type=element_type):
+                new_type = (
+                    element_type  # TODO: we only temporarily allow any index for vararg types
+                )
             case ts.OffsetType(source=source, target=(target1, target2)):
                 if not target2.kind == DimensionKind.LOCAL:
                     raise errors.DSLError(
@@ -715,6 +720,90 @@ def visit_TupleExpr(self, node: foast.TupleExpr, **kwargs: Any) -> foast.TupleEx
         new_type = ts.TupleType(types=[element.type for element in new_elts])
         return foast.TupleExpr(elts=new_elts, type=new_type, location=node.location)
 
+    def visit_TupleComprehension(
+        self, node: foast.TupleComprehension, **kwargs: Any
+    ) -> foast.TupleComprehension:
+        target = self.visit(node.inner.target, **kwargs)
+        iterable = self.visit(node.iterable, **kwargs)
+
+        def deduce_target_type(
+            target: NestedTuple[foast.Symbol] | foast.Symbol,
+            element_type: ts.TypeSpec,
+            inner_kwargs: dict[str, Any],
+        ) -> NestedTuple[foast.Symbol] | foast.Symbol:
+            @tree_map(with_path_arg=True)
+            def process_target(target_el: foast.Symbol, path: tuple[int, ...]) -> foast.Symbol:
+                try:
+                    type_ = element_type
+                    for i in path:
+                        if not isinstance(type_, ts.TupleType) or len(type_.types) <= i:
+                            raise IndexError()
+                        type_ = type_.types[i]
+                    return self.visit(target_el, refine_type=type_, **inner_kwargs)
+                except IndexError:
+                    raise errors.DSLError(
+                        target_el.location, f"Cannot unpack non-iterable '{type_}' object."
+                    ) from None
+
+            return process_target(target)
+
+        def deduce_mapper(
+            element_type: ts.DataType,
+        ) -> foast.TupleComprehensionMapper:
+            inner_kwargs = {**kwargs, "symtable": kwargs["symtable"].new_child()}
+            new_target = deduce_target_type(target, element_type, inner_kwargs)
+            return foast.TupleComprehensionMapper(
+                target=new_target,
+                element_expr=self.visit(node.inner.element_expr, **inner_kwargs),
+                location=node.location,
+            )
+
+        if isinstance(iterable.type, ts.TupleType):
+            if len(iterable.type.types) == 0:
+                raise errors.DSLError(
+                    iterable.location,
+                    "Cannot iterate over an empty tuple in a tuple comprehension.",
+                )
+            if not all(
+                isinstance(element_type, ts.DataType) for element_type in iterable.type.types
+            ):
+                raise errors.DSLError(
+                    iterable.location,
+                    "Tuple comprehension iterable elements must be data types.",
+                )
+
+            element_types = cast(list[ts.DataType], iterable.type.types)
+            if not all(element_type == element_types[0] for element_type in element_types):
+                raise NotImplementedError(
+                    "Tuple comprehensions over fixed-length tuples require all iterable "
+                    "elements to have the same type."
+                )
+            new_mapper = deduce_mapper(element_types[0])
+            result = foast.TupleComprehension(
+                inner=new_mapper,
+                iterable=iterable,
+                location=node.location,
+                type=ts.TupleType(types=[new_mapper.element_expr.type for _ in element_types]),
+            )
+            return result
+        elif isinstance(iterable.type, ts.VarArgType):
+            element_type = iterable.type.element_type
+            new_mapper = deduce_mapper(element_type)
+            element_expr = new_mapper.element_expr
+            return_type = ts.VarArgType(element_type=element_expr.type)
+
+            return foast.TupleComprehension(
+                inner=new_mapper,
+                iterable=iterable,
+                location=node.location,
+                type=return_type,
+            )
+        else:
+            raise errors.DSLError(
+                iterable.location,
+                f"Iterable in generator expression must be a tuple, got '{iterable.type}'.",
+            )
+
     def visit_Call(self, node: foast.Call, **kwargs: Any) -> foast.Call:
         new_func = self.visit(node.func, **kwargs)
         new_args = self.visit(node.args, **kwargs)
@@ -998,7 +1087,9 @@ def deduce_return_type(
                     f"Field arguments to '{func_name}' must be of same dtype, got '{t_dtype}' != "
                     f"'{f_dtype}'.",
                 )
-            return_dims = promote_dims(cond_dims, type_info.extract_dims(type_info.promote(tb, fb)))
+            return_dims = promote_dims(
+                cond_dims, type_info.extract_dims(tb), type_info.extract_dims(fb)
+            )
             return_type = ts.FieldType(dims=return_dims, dtype=t_dtype)
             return return_type
 

src/gt4py/next/ffront/foast_pretty_printer.py

@@ -120,6 +120,18 @@ def apply(cls, node: foast.LocatedNode, **kwargs: Any) -> str:  # type: ignore[o
 
     UnaryOp = as_fmt("{op}{operand}")
 
+    def visit_TupleComprehensionMapper(
+        self, node: foast.TupleComprehensionMapper, **kwargs: Any
+    ) -> str:
+        element_expr = self.visit(node.element_expr, **kwargs)
+        target = self.visit(node.target, **kwargs)
+        return f"{element_expr} for {target}"
+
+    def visit_TupleComprehension(self, node: foast.TupleComprehension, **kwargs: Any) -> str:
+        mapper = self.visit(node.inner, **kwargs)
+        iterable = self.visit(node.iterable, **kwargs)
+        return f"tuple(({mapper} in {iterable}))"
+
     def visit_UnaryOp(self, node: foast.UnaryOp, **kwargs: Any) -> str:
         if node.op is dialect_ast_enums.UnaryOperator.NOT:
             op = "not "

src/gt4py/next/ffront/foast_to_gtir.py

@@ -8,11 +8,12 @@
 
 
 import dataclasses
+import functools
 import warnings
 from typing import Any, Callable, Optional
 
 from gt4py import eve
-from gt4py.eve.extended_typing import Never, cast
+from gt4py.eve.extended_typing import Never
 from gt4py.next import utils
 from gt4py.next.ffront import (
     dialect_ast_enums,
@@ -260,6 +261,73 @@ def visit_Subscript(self, node: foast.Subscript, **kwargs: Any) -> itir.Expr:
     def visit_TupleExpr(self, node: foast.TupleExpr, **kwargs: Any) -> itir.Expr:
         return im.make_tuple(*[self.visit(el, **kwargs) for el in node.elts])
 
+    def _bind_tuple_comprehension_target(
+        self,
+        comprehension_target: itir.Sym | tuple,
+        element_expr: itir.Expr,
+        iterable_element: itir.Expr | str,
+    ) -> itir.Expr:
+        """Return ``element_expr`` with the comprehension target bound to one element."""
+        # For `2.0 * local_el + scalar_el for local_el, scalar_el in iterable`:
+        # - `comprehension_target`: `(local_el, scalar_el)`
+        # - `element_expr`: `2.0 * local_el + scalar_el`
+        # - `iterable_element` is the current element from `iterable`
+        # Returns `let local_el = iterable_element[0], scalar_el = iterable_element[1]
+        # in element_expr`.
+        if not isinstance(comprehension_target, tuple):
+            return im.let(comprehension_target, iterable_element)(element_expr)
+
+        flat_targets = utils.flatten_nested_tuple(comprehension_target)
+        nested_target_values = utils.tree_map(
+            lambda _, path: functools.reduce(
+                lambda element, index: im.tuple_get(index, element), path, iterable_element
+            ),
+            with_path_arg=True,
+        )(comprehension_target)
+
+        flat_target_values = utils.flatten_nested_tuple(nested_target_values)  # type: ignore[arg-type]
+
+        target_bindings = tuple(zip(flat_targets, flat_target_values, strict=True))
+        return im.let(*target_bindings)(element_expr)  # type: ignore[arg-type]
+
+    def visit_TupleComprehension(self, node: foast.TupleComprehension, **kwargs: Any) -> itir.Expr:
+        # e.g. tuple(2.0 * el for el in (a, a))` or `tuple(2.0 * el for el in (a(V2E), a(V2E)))`
+        # `tuple(2.0 * local_el + scalar_el for local_el, scalar_el in ((a(V2E), b), (c(V2E), d)))`.
+        # Only homogeneous (fixed-length and variable-length) tuples are supported.
+        comprehension_target = self.visit(node.inner.target, **kwargs)
+        element_expr = self.visit(node.inner.element_expr, **kwargs)
+        iterable_expr = self.visit(node.iterable, **kwargs)
+        iterable_type = node.iterable.type
+
+        def lower_body_for_iterable_element(iterable_element: itir.Expr | str) -> itir.Expr:
+            return self._bind_tuple_comprehension_target(
+                comprehension_target, element_expr, iterable_element
+            )
+
+        if isinstance(iterable_type, ts.TupleType):
+            assert isinstance(node.type, ts.TupleType)
+            iterable_value_name = next(self.uid_generator["__tuple_comprh"])
+
+            fixed_tuple_elements = [
+                lower_body_for_iterable_element(im.tuple_get(element_index, iterable_value_name))
+                for element_index in range(len(iterable_type.types))
+            ]
+
+            result_tuple = im.make_tuple(*fixed_tuple_elements)
+            return im.let(iterable_value_name, iterable_expr)(result_tuple)
+
+        assert isinstance(iterable_type, ts.VarArgType)
+        assert isinstance(node.type, ts.VarArgType)
+        if not isinstance(comprehension_target, tuple):
+            map_tuple_lambda = im.lambda_(comprehension_target)(element_expr)
+        else:
+            iterable_element_param = next(self.uid_generator["__tuple_comprh"])
+            map_tuple_lambda = im.lambda_(iterable_element_param)(
+                lower_body_for_iterable_element(iterable_element_param)
+            )
+
+        return im.call(im.call("map_tuple")(map_tuple_lambda))(iterable_expr)
+
     def visit_UnaryOp(self, node: foast.UnaryOp, **kwargs: Any) -> itir.Expr:
         # TODO(tehrengruber): extend iterator ir to support unary operators
         dtype = type_info.extract_dtype(node.type)
@@ -521,10 +589,7 @@ def _visit_type_constr(self, node: foast.Call, **kwargs: Any) -> itir.Expr:
         return im.literal(str(val), target_type)
 
     def _make_literal(self, val: Any, type_: ts.TypeSpec) -> itir.Expr:
-        if isinstance(type_, ts.COLLECTION_TYPE_SPECS):
-            type_ = cast(
-                ts.CollectionTypeSpec, type_
-            )  # This shouldn't be needed after the previous isinstance() check
+        if isinstance(type_, (ts.TupleType, ts.NamedCollectionType)):
             # This code-path is only active in the init of a scan,
             # as otherwise the frontend generates tuple expressions of `Constant`s.
             val = arguments.extract(val) if isinstance(type_, ts.NamedCollectionType) else val

src/gt4py/next/ffront/foast_to_past.py

@@ -113,9 +113,10 @@ def __call__(self, inp: ConcreteFOASTOperatorDef) -> ConcretePASTProgramDef:
             *partial_program_type.definition.kw_only_args.keys(),
         ]
         assert isinstance(type_, ts.CallableType)
-        assert arg_types[-1] == type_info.return_type(
+        return_type = type_info.return_type(
             type_, with_args=list(arg_types), with_kwargs=kwarg_types
         )
+        assert type_info.is_concretizable(return_type, arg_types[-1])
         assert args_names[-1] == "out"
 
         params_decl: list[past.Symbol] = [