Repository: GridTools/gt4py
Pull request #1118 (2023-01-06 01:29:05)

Add BlockStmt
During development of if-statements #1079 a new FOAST node `BlockStmt` was introduced. The node is used to group multiple `foast.Stmt`s into a single node (e.g. each branch of a statement is a single `BlockStmt`). Primarily to avoid #1079 to rot before the conceptual problems on ITIR are solved these changes were factored out into this PR. However the changes are meaningful on their own as they introduce a cleaner lowering mechanism from the somewhat-stateful FOAST to the functional ITIR.

Additional changes in this PR:
- Fixed small bug in `test_closure_symbols`
- Add `if_` helper function to `itir_makers.py`


Patch:
diff --git a/src/functional/ffront/field_operator_ast.py b/src/functional/ffront/field_operator_ast.py
index b1873094..59278b68 100644
--- a/src/functional/ffront/field_operator_ast.py
+++ b/src/functional/ffront/field_operator_ast.py
@@ -11,6 +11,9 @@
 # distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
+
+from __future__ import annotations
+
 from typing import Any, Generic, TypeVar, Union
 
 from eve import Coerced, Node, SourceLocation, SymbolName, SymbolRef
@@ -171,10 +174,14 @@ class Return(Stmt):
     value: Expr
 
 
+class BlockStmt(Stmt, SymbolTableTrait):
+    stmts: list[Stmt]
+
+
 class FunctionDefinition(LocatedNode, SymbolTableTrait):
     id: Coerced[SymbolName]  # noqa: A003  # shadowing a python builtin
     params: list[DataSymbol]
-    body: list[Stmt]
+    body: BlockStmt
     closure_vars: list[Symbol]
     type: Union[ts.FunctionType, ts.DeferredType] = ts.DeferredType(  # noqa: A003
         constraint=ts.FunctionType
diff --git a/src/functional/ffront/foast_introspection.py b/src/functional/ffront/foast_introspection.py
new file mode 100644
index 00000000..00b34c60
--- /dev/null
+++ b/src/functional/ffront/foast_introspection.py
@@ -0,0 +1,53 @@
+# GT4Py Project - GridTools Framework
+#
+# Copyright (c) 2014-2023, ETH Zurich
+# All rights reserved.
+#
+# This file is part of the GT4Py project and the GridTools framework.
+# GT4Py is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+import enum
+
+from functional.ffront import field_operator_ast as foast
+
+
+class StmtReturnKind(enum.IntEnum):
+    UNCONDITIONAL_RETURN = 0
+    NO_RETURN = 2
+
+
+def deduce_stmt_return_kind(node: foast.Stmt) -> StmtReturnKind:
+    """
+    Deduce if a statement returns and if so, whether it does unconditionally.
+
+    Example with ``StmtReturnKind.UNCONDITIONAL_RETURN``::
+
+        if cond:
+          return 1
+        else:
+          return 2
+
+    Example with ``StmtReturnKind.NO_RETURN``::
+
+        if cond:
+          result = 1
+        else:
+          result = 2
+    """
+    if isinstance(node, foast.Return):
+        return StmtReturnKind.UNCONDITIONAL_RETURN
+    elif isinstance(node, foast.BlockStmt):
+        for stmt in node.stmts:
+            return_kind = deduce_stmt_return_kind(stmt)
+            if return_kind != StmtReturnKind.NO_RETURN:
+                return return_kind
+        return StmtReturnKind.NO_RETURN
+    elif isinstance(node, (foast.Assign, foast.TupleTargetAssign)):
+        return StmtReturnKind.NO_RETURN
+    else:
+        raise AssertionError(f"Statements of type `{type(node).__name__}` not understood.")
diff --git a/src/functional/ffront/foast_passes/iterable_unpack.py b/src/functional/ffront/foast_passes/iterable_unpack.py
index a45f8ccd..135d5fdd 100644
--- a/src/functional/ffront/foast_passes/iterable_unpack.py
+++ b/src/functional/ffront/foast_passes/iterable_unpack.py
@@ -34,20 +34,6 @@ class UnpackedAssignPass(NodeTranslator, traits.VisitorWithSymbolTableTrait):
         node = cls().visit(node)
         return node
 
-    def visit_FunctionDefinition(self, node: foast.FunctionDefinition, **kwargs):
-        new_body = self.visit(node.body, **kwargs)
-        unrolled_body = self._unroll_tuple_target_assign(new_body)
-        assert isinstance(unrolled_body[-1], foast.Return)
-
-        return foast.FunctionDefinition(
-            id=node.id,
-            params=self.visit(node.params, **kwargs),
-            body=unrolled_body,
-            closure_vars=self.visit(node.closure_vars, **kwargs),
-            type=node.type,
-            location=node.location,
-        )
-
     def _unique_tuple_symbol(self, node: foast.TupleTargetAssign) -> foast.Symbol[Any]:
         sym: foast.Symbol = foast.Symbol(
             id=f"__tuple_tmp_{self.unique_tuple_symbol_id}",
@@ -57,18 +43,16 @@ class UnpackedAssignPass(NodeTranslator, traits.VisitorWithSymbolTableTrait):
         self.unique_tuple_symbol_id += 1
         return sym
 
-    def _unroll_tuple_target_assign(
-        self, body: list[foast.LocatedNode]
-    ) -> list[foast.Assign | foast.LocatedNode]:
-        unrolled: list[foast.Assign | foast.LocatedNode] = []
+    def visit_BlockStmt(self, node: foast.BlockStmt, **kwargs) -> foast.BlockStmt:
+        unrolled_stmts: list[foast.Assign | foast.BlockStmt | foast.Return] = []
 
-        for node in body:
-            if isinstance(node, foast.TupleTargetAssign):
-                num_elts, targets = len(node.value.type.types), node.targets  # type: ignore
+        for stmt in node.stmts:
+            if isinstance(stmt, foast.TupleTargetAssign):
+                num_elts, targets = len(stmt.value.type.types), stmt.targets  # type: ignore
                 indices = compute_assign_indices(targets, num_elts)
-                tuple_symbol = self._unique_tuple_symbol(node)
-                unrolled.append(
-                    foast.Assign(target=tuple_symbol, value=node.value, location=node.location)
+                tuple_symbol = self._unique_tuple_symbol(stmt)
+                unrolled_stmts.append(
+                    foast.Assign(target=tuple_symbol, value=stmt.value, location=stmt.location)
                 )
 
                 for (index, subtarget) in zip(indices, targets):
@@ -81,7 +65,7 @@ class UnpackedAssignPass(NodeTranslator, traits.VisitorWithSymbolTableTrait):
                         slice_indices = list(range(lower, upper))
                         tuple_slice = [
                             foast.Subscript(
-                                value=tuple_name, index=i, type=el_type, location=node.location
+                                value=tuple_name, index=i, type=el_type, location=stmt.location
                             )
                             for i in slice_indices
                         ]
@@ -89,23 +73,23 @@ class UnpackedAssignPass(NodeTranslator, traits.VisitorWithSymbolTableTrait):
                         new_tuple = foast.TupleExpr(
                             elts=tuple_slice,
                             type=el_type,
-                            location=node.location,
+                            location=stmt.location,
                         )
                         new_assign = foast.Assign(
                             target=subtarget.id,
                             value=new_tuple,
-                            location=node.location,
+                            location=stmt.location,
                         )
                     else:
                         new_assign = foast.Assign(
                             target=subtarget,
                             value=foast.Subscript(
-                                value=tuple_name, index=index, type=el_type, location=node.location
+                                value=tuple_name, index=index, type=el_type, location=stmt.location
                             ),
-                            location=node.location,
+                            location=stmt.location,
                         )
-                    unrolled.append(new_assign)
+                    unrolled_stmts.append(new_assign)
             else:
-                unrolled.append(node)
+                unrolled_stmts.append(self.generic_visit(stmt, **kwargs))
 
-        return unrolled
+        return foast.BlockStmt(stmts=unrolled_stmts, location=node.location)
diff --git a/src/functional/ffront/foast_passes/type_deduction.py b/src/functional/ffront/foast_passes/type_deduction.py
index 90f9fd69..4642a6f7 100644
--- a/src/functional/ffront/foast_passes/type_deduction.py
+++ b/src/functional/ffront/foast_passes/type_deduction.py
@@ -110,6 +110,20 @@ def promote_to_mask_type(
         return input_type
 
 
+def deduce_stmt_return_type(node: foast.BlockStmt) -> ts.TypeSpec:
+    """Deduce type of value returned inside a block statement."""
+    for stmt in node.stmts:
+        if isinstance(stmt, foast.Return):
+            return stmt.value.type
+
+    # If the node was constructed by the foast parsing we should never get here, but instead
+    # have gotten an error there.
+    raise AssertionError(
+        "Malformed block statement. Expected a return statement in this context, "
+        "but none was found. Please submit a bug report."
+    )
+
+
 class FieldOperatorTypeDeductionCompletnessValidator(NodeVisitor):
     """Validate an FOAST expression is fully typed."""
 
@@ -151,11 +165,11 @@ class FieldOperatorTypeDeduction(traits.VisitorWithSymbolTableTrait, NodeTransla
     >>> untyped_fieldop = FieldOperatorParser(
     ...     source_definition=source_definition, closure_vars=closure_vars, annotations=annotations
     ... ).visit(ast.parse(source_definition.source).body[0])
-    >>> untyped_fieldop.body[0].value.type
+    >>> untyped_fieldop.body.stmts[0].value.type
     DeferredType(constraint=None)
 
     >>> typed_fieldop = FieldOperatorTypeDeduction.apply(untyped_fieldop)
-    >>> assert typed_fieldop.body[0].value.type == ts.FieldType(dtype=ts.ScalarType(
+    >>> assert typed_fieldop.body.stmts[0].value.type == ts.FieldType(dtype=ts.ScalarType(
     ...     kind=ts.ScalarKind.FLOAT64), dims=[IDim])
     """
 
@@ -171,8 +185,7 @@ class FieldOperatorTypeDeduction(traits.VisitorWithSymbolTableTrait, NodeTransla
         new_params = self.visit(node.params, **kwargs)
         new_body = self.visit(node.body, **kwargs)
         new_closure_vars = self.visit(node.closure_vars, **kwargs)
-        assert isinstance(new_body[-1], foast.Return)
-        return_type = new_body[-1].value.type
+        return_type = deduce_stmt_return_type(new_body)
         if not isinstance(return_type, (ts.DataType, ts.DeferredType, ts.VoidType)):
             raise FieldOperatorTypeDeductionError.from_foast_node(
                 node,
@@ -246,7 +259,7 @@ class FieldOperatorTypeDeduction(traits.VisitorWithSymbolTableTrait, NodeTransla
         symtable = kwargs["symtable"]
         if node.id not in symtable or symtable[node.id].type is None:
             raise FieldOperatorTypeDeductionError.from_foast_node(
-                node, msg=f"Undeclared symbol {node.id}"
+                node, msg=f"Undeclared symbol `{node.id}`."
             )
 
         symbol = symtable[node.id]
diff --git a/src/functional/ffront/foast_pretty_printer.py b/src/functional/ffront/foast_pretty_printer.py
index 027d289e..8a1b2088 100644
--- a/src/functional/ffront/foast_pretty_printer.py
+++ b/src/functional/ffront/foast_pretty_printer.py
@@ -11,6 +11,9 @@ from functional.ffront import dialect_ast_enums, type_specifications as ts
 
 PropertyIdentifier: TypeAlias = Union[type[foast.LocatedNode], tuple[type[foast.LocatedNode], str]]
 
+INDENTATION_PREFIX: Final[str] = "  "
+
+
 # See https://docs.python.org/3/reference/expressions.html#operator-precedence
 # The following list contains all entries from the above link. Operators that
 # are not modeled in FOAST are kept for ease of comparability and future
@@ -139,11 +142,13 @@ class _PrettyPrinter(TemplatedGenerator):
 
     Return = as_fmt("return {value}")
 
+    BlockStmt = as_mako("${'\\n'.join(stmts)}")
+
     FunctionDefinition = as_mako(
         textwrap.dedent(
             """
             def ${id}(${', '.join(params_annotated)})${return_type}:
-              ${'\\n  '.join(body)}
+            ${indented_body}
             """
         ).strip()
     )
@@ -160,7 +165,13 @@ class _PrettyPrinter(TemplatedGenerator):
         return_type = (
             f" -> {node.type.returns}" if not isinstance(node.type, ts.DeferredType) else ""
         )
-        return self.generic_visit(node, params_annotated=params_annotated, return_type=return_type)
+        indented_body = textwrap.indent(self.visit(node.body), INDENTATION_PREFIX)
+        return self.generic_visit(
+            node,
+            indented_body=indented_body,
+            params_annotated=params_annotated,
+            return_type=return_type,
+        )
 
     FieldOperator = as_fmt("@field_operator\n{definition}")
 
diff --git a/src/functional/ffront/foast_to_itir.py b/src/functional/ffront/foast_to_itir.py
index 73491be0..85f6d9c7 100644
--- a/src/functional/ffront/foast_to_itir.py
+++ b/src/functional/ffront/foast_to_itir.py
@@ -15,7 +15,7 @@
 import enum
 import itertools
 from dataclasses import dataclass, field
-from typing import Any, Callable, Optional, cast
+from typing import Any, Callable
 
 import numpy as np
 
@@ -115,7 +115,7 @@ def to_value(node: foast.Expr) -> Callable[[itir.Expr], itir.Expr]:
     ...    return a, b
 
     >>> parsed = FieldOperatorParser.apply_to_function(foo)
-    >>> field_a, scalar_b = parsed.body[-1].value.elts
+    >>> field_a, scalar_b = parsed.body.stmts[-1].value.elts
     >>> to_value(field_a)(im.ref("a"))
     FunCall(fun=SymRef(id=SymbolRef('deref')), args=[SymRef(id=SymbolRef('a'))])
     >>> to_value(scalar_b)(im.ref("a"))
@@ -178,6 +178,10 @@ class FieldOperatorLowering(NodeTranslator):
             if isinstance(node.definition.params[i].type, ts.ScalarType):
                 new_body = im.let(param.id, im.deref_(param.id))(new_body)
 
+        assert isinstance(node.type, ts.FieldOperatorType)
+        if iterator_type_kind(node.type.definition.returns) == ITIRTypeKind.ITERATOR:
+            new_body = im.deref_(new_body)
+
         return itir.FunctionDefinition(
             id=func_definition.id,
             params=func_definition.params,
@@ -192,7 +196,7 @@ class FieldOperatorLowering(NodeTranslator):
         return itir.FunctionDefinition(
             id=node.id,
             params=params,
-            expr=self._visit_body(node.body, params=params, symtable=symtable),
+            expr=self.visit_BlockStmt(node.body, inner_expr=None, symtable=symtable),
         )
 
     def visit_ScanOperator(self, node: foast.ScanOperator, **kwargs) -> itir.FunctionDefinition:
@@ -218,27 +222,34 @@ class FieldOperatorLowering(NodeTranslator):
             expr=body,
         )
 
-    def _visit_body(
-        self, body: list[foast.Stmt], params: Optional[list[itir.Sym]] = None, **kwargs
-    ) -> itir.FunCall:
-        *assigns, return_stmt = body
-        current_expr = self.visit(return_stmt, **kwargs)
-
-        for assign in reversed(assigns):
-            current_expr = im.let(*self._visit_assign(cast(foast.Assign, assign), **kwargs))(
-                current_expr
-            )
-
-        return current_expr
+    def visit_Stmt(self, node: foast.Stmt, **kwargs):
+        raise AssertionError("Statements must always be visited in the context of a function.")
+
+    def visit_Return(
+        self, node: foast.Return, *, inner_expr: itir.Expr | None, **kwargs
+    ) -> itir.Expr:
+        return self.visit(node.value, **kwargs)
+
+    def visit_BlockStmt(
+        self, node: foast.BlockStmt, *, inner_expr: itir.Expr | None, **kwargs
+    ) -> itir.Expr:
+        for stmt in reversed(node.stmts):
+            inner_expr = self.visit(stmt, inner_expr=inner_expr, **kwargs)
+        assert inner_expr
+        return inner_expr
+
+    def visit_Assign(
+        self, node: foast.Assign, *, inner_expr: itir.Expr | None, **kwargs
+    ) -> itir.Expr:
+        return im.let(self.visit(node.target, **kwargs), self.visit(node.value, **kwargs))(
+            inner_expr
+        )
 
     def _visit_assign(self, node: foast.Assign, **kwargs) -> tuple[itir.Sym, itir.Expr]:
         sym = self.visit(node.target, **kwargs)
         expr = self.visit(node.value, **kwargs)
         return sym, expr
 
-    def visit_Return(self, node: foast.Return, **kwargs) -> itir.Expr:
-        return to_value(node.value)(self.visit(node.value, **kwargs))
-
     def visit_Symbol(self, node: foast.Symbol, **kwargs) -> itir.Sym:
         return im.sym(node.id)
 
@@ -328,7 +339,7 @@ class FieldOperatorLowering(NodeTranslator):
         lowered_false_expr = self.visit(node.false_expr, **kwargs)
 
         return self._lift_if_field(node)(
-            im.call_("if_")(
+            im.if_(
                 lowered_node_cond,
                 to_value(node.true_expr)(lowered_true_expr),
                 to_value(node.false_expr)(lowered_false_expr),
diff --git a/src/functional/ffront/func_to_foast.py b/src/functional/ffront/func_to_foast.py
index 5f08811e..1e52428a 100644
--- a/src/functional/ffront/func_to_foast.py
+++ b/src/functional/ffront/func_to_foast.py
@@ -35,6 +35,7 @@ from functional.ffront.ast_passes import (
     UnchainComparesPass,
 )
 from functional.ffront.dialect_parser import DialectParser, DialectSyntaxError
+from functional.ffront.foast_introspection import StmtReturnKind, deduce_stmt_return_kind
 from functional.ffront.foast_passes.closure_var_folding import ClosureVarFolding
 from functional.ffront.foast_passes.closure_var_type_deduction import ClosureVarTypeDeduction
 from functional.ffront.foast_passes.dead_closure_var_elimination import DeadClosureVarElimination
@@ -67,7 +68,7 @@ class FieldOperatorParser(DialectParser[foast.FunctionDefinition]):
     FunctionDefinition(..., id=SymbolName('field_op'), ...)
     >>> foast_tree.params  # doctest: +ELLIPSIS
     [Symbol(..., id=SymbolName('inp'), type=FieldType(...), ...)]
-    >>> foast_tree.body  # doctest: +ELLIPSIS
+    >>> foast_tree.body.stmts  # doctest: +ELLIPSIS
     [Return(..., value=Name(..., id=SymbolRef('inp')))]
 
 
@@ -170,10 +171,17 @@ class FieldOperatorParser(DialectParser[foast.FunctionDefinition]):
                 )
             )
 
+        new_body = self._visit_stmts(node.body, self._make_loc(node), **kwargs)
+
+        if deduce_stmt_return_kind(new_body) != StmtReturnKind.UNCONDITIONAL_RETURN:
+            raise FieldOperatorSyntaxError.from_AST(
+                node, msg="Function must return a value, but no return statement was found."
+            )
+
         return foast.FunctionDefinition(
             id=node.name,
             params=self.visit(node.args, **kwargs),
-            body=self.visit_stmt_list(node.body, **kwargs),
+            body=new_body,
             closure_vars=closure_var_symbols,
             location=self._make_loc(node),
         )
@@ -314,13 +322,6 @@ class FieldOperatorParser(DialectParser[foast.FunctionDefinition]):
             raise FieldOperatorSyntaxError.from_AST(node, msg="Empty return not allowed")
         return foast.Return(value=self.visit(node.value), location=self._make_loc(node))
 
-    def visit_stmt_list(self, nodes: list[ast.stmt], **kwargs) -> list[foast.Expr]:
-        if not isinstance(last_node := nodes[-1], ast.Return):
-            raise FieldOperatorSyntaxError.from_AST(
-                last_node, msg="Field operator must return a field expression on the last line!"
-            )
-        return [self.visit(node, **kwargs) for node in nodes]
-
     def visit_Expr(self, node: ast.Expr) -> foast.Expr:
         return self.visit(node.value)
 
@@ -394,6 +395,14 @@ class FieldOperatorParser(DialectParser[foast.FunctionDefinition]):
             type=ts.DeferredType(constraint=ts.DataType),
         )
 
+    def _visit_stmts(
+        self, stmts: list[ast.stmt], location: eve.SourceLocation, **kwargs
+    ) -> foast.BlockStmt:
+        return foast.BlockStmt(
+            stmts=[self.visit(el, **kwargs) for el in stmts if not isinstance(el, ast.Pass)],
+            location=location,
+        )
+
     def visit_Compare(self, node: ast.Compare, **kwargs) -> foast.Compare:
         if len(node.ops) != 1 or len(node.comparators) != 1:
             raise FieldOperatorSyntaxError.from_AST(
diff --git a/src/functional/ffront/itir_makers.py b/src/functional/ffront/itir_makers.py
index 899432e5..90b0abd7 100644
--- a/src/functional/ffront/itir_makers.py
+++ b/src/functional/ffront/itir_makers.py
@@ -223,6 +223,11 @@ def tuple_get_(tuple_expr, index):
     return call_("tuple_get")(tuple_expr, index)
 
 
+def if_(cond, true_val, false_val):
+    """Create a not_ FunCall, shorthand for ``call("if_")(expr)``."""
+    return call_("if_")(cond, true_val, false_val)
+
+
 def lift_(expr):
     """Create a lift FunCall, shorthand for ``call(call("lift")(expr))``."""
     return call_(call_("lift")(expr))


Test patch:
diff --git a/tests/functional_tests/ffront_tests/test_execution.py b/tests/functional_tests/ffront_tests/test_execution.py
index d4d4ab8d..e409f5c2 100644
--- a/tests/functional_tests/ffront_tests/test_execution.py
+++ b/tests/functional_tests/ffront_tests/test_execution.py
@@ -744,8 +744,6 @@ def test_where_k_offset(fieldview_backend):
 
 
 def test_undefined_symbols():
-    from functional.ffront.foast_passes.type_deduction import FieldOperatorTypeDeductionError
-
     with pytest.raises(FieldOperatorTypeDeductionError, match="Undeclared symbol"):
 
         @field_operator
@@ -800,7 +798,7 @@ def test_tuple_unpacking(fieldview_backend):
     out4 = np_as_located_field(IDim)(np.ones((size)))
 
     @field_operator
-    def _unpack(
+    def unpack(
         inp: Field[[IDim], float64],
     ) -> tuple[
         Field[[IDim], float64],
@@ -811,17 +809,7 @@ def test_tuple_unpacking(fieldview_backend):
         a, b, c, d = (inp + 2.0, inp + 3.0, inp + 5.0, inp + 7.0)
         return a, b, c, d
 
-    @program(backend=fieldview_backend)
-    def unpack(
-        inp: Field[[IDim], float64],
-        out1: Field[[IDim], float64],
-        out2: Field[[IDim], float64],
-        out3: Field[[IDim], float64],
-        out4: Field[[IDim], float64],
-    ):
-        _unpack(inp, out=(out1, out2, out3, out4))
-
-    unpack(inp, out1, out2, out3, out4, offset_provider={})
+    unpack(inp, out=(out1, out2, out3, out4), offset_provider={})
 
     arr = inp.array()
 
diff --git a/tests/functional_tests/ffront_tests/test_foast_lowering.py b/tests/functional_tests/ffront_tests/test_foast_lowering.py
index 57dfb0ff..81e4ede8 100644
--- a/tests/functional_tests/ffront_tests/test_foast_lowering.py
+++ b/tests/functional_tests/ffront_tests/test_foast_lowering.py
@@ -60,7 +60,7 @@ def test_copy():
     lowered = FieldOperatorLowering.apply(parsed)
 
     assert lowered.id == "copy_field"
-    assert lowered.expr == im.deref_("inp")
+    assert lowered.expr == im.ref("inp")
 
 
 def test_scalar_arg():
@@ -71,9 +71,7 @@ def test_scalar_arg():
     parsed = FieldOperatorParser.apply_to_function(scalar_arg)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.deref_(
-        im.lift_(im.lambda__("bar")(im.multiplies_("alpha", im.deref_("bar"))))("bar")
-    )
+    reference = im.lift_(im.lambda__("bar")(im.multiplies_("alpha", im.deref_("bar"))))("bar")
 
     assert lowered.expr == reference
 
@@ -85,11 +83,9 @@ def test_multicopy():
     parsed = FieldOperatorParser.apply_to_function(multicopy)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.deref_(
-        im.lift_(im.lambda__("inp1", "inp2")(im.make_tuple_(im.deref_("inp1"), im.deref_("inp2"))))(
-            "inp1", "inp2"
-        )
-    )
+    reference = im.lift_(
+        im.lambda__("inp1", "inp2")(im.make_tuple_(im.deref_("inp1"), im.deref_("inp2")))
+    )("inp1", "inp2")
 
     assert lowered.expr == reference
 
@@ -102,11 +98,9 @@ def test_arithmetic():
     parsed = FieldOperatorParser.apply_to_function(arithmetic)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.deref_(
-        im.lift_(im.lambda__("inp1", "inp2")(im.plus_(im.deref_("inp1"), im.deref_("inp2"))))(
-            "inp1", "inp2"
-        )
-    )
+    reference = im.lift_(
+        im.lambda__("inp1", "inp2")(im.plus_(im.deref_("inp1"), im.deref_("inp2")))
+    )("inp1", "inp2")
 
     assert lowered.expr == reference
 
@@ -121,7 +115,7 @@ def test_shift():
     parsed = FieldOperatorParser.apply_to_function(shift_by_one)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.deref_(im.shift_("Ioff", 1)("inp"))
+    reference = im.shift_("Ioff", 1)("inp")
 
     assert lowered.expr == reference
 
@@ -136,7 +130,7 @@ def test_negative_shift():
     parsed = FieldOperatorParser.apply_to_function(shift_by_one)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.deref_(im.shift_("Ioff", -1)("inp"))
+    reference = im.shift_("Ioff", -1)("inp")
 
     assert lowered.expr == reference
 
@@ -152,7 +146,7 @@ def test_temp_assignment():
     lowered = FieldOperatorLowering.apply(parsed)
 
     reference = im.let("tmp__0", "inp")(
-        im.let("inp__0", "tmp__0")(im.let("tmp2__0", "inp__0")(im.deref_("tmp2__0")))
+        im.let("inp__0", "tmp__0")(im.let("tmp2__0", "inp__0")("tmp2__0"))
     )
 
     assert lowered.expr == reference
@@ -174,7 +168,7 @@ def test_unary_ops():
         im.let(
             "tmp__1",
             im.lift_(im.lambda__("tmp__0")(im.minus_(0, im.deref_("tmp__0"))))("tmp__0"),
-        )(im.deref_("tmp__1"))
+        )("tmp__1")
     )
 
     assert lowered.expr == reference
@@ -203,7 +197,7 @@ def test_unpacking():
     )("__tuple_tmp_0")
 
     reference = im.let("__tuple_tmp_0", tuple_expr)(
-        im.let("tmp1__0", tuple_access_0)(im.let("tmp2__0", tuple_access_1)(im.deref_("tmp1__0")))
+        im.let("tmp1__0", tuple_access_0)(im.let("tmp2__0", tuple_access_1)("tmp1__0"))
     )
     assert lowered.expr == reference
 
@@ -218,7 +212,7 @@ def test_annotated_assignment():
     parsed = FieldOperatorParser.apply_to_function(copy_field)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.let("tmp__0", "inp")(im.deref_("tmp__0"))
+    reference = im.let("tmp__0", "inp")("tmp__0")
 
     assert lowered.expr == reference
 
@@ -240,7 +234,7 @@ def test_call():
     parsed = FieldOperatorParser.apply_to_function(call)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.deref_(im.lift_(im.lambda__("inp")(im.call_("identity")("inp")))("inp"))
+    reference = im.lift_(im.lambda__("inp")(im.call_("identity")("inp")))("inp")
 
     assert lowered.expr == reference
 
@@ -258,7 +252,7 @@ def test_temp_tuple():
     tuple_expr = im.lift_(im.lambda__("a", "b")(im.make_tuple_(im.deref_("a"), im.deref_("b"))))(
         "a", "b"
     )
-    reference = im.let("tmp__0", tuple_expr)(im.deref_("tmp__0"))
+    reference = im.let("tmp__0", tuple_expr)("tmp__0")
 
     assert lowered.expr == reference
 
@@ -270,7 +264,7 @@ def test_unary_not():
     parsed = FieldOperatorParser.apply_to_function(unary_not)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.deref_(im.lift_(im.lambda__("cond")(im.not__(im.deref_("cond"))))("cond"))
+    reference = im.lift_(im.lambda__("cond")(im.not__(im.deref_("cond"))))("cond")
 
     assert lowered.expr == reference
 
@@ -282,9 +276,7 @@ def test_binary_plus():
     parsed = FieldOperatorParser.apply_to_function(plus)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.deref_(
-        im.lift_(im.lambda__("a", "b")(im.plus_(im.deref_("a"), im.deref_("b"))))("a", "b")
-    )
+    reference = im.lift_(im.lambda__("a", "b")(im.plus_(im.deref_("a"), im.deref_("b"))))("a", "b")
 
     assert lowered.expr == reference
 
@@ -296,8 +288,8 @@ def test_add_scalar_literal_to_field():
     parsed = FieldOperatorParser.apply_to_function(scalar_plus_field)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.deref_(
-        im.lift_(im.lambda__("a")(im.plus_(im.literal_("2.0", "float64"), im.deref_("a"))))("a")
+    reference = im.lift_(im.lambda__("a")(im.plus_(im.literal_("2.0", "float64"), im.deref_("a"))))(
+        "a"
     )
 
     assert lowered.expr == reference
@@ -317,7 +309,7 @@ def test_add_scalar_literals():
             im.literal_("1", "int32"),
             im.literal_("1", "int32"),
         ),
-    )(im.deref_(im.lift_(im.lambda__("a")(im.plus_(im.deref_("a"), "tmp__0")))("a")))
+    )(im.lift_(im.lambda__("a")(im.plus_(im.deref_("a"), "tmp__0")))("a"))
 
     assert lowered.expr == reference
 
@@ -329,8 +321,8 @@ def test_binary_mult():
     parsed = FieldOperatorParser.apply_to_function(mult)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.deref_(
-        im.lift_(im.lambda__("a", "b")(im.multiplies_(im.deref_("a"), im.deref_("b"))))("a", "b")
+    reference = im.lift_(im.lambda__("a", "b")(im.multiplies_(im.deref_("a"), im.deref_("b"))))(
+        "a", "b"
     )
 
     assert lowered.expr == reference
@@ -343,9 +335,7 @@ def test_binary_minus():
     parsed = FieldOperatorParser.apply_to_function(minus)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.deref_(
-        im.lift_(im.lambda__("a", "b")(im.minus_(im.deref_("a"), im.deref_("b"))))("a", "b")
-    )
+    reference = im.lift_(im.lambda__("a", "b")(im.minus_(im.deref_("a"), im.deref_("b"))))("a", "b")
 
     assert lowered.expr == reference
 
@@ -357,8 +347,8 @@ def test_binary_div():
     parsed = FieldOperatorParser.apply_to_function(division)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.deref_(
-        im.lift_(im.lambda__("a", "b")(im.divides_(im.deref_("a"), im.deref_("b"))))("a", "b")
+    reference = im.lift_(im.lambda__("a", "b")(im.divides_(im.deref_("a"), im.deref_("b"))))(
+        "a", "b"
     )
 
     assert lowered.expr == reference
@@ -371,9 +361,7 @@ def test_binary_and():
     parsed = FieldOperatorParser.apply_to_function(bit_and)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.deref_(
-        im.lift_(im.lambda__("a", "b")(im.and__(im.deref_("a"), im.deref_("b"))))("a", "b")
-    )
+    reference = im.lift_(im.lambda__("a", "b")(im.and__(im.deref_("a"), im.deref_("b"))))("a", "b")
 
     assert lowered.expr == reference
 
@@ -385,8 +373,8 @@ def test_scalar_and():
     parsed = FieldOperatorParser.apply_to_function(scalar_and)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.deref_(
-        im.lift_(im.lambda__("a")(im.and__(im.deref_("a"), im.literal_("False", "bool"))))("a")
+    reference = im.lift_(im.lambda__("a")(im.and__(im.deref_("a"), im.literal_("False", "bool"))))(
+        "a"
     )
 
     assert lowered.expr == reference
@@ -399,9 +387,7 @@ def test_binary_or():
     parsed = FieldOperatorParser.apply_to_function(bit_or)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.deref_(
-        im.lift_(im.lambda__("a", "b")(im.or__(im.deref_("a"), im.deref_("b"))))("a", "b")
-    )
+    reference = im.lift_(im.lambda__("a", "b")(im.or__(im.deref_("a"), im.deref_("b"))))("a", "b")
 
     assert lowered.expr == reference
 
@@ -425,8 +411,8 @@ def test_compare_gt():
     parsed = FieldOperatorParser.apply_to_function(comp_gt)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.deref_(
-        im.lift_(im.lambda__("a", "b")(im.greater_(im.deref_("a"), im.deref_("b"))))("a", "b")
+    reference = im.lift_(im.lambda__("a", "b")(im.greater_(im.deref_("a"), im.deref_("b"))))(
+        "a", "b"
     )
 
     assert lowered.expr == reference
@@ -439,9 +425,7 @@ def test_compare_lt():
     parsed = FieldOperatorParser.apply_to_function(comp_lt)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.deref_(
-        im.lift_(im.lambda__("a", "b")(im.less_(im.deref_("a"), im.deref_("b"))))("a", "b")
-    )
+    reference = im.lift_(im.lambda__("a", "b")(im.less_(im.deref_("a"), im.deref_("b"))))("a", "b")
 
     assert lowered.expr == reference
 
@@ -453,9 +437,7 @@ def test_compare_eq():
     parsed = FieldOperatorParser.apply_to_function(comp_eq)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.deref_(
-        im.lift_(im.lambda__("a", "b")(im.eq_(im.deref_("a"), im.deref_("b"))))("a", "b")
-    )
+    reference = im.lift_(im.lambda__("a", "b")(im.eq_(im.deref_("a"), im.deref_("b"))))("a", "b")
 
     assert lowered.expr == reference
 
@@ -469,24 +451,22 @@ def test_compare_chain():
     parsed = FieldOperatorParser.apply_to_function(compare_chain)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.deref_(
-        im.lift_(
-            im.lambda__("a", "b", "c")(
-                im.and__(
-                    im.deref_(
-                        im.lift_(
-                            im.lambda__("a", "b")(im.greater_(im.deref_("a"), im.deref_("b")))
-                        )("a", "b")
-                    ),
-                    im.deref_(
-                        im.lift_(
-                            im.lambda__("b", "c")(im.greater_(im.deref_("b"), im.deref_("c")))
-                        )("b", "c")
-                    ),
-                )
+    reference = im.lift_(
+        im.lambda__("a", "b", "c")(
+            im.and__(
+                im.deref_(
+                    im.lift_(im.lambda__("a", "b")(im.greater_(im.deref_("a"), im.deref_("b"))))(
+                        "a", "b"
+                    )
+                ),
+                im.deref_(
+                    im.lift_(im.lambda__("b", "c")(im.greater_(im.deref_("b"), im.deref_("c"))))(
+                        "b", "c"
+                    )
+                ),
             )
-        )("a", "b", "c")
-    )
+        )
+    )("a", "b", "c")
 
     assert lowered.expr == reference
 
@@ -498,14 +478,12 @@ def test_reduction_lowering_simple():
     parsed = FieldOperatorParser.apply_to_function(reduction)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.deref_(
-        im.lift_(
-            im.call_("reduce")(
-                im.lambda__("acc", "edge_f__0")(im.plus_("acc", "edge_f__0")),
-                0,
-            )
-        )(im.shift_("V2E")("edge_f"))
-    )
+    reference = im.lift_(
+        im.call_("reduce")(
+            im.lambda__("acc", "edge_f__0")(im.plus_("acc", "edge_f__0")),
+            0,
+        )
+    )(im.shift_("V2E")("edge_f"))
 
     assert lowered.expr == reference
 
@@ -519,21 +497,19 @@ def test_reduction_lowering_expr():
     lowered = FieldOperatorLowering.apply(parsed)
 
     reference = im.let("e1_nbh__0", im.shift_("V2E")("e1"))(
-        im.deref_(
-            im.lift_(
-                im.call_("reduce")(
-                    im.lambda__("acc", "e1_nbh__0__0", "e2__1")(
-                        im.plus_(
-                            "acc",
-                            im.multiplies_(
-                                im.literal_("1.1", "float64"), im.plus_("e1_nbh__0__0", "e2__1")
-                            ),
-                        )
-                    ),
-                    0,
-                )
-            )("e1_nbh__0", "e2")
-        )
+        im.lift_(
+            im.call_("reduce")(
+                im.lambda__("acc", "e1_nbh__0__0", "e2__1")(
+                    im.plus_(
+                        "acc",
+                        im.multiplies_(
+                            im.literal_("1.1", "float64"), im.plus_("e1_nbh__0__0", "e2__1")
+                        ),
+                    )
+                ),
+                0,
+            )
+        )("e1_nbh__0", "e2")
     )
 
     assert lowered.expr == reference
diff --git a/tests/functional_tests/ffront_tests/test_foast_parsing.py b/tests/functional_tests/ffront_tests/test_foast_parsing.py
index 00eb16af..02a64eb6 100644
--- a/tests/functional_tests/ffront_tests/test_foast_parsing.py
+++ b/tests/functional_tests/ffront_tests/test_foast_parsing.py
@@ -117,7 +117,7 @@ def test_return_type():
 
     parsed = FieldOperatorParser.apply_to_function(rettype)
 
-    assert parsed.body[-1].value.type == ts.FieldType(
+    assert parsed.body.stmts[-1].value.type == ts.FieldType(
         dims=[TDim],
         dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64, shape=None),
     )
@@ -131,7 +131,7 @@ def test_invalid_syntax_no_return():
 
     with pytest.raises(
         FieldOperatorSyntaxError,
-        match="Field operator must return a field expression on the last line!",
+        match="Function must return a value, but no return statement was found\.",
     ):
         _ = FieldOperatorParser.apply_to_function(no_return)
 
@@ -157,7 +157,7 @@ def test_temp_assignment():
 
     parsed = FieldOperatorParser.apply_to_function(copy_field)
 
-    assert parsed.annex.symtable["tmp__0"].type == ts.FieldType(
+    assert parsed.body.annex.symtable["tmp__0"].type == ts.FieldType(
         dims=[TDim],
         dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64, shape=None),
     )
@@ -180,7 +180,7 @@ def test_binary_pow():
 
     parsed = FieldOperatorParser.apply_to_function(power)
 
-    assert parsed.body[-1].value.type == ts.FieldType(
+    assert parsed.body.stmts[-1].value.type == ts.FieldType(
         dims=[TDim],
         dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64, shape=None),
     )
@@ -192,7 +192,7 @@ def test_binary_mod():
 
     parsed = FieldOperatorParser.apply_to_function(modulo)
 
-    assert parsed.body[-1].value.type == ts.FieldType(
+    assert parsed.body.stmts[-1].value.type == ts.FieldType(
         dims=[TDim],
         dtype=ts.ScalarType(kind=ts.ScalarKind.INT64, shape=None),
     )
@@ -226,7 +226,7 @@ def test_bool_xor():
 
     parsed = FieldOperatorParser.apply_to_function(bool_xor)
 
-    assert parsed.body[-1].value.type == ts.FieldType(
+    assert parsed.body.stmts[-1].value.type == ts.FieldType(
         dims=[TDim],
         dtype=ts.ScalarType(kind=ts.ScalarKind.BOOL, shape=None),
     )
@@ -238,7 +238,7 @@ def test_unary_tilde():
 
     parsed = FieldOperatorParser.apply_to_function(unary_tilde)
 
-    assert parsed.body[-1].value.type == ts.FieldType(
+    assert parsed.body.stmts[-1].value.type == ts.FieldType(
         dims=[TDim],
         dtype=ts.ScalarType(kind=ts.ScalarKind.BOOL, shape=None),
     )
@@ -333,7 +333,7 @@ def test_astype():
 
     parsed = FieldOperatorParser.apply_to_function(astype_fieldop)
 
-    assert parsed.body[-1].value.type == ts.FieldType(
+    assert parsed.body.stmts[-1].value.type == ts.FieldType(
         dims=[TDim],
         dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64, shape=None),
     )
@@ -354,22 +354,25 @@ def test_closure_symbols():
         return a, b
 
     parsed = FieldOperatorParser.apply_to_function(operator_with_refs)
-    assert "nonlocals_unused" not in parsed.annex.symtable
-    assert "nonlocals" not in parsed.annex.symtable
+    assert "nonlocals_unreferenced" not in {**parsed.annex.symtable, **parsed.body.annex.symtable}
+    assert "nonlocals" not in {**parsed.annex.symtable, **parsed.body.annex.symtable}
 
     pattern_node = P(
         foast.FunctionDefinition,
-        body=[
-            P(
-                foast.Assign,
-                value=P(foast.BinOp, right=P(foast.Constant, value=nonlocals.float_value)),
-            ),
-            P(
-                foast.Assign,
-                value=P(foast.BinOp, right=P(foast.Constant, value=nonlocals.np_value)),
-            ),
-            P(foast.Return),
-        ],
+        body=P(
+            foast.BlockStmt,
+            stmts=[
+                P(
+                    foast.Assign,
+                    value=P(foast.BinOp, right=P(foast.Constant, value=nonlocals.float_value)),
+                ),
+                P(
+                    foast.Assign,
+                    value=P(foast.BinOp, right=P(foast.Constant, value=nonlocals.np_value)),
+                ),
+                P(foast.Return),
+            ],
+        ),
     )
     assert pattern_node.match(parsed, raise_exception=True)
 
diff --git a/tests/functional_tests/ffront_tests/test_foast_pretty_printer.py b/tests/functional_tests/ffront_tests/test_foast_pretty_printer.py
index 91671f4e..39659f7d 100644
--- a/tests/functional_tests/ffront_tests/test_foast_pretty_printer.py
+++ b/tests/functional_tests/ffront_tests/test_foast_pretty_printer.py
@@ -62,12 +62,14 @@ def test_scanop():
 
     @scan_operator(axis=KDim, forward=False, init=1.0)
     def scan(inp: int64) -> int64:
+        foo = inp
         return inp
 
     expected = textwrap.dedent(
         """
         @scan_operator(axis=Dimension(value="KDim", kind=DimensionKind.VERTICAL), forward=False, init=1.0)
         def scan(inp: int64) -> int64:
+          foo__0 = inp
           return inp
         """
     ).strip()
diff --git a/tests/functional_tests/ffront_tests/test_math_builtin_execution.py b/tests/functional_tests/ffront_tests/test_math_builtin_execution.py
index 2360f392..ec158678 100644
--- a/tests/functional_tests/ffront_tests/test_math_builtin_execution.py
+++ b/tests/functional_tests/ffront_tests/test_math_builtin_execution.py
@@ -74,17 +74,20 @@ def make_builtin_field_operator(builtin_name: str):
         id=builtin_name + "_field_operator",
         definition=foast.FunctionDefinition(
             id=builtin_name + "_field_operator",
-            body=[
-                foast.Return(
-                    value=foast.Call(
-                        func=foast.Name(id=builtin_name, location=loc),
-                        args=args,
-                        kwargs={},
+            body=foast.BlockStmt(
+                stmts=[
+                    foast.Return(
+                        value=foast.Call(
+                            func=foast.Name(id=builtin_name, location=loc),
+                            args=args,
+                            kwargs={},
+                            location=loc,
+                        ),
                         location=loc,
-                    ),
-                    location=loc,
-                )
-            ],
+                    )
+                ],
+                location=loc,
+            ),
             closure_vars=closure_var_symbols,
             params=params,
             location=loc,
diff --git a/tests/functional_tests/ffront_tests/test_type_deduction.py b/tests/functional_tests/ffront_tests/test_type_deduction.py
index c117e83d..912c48bf 100644
--- a/tests/functional_tests/ffront_tests/test_type_deduction.py
+++ b/tests/functional_tests/ffront_tests/test_type_deduction.py
@@ -300,11 +300,11 @@ def test_unpack_assign():
 
     parsed = FieldOperatorParser.apply_to_function(unpack_explicit_tuple)
 
-    assert parsed.annex.symtable["tmp_a__0"].type == ts.FieldType(
+    assert parsed.body.annex.symtable["tmp_a__0"].type == ts.FieldType(
         dims=[TDim],
         dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64, shape=None),
     )
-    assert parsed.annex.symtable["tmp_b__0"].type == ts.FieldType(
+    assert parsed.body.annex.symtable["tmp_b__0"].type == ts.FieldType(
         dims=[TDim],
         dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64, shape=None),
     )
@@ -362,7 +362,7 @@ def test_assign_tuple():
 
     parsed = FieldOperatorParser.apply_to_function(temp_tuple)
 
-    assert parsed.annex.symtable["tmp__0"].type == ts.TupleType(
+    assert parsed.body.annex.symtable["tmp__0"].type == ts.TupleType(
         types=[
             ts.FieldType(
                 dims=[TDim],
@@ -456,7 +456,7 @@ def test_remap(remap_setup):
 
     parsed = FieldOperatorParser.apply_to_function(remap_fo)
 
-    assert parsed.body[0].value.type == ts.FieldType(
+    assert parsed.body.stmts[0].value.type == ts.FieldType(
         dims=[Y], dtype=ts.ScalarType(kind=ts.ScalarKind.INT64)
     )
 
@@ -469,7 +469,7 @@ def test_remap_nbfield(remap_setup):
 
     parsed = FieldOperatorParser.apply_to_function(remap_fo)
 
-    assert parsed.body[0].value.type == ts.FieldType(
+    assert parsed.body.stmts[0].value.type == ts.FieldType(
         dims=[Y, Y2XDim], dtype=ts.ScalarType(kind=ts.ScalarKind.INT64)
     )
 
@@ -482,7 +482,7 @@ def test_remap_reduce(remap_setup):
 
     parsed = FieldOperatorParser.apply_to_function(remap_fo)
 
-    assert parsed.body[0].value.type == ts.FieldType(
+    assert parsed.body.stmts[0].value.type == ts.FieldType(
         dims=[Y], dtype=ts.ScalarType(kind=ts.ScalarKind.INT64)
     )
 
@@ -495,7 +495,7 @@ def test_remap_reduce_sparse(remap_setup):
 
     parsed = FieldOperatorParser.apply_to_function(remap_fo)
 
-    assert parsed.body[0].value.type == ts.FieldType(
+    assert parsed.body.stmts[0].value.type == ts.FieldType(
         dims=[Y], dtype=ts.ScalarType(kind=ts.ScalarKind.INT64)
     )
 
@@ -521,7 +521,7 @@ def test_broadcast_multi_dim():
 
     parsed = FieldOperatorParser.apply_to_function(simple_broadcast)
 
-    assert parsed.body[0].value.type == ts.FieldType(
+    assert parsed.body.stmts[0].value.type == ts.FieldType(
         dims=[ADim, BDim, CDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64)
     )
 
@@ -565,7 +565,7 @@ def test_where_dim():
 
     parsed = FieldOperatorParser.apply_to_function(simple_where)
 
-    assert parsed.body[0].value.type == ts.FieldType(
+    assert parsed.body.stmts[0].value.type == ts.FieldType(
         dims=[ADim, BDim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64)
     )
 
@@ -578,7 +578,7 @@ def test_where_broadcast_dim():
 
     parsed = FieldOperatorParser.apply_to_function(simple_where)
 
-    assert parsed.body[0].value.type == ts.FieldType(
+    assert parsed.body.stmts[0].value.type == ts.FieldType(
         dims=[ADim], dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64)
     )
 
@@ -591,7 +591,7 @@ def test_where_tuple_dim():
 
     parsed = FieldOperatorParser.apply_to_function(tuple_where)
 
-    assert parsed.body[0].value.type == ts.TupleType(
+    assert parsed.body.stmts[0].value.type == ts.TupleType(
         types=[
             ts.TupleType(
                 types=[
@@ -633,7 +633,7 @@ def test_where_mixed_dims():
 
     parsed = FieldOperatorParser.apply_to_function(tuple_where_mix_dims)
 
-    assert parsed.body[0].value.type == ts.TupleType(
+    assert parsed.body.stmts[0].value.type == ts.TupleType(
         types=[
             ts.TupleType(
                 types=[
@@ -661,7 +661,7 @@ def test_astype_dtype():
 
     parsed = FieldOperatorParser.apply_to_function(simple_astype)
 
-    assert parsed.body[0].value.type == ts.FieldType(
+    assert parsed.body.stmts[0].value.type == ts.FieldType(
         dims=[ADim], dtype=ts.ScalarType(kind=ts.ScalarKind.BOOL)
     )