Repository: GridTools/gt4py
Pull request #1026 (2022-11-02 16:08:03)

Functional: scan allowing for scalar inputs
Allow scalar floats in scan_operators:

```python
@scan_operator(axis=KDim, forward=True, init=(0.0))
    def _scan_scalar(carry: float, qc_in: float, scalar: float) -> float:
        qc = qc_in + carry + scalar
        return qc

@program
def scan_scalar(qc: Field[[IDim, KDim], float], scalar: float):
    _scan_scalar(qc, scalar, out=qc)
```


Patch:
diff --git a/.gitpod/.vscode/settings.json b/.gitpod/.vscode/settings.json
index 7a666b7f..1591bdb2 100644
--- a/.gitpod/.vscode/settings.json
+++ b/.gitpod/.vscode/settings.json
@@ -1,5 +1,6 @@
 {
     "python.formatting.provider": "black",
     "python.testing.pytestEnabled": true,
-    "python.defaultInterpreterPath": "/workspace/gt4py/.venv/bin/python"
-}
\ No newline at end of file
+    "python.defaultInterpreterPath": "/workspace/gt4py/.venv/bin/python",
+    "python.terminal.activateEnvironment": true
+}
diff --git a/setup.cfg b/setup.cfg
index a0070073..93aaaeac 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -132,6 +132,7 @@ max-complexity = 15
 doctests = true
 ignore =
     B008    # Do not perform function calls in argument defaults
+    B028    # Consider replacing f"'{foo}'" with f"{foo!r}"  # TODO: review this ignore
     D1      # Public code object needs docstring
     DAR     # Disable dargling errors by default
     E203    # Whitespace before ':' (black formatter breaks this sometimes)
@@ -237,9 +238,6 @@ ignore_errors = True
 [mypy-functional.iterator.transforms.global_tmps]
 ignore_errors = True
 
-[mypy-functional.iterator.transforms.remap_symbols]
-ignore_errors = True
-
 [mypy-functional.iterator.transforms.popup_tmps]
 ignore_errors = True
 
diff --git a/src/functional/common.py b/src/functional/common.py
index 331eab31..fec59ab1 100644
--- a/src/functional/common.py
+++ b/src/functional/common.py
@@ -17,7 +17,17 @@ from __future__ import annotations
 import enum
 from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import Any, Generic, Optional, Protocol, TypeVar, runtime_checkable
+from typing import (
+    Any,
+    Generic,
+    Optional,
+    Protocol,
+    SupportsFloat,
+    SupportsInt,
+    TypeAlias,
+    TypeVar,
+    runtime_checkable,
+)
 
 import numpy as np
 
@@ -28,6 +38,8 @@ DimT = TypeVar("DimT", bound="Dimension")
 DimsT = TypeVar("DimsT", bound=Sequence["Dimension"])
 DT = TypeVar("DT", bound="DType")
 
+Scalar: TypeAlias = SupportsInt | SupportsFloat | np.int32 | np.int64 | np.float32 | np.float64
+
 
 @enum.unique
 class DimensionKind(StrEnum):
diff --git a/src/functional/ffront/decorator.py b/src/functional/ffront/decorator.py
index 790760de..8d2c4551 100644
--- a/src/functional/ffront/decorator.py
+++ b/src/functional/ffront/decorator.py
@@ -26,14 +26,13 @@ import types
 import typing
 import warnings
 from collections.abc import Callable, Iterable
-from typing import Generator, Generic, SupportsFloat, SupportsInt, TypeAlias, TypeVar
+from typing import Generator, Generic, TypeVar
 
-import numpy as np
 from devtools import debug
 
 from eve.extended_typing import Any, Optional
 from eve.utils import UIDGenerator
-from functional.common import DimensionKind, GridType, GTTypeError
+from functional.common import DimensionKind, GridType, GTTypeError, Scalar
 from functional.ffront import (
     dialect_ast_enums,
     field_operator_ast as foast,
@@ -59,8 +58,6 @@ from functional.program_processors import processor_interface as ppi
 from functional.program_processors.runners import roundtrip
 
 
-Scalar: TypeAlias = SupportsInt | SupportsFloat | np.int32 | np.int64 | np.float32 | np.float64
-
 DEFAULT_BACKEND: Callable = roundtrip.executor
 
 
@@ -166,6 +163,8 @@ def _field_constituents_shape_and_dims(
             yield (arg.shape, dims)
         else:
             yield (None, dims)
+    elif isinstance(arg_type, ts.ScalarType):
+        yield (None, [])
     else:
         raise ValueError("Expected `FieldType` or `TupleType` thereof.")
 
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..63ce98bc 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,
@@ -215,7 +228,7 @@ class FieldOperatorTypeDeduction(traits.VisitorWithSymbolTableTrait, NodeTransla
         new_forward = self.visit(node.forward, **kwargs)
         if not new_forward.type.kind == ts.ScalarKind.BOOL:
             raise FieldOperatorTypeDeductionError.from_foast_node(
-                node, msg=f"Argument `forward` to scan operator `{node.id}` must" f"be a boolean."
+                node, msg=f"Argument `forward` to scan operator `{node.id}` must be a boolean."
             )
         new_init = self.visit(node.init, **kwargs)
         if not all(
@@ -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..92141940 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"))
@@ -175,9 +175,13 @@ class FieldOperatorLowering(NodeTranslator):
         #  rest of the lowering inside the body. See ADR-0002 for more details.
         new_body = func_definition.expr
         for i, param in enumerate(func_definition.params):
-            if isinstance(node.definition.params[i].type, ts.ScalarType):
+            if iterator_type_kind(node.definition.params[i].type) == ITIRTypeKind.VALUE:
                 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))
diff --git a/src/functional/ffront/past_to_itir.py b/src/functional/ffront/past_to_itir.py
index 2bc5d88f..c12d7dc7 100644
--- a/src/functional/ffront/past_to_itir.py
+++ b/src/functional/ffront/past_to_itir.py
@@ -95,7 +95,7 @@ class ProgramLowering(traits.VisitorWithSymbolTableTrait, NodeTranslator):
         """Generate symbols for each field param and dimension."""
         size_params = []
         for param in node.params:
-            if type_info.is_field_type_or_tuple_of_field_type(param.type):
+            if type_info.is_type_or_tuple_of_type(param.type, ts.FieldType):
                 fields_dims: list[list[Dimension]] = (
                     type_info.primitive_constituents(param.type).getattr("dims").to_list()
                 )
@@ -131,7 +131,7 @@ class ProgramLowering(traits.VisitorWithSymbolTableTrait, NodeTranslator):
 
     def _visit_stencil_call(self, node: past.Call, **kwargs) -> itir.StencilClosure:
         assert isinstance(node.kwargs["out"].type, ts.TypeSpec)
-        assert type_info.is_field_type_or_tuple_of_field_type(node.kwargs["out"].type)
+        assert type_info.is_type_or_tuple_of_type(node.kwargs["out"].type, ts.FieldType)
 
         domain = node.kwargs.get("domain", None)
         output, lowered_domain = self._visit_stencil_call_out_arg(
diff --git a/src/functional/ffront/type_info.py b/src/functional/ffront/type_info.py
index 7b5636e5..a617fe67 100644
--- a/src/functional/ffront/type_info.py
+++ b/src/functional/ffront/type_info.py
@@ -28,11 +28,11 @@ from functional.type_system.type_info import (  # noqa: F401
     is_arithmetic as is_arithmetic,
     is_concrete as is_concrete,
     is_concretizable as is_concretizable,
-    is_field_type_or_tuple_of_field_type as is_field_type_or_tuple_of_field_type,
     is_floating_point as is_floating_point,
     is_integral as is_integral,
     is_logical as is_logical,
     is_number as is_number,
+    is_type_or_tuple_of_type as is_type_or_tuple_of_type,
     primitive_constituents as primitive_constituents,
     promote as promote,
     promote_dims as promote_dims,
@@ -52,17 +52,29 @@ def _is_zero_dim_field(field: ts.TypeSpec) -> bool:
 
 
 def promote_zero_dims(
-    args: list[ts.TypeSpec], function_type: ts.FieldOperatorType | ts.ProgramType
+    args: list[ts.TypeSpec], function_type: ts.FieldOperatorType | ts.ProgramType | ts.FunctionType
 ) -> list:
     """Promote arg types to zero dimensional fields if compatible and required by function signature."""
-    new_args = args.copy()
+    new_args = []
     for arg_i, arg in enumerate(args):
-        def_type = function_type.definition.args[arg_i]
-        if _is_zero_dim_field(def_type) and is_number(arg):
-            if extract_dtype(def_type) == extract_dtype(arg):
-                new_args[arg_i] = def_type
-            else:
-                raise GTTypeError(f"{arg} is not compatible with {def_type}.")
+        def_type = (
+            function_type.args[arg_i]
+            if isinstance(function_type, ts.FunctionType)
+            else function_type.definition.args[arg_i]
+        )
+
+        def _as_field(arg: ts.TypeSpec, path: tuple):
+            el_def_type = reduce(lambda type_, idx: type_.types[idx], path, def_type)  # noqa: B023
+
+            if _is_zero_dim_field(el_def_type) and is_number(arg):
+                if extract_dtype(el_def_type) == extract_dtype(arg):
+                    return el_def_type
+                else:
+                    raise GTTypeError(f"{arg} is not compatible with {el_def_type}.")
+            return arg
+
+        new_args.append(apply_to_primitive_constituents(arg, _as_field, with_path_arg=True))
+
     return new_args
 
 
@@ -89,8 +101,8 @@ def function_signature_incompatibilities_fieldop(
 def function_signature_incompatibilities_scanop(
     scanop_type: ts.ScanOperatorType, args: list[ts.TypeSpec], kwargs: dict[str, ts.TypeSpec]
 ) -> Iterator[str]:
-    if not all(is_field_type_or_tuple_of_field_type(arg) for arg in args):
-        yield "Arguments to scan operator must be fields or tuples thereof."
+    if not all(is_type_or_tuple_of_type(arg, (ts.ScalarType, ts.FieldType)) for arg in args):
+        yield "Arguments to scan operator must be fields, scalars or tuples thereof."
         return
 
     arg_dims = [extract_dims(el) for arg in args for el in primitive_constituents(arg)]
@@ -132,7 +144,9 @@ def function_signature_incompatibilities_scanop(
         returns=ts.DeferredType(constraint=None),
     )
 
-    yield from function_signature_incompatibilities(function_type, args, kwargs)
+    yield from function_signature_incompatibilities(
+        function_type, promote_zero_dims(args, function_type), kwargs
+    )
 
 
 @function_signature_incompatibilities.register
diff --git a/src/functional/iterator/embedded.py b/src/functional/iterator/embedded.py
index fd84c6d2..6ccc79a8 100644
--- a/src/functional/iterator/embedded.py
+++ b/src/functional/iterator/embedded.py
@@ -163,6 +163,11 @@ class MutableLocatedField(LocatedField, Protocol):
         ...
 
 
+_column_range: Optional[
+    range
+] = None  # TODO this is a bit ugly, alternative: pass scan range via iterator
+
+
 class Column(np.lib.mixins.NDArrayOperatorsMixin):
     """Represents a column when executed in column mode (`column_axis != None`).
 
@@ -170,13 +175,13 @@ class Column(np.lib.mixins.NDArrayOperatorsMixin):
     and simplify dispatching in iterator ir builtins.
     """
 
-    def __init__(self, kstart: int, data: np.ndarray) -> None:
+    def __init__(self, kstart: int, data: np.ndarray | Scalar) -> None:
         self.kstart = kstart
-        self.data = data
+        assert isinstance(data, (np.ndarray, Scalar))  # type: ignore # mypy bug
+        self.data = data if isinstance(data, np.ndarray) else np.full(len(_column_range), data)  # type: ignore[arg-type]
 
     def __getitem__(self, i: int) -> Any:
         result = self.data[i - self.kstart]
-        # if the element type is a tuple return a regular type instead of a
         #  numpy type
         if self.data.dtype.names:
             return tuple(result)
@@ -702,7 +707,7 @@ def _get_axes(
 
 
 def _make_tuple(
-    field_or_tuple: LocatedField | tuple,  # arbitrary nesting of tuples of LocatedField
+    field_or_tuple: LocatedField | tuple,
     indices: FieldIndexOrIndices,
     *,
     as_column: bool = False,
@@ -999,7 +1004,6 @@ class ScanArgIterator:
     def deref(self) -> Any:
         if not self.can_deref():
             return _UNDEFINED
-        # TODO(tehrengruber): _make_tuple is for fields
         return _make_tuple(self.wrapped_iter.deref(), self.k_pos)
 
     def can_deref(self) -> bool:
@@ -1108,11 +1112,6 @@ def as_tuple_field(field):
     return field
 
 
-_column_range: Optional[
-    range
-] = None  # TODO this is a bit ugly, alternative: pass scan range via iterator
-
-
 def _column_dtype(elem: Any) -> np.dtype:
     if isinstance(elem, tuple):
         return np.dtype([(f"f{i}", _column_dtype(e)) for i, e in enumerate(elem)])
diff --git a/src/functional/iterator/transforms/remap_symbols.py b/src/functional/iterator/transforms/remap_symbols.py
index af11b61b..68ec6660 100644
--- a/src/functional/iterator/transforms/remap_symbols.py
+++ b/src/functional/iterator/transforms/remap_symbols.py
@@ -1,6 +1,6 @@
 from typing import Any, Dict, Optional, Set
 
-from eve import NodeTranslator
+from eve import NodeTranslator, SymbolTableTrait
 from functional.iterator import ir
 
 
@@ -17,7 +17,7 @@ class RemapSymbolRefs(NodeTranslator):
         )
 
     def generic_visit(self, node: ir.Node, **kwargs: Any):  # type: ignore[override]
-        assert isinstance(node, ir.SymbolTableTrait) == isinstance(
+        assert isinstance(node, SymbolTableTrait) == isinstance(
             node, ir.Lambda
         ), "found unexpected new symbol scope"
         return super().generic_visit(node, **kwargs)
@@ -41,7 +41,7 @@ class RenameSymbols(NodeTranslator):
     def generic_visit(  # type: ignore[override]
         self, node: ir.Node, *, name_map: Dict[str, str], active: Optional[Set[str]] = None
     ):
-        if isinstance(node, ir.SymbolTableTrait):
+        if isinstance(node, SymbolTableTrait):
             if active is None:
                 active = set()
             active = active | set(node.annex.symtable)
diff --git a/src/functional/type_system/type_info.py b/src/functional/type_system/type_info.py
index 758fe9bd..1e61acab 100644
--- a/src/functional/type_system/type_info.py
+++ b/src/functional/type_system/type_info.py
@@ -227,22 +227,24 @@ def is_arithmetic(symbol_type: ts.TypeSpec) -> bool:
     return is_floating_point(symbol_type) or is_integral(symbol_type)
 
 
-def is_field_type_or_tuple_of_field_type(type_: ts.TypeSpec) -> bool:
+def is_type_or_tuple_of_type(type_: ts.TypeSpec, expected_type: type | tuple) -> bool:
     """
-     Return True if ``type_`` is FieldType or FieldType nested in TupleType.
+    Return True if ``type_`` matches any of the expected.
 
-     Examples:
-     ---------
+    Examples:
+    ---------
     >>> scalar_type = ts.ScalarType(kind=ts.ScalarKind.INT)
     >>> field_type = ts.FieldType(dims=[], dtype=scalar_type)
-     >>> is_field_type_or_tuple_of_field_type(field_type)
-     True
-     >>> is_field_type_or_tuple_of_field_type(ts.TupleType(types=[field_type, field_type]))
-     True
-     >>> is_field_type_or_tuple_of_field_type(ts.TupleType(types=[field_type, scalar_type]))
-     False
+    >>> is_type_or_tuple_of_type(field_type, ts.FieldType)
+    True
+    >>> is_type_or_tuple_of_type(ts.TupleType(types=[scalar_type, field_type]), (ts.ScalarType, ts.FieldType))
+    True
+    >>> is_type_or_tuple_of_type(scalar_type, ts.FieldType)
+    False
+    >>> is_type_or_tuple_of_type(ts.TupleType(types=[scalar_type, field_type]), ts.FieldType)
+    False
     """
-    return all(isinstance(t, ts.FieldType) for t in primitive_constituents(type_))
+    return all(isinstance(t, expected_type) for t in primitive_constituents(type_))
 
 
 def extract_dims(symbol_type: ts.TypeSpec) -> list[Dimension]:


Test patch:
diff --git a/tests/functional_tests/ffront_tests/test_execution.py b/tests/functional_tests/ffront_tests/test_execution.py
index d4d4ab8d..25ab0169 100644
--- a/tests/functional_tests/ffront_tests/test_execution.py
+++ b/tests/functional_tests/ffront_tests/test_execution.py
@@ -15,6 +15,7 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 from functools import reduce
 
+import numpy as np
 import pytest as pytest
 
 from functional.ffront.decorator import field_operator, program, scan_operator
@@ -225,6 +226,49 @@ def test_scalar_arg_with_field(fieldview_backend):
     assert np.allclose(ref, out.array())
 
 
+def test_scalar_scan():
+    size = 10
+    KDim = Dimension("K", kind=DimensionKind.VERTICAL)
+    qc = np_as_located_field(IDim, KDim)(np.zeros((size, size)))
+    scalar = 1.0
+    expected = np.full((size, size), np.arange(start=1, stop=11, step=1).astype(float64))
+
+    @scan_operator(axis=KDim, forward=True, init=(0.0))
+    def _scan_scalar(carry: float, qc_in: float, scalar: float) -> float:
+        qc = qc_in + carry + scalar
+        return qc
+
+    @program
+    def scan_scalar(qc: Field[[IDim, KDim], float], scalar: float):
+        _scan_scalar(qc, scalar, out=qc)
+
+    scan_scalar(qc, scalar, offset_provider={})
+    assert np.allclose(np.asarray(qc), expected)
+
+
+def test_tuple_scalar_scan():
+    size = 10
+    KDim = Dimension("K", kind=DimensionKind.VERTICAL)
+    qc = np_as_located_field(IDim, KDim)(np.zeros((size, size)))
+    tuple_scalar = (1.0, (1.0, 0.0))
+    expected = np.full((size, size), np.arange(start=1, stop=11, step=1).astype(float64))
+
+    @scan_operator(axis=KDim, forward=True, init=0.0)
+    def _scan_tuple_scalar(
+        state: float, qc_in: float, tuple_scalar: tuple[float, tuple[float, float]]
+    ) -> float:
+        return (qc_in + state + tuple_scalar[1][0] + tuple_scalar[1][1]) / tuple_scalar[0]
+
+    @field_operator
+    def scan_tuple_scalar(
+        qc: Field[[IDim, KDim], float], tuple_scalar: tuple[float, tuple[float, float]]
+    ) -> Field[[IDim, KDim], float]:
+        return _scan_tuple_scalar(qc, tuple_scalar)
+
+    scan_tuple_scalar(qc, tuple_scalar, out=qc, offset_provider={})
+    assert np.allclose(np.asarray(qc), expected)
+
+
 def test_astype_int(fieldview_backend):
     size = 10
     b_float_64 = np_as_located_field(IDim)(np.ones((size), dtype=np.float64))
@@ -744,8 +788,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 +842,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 +853,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)
     )