fix[next]: collect closure vars via the compiler's scope analysis

src/gt4py/next/ffront/source_utils.py

@@ -8,11 +8,13 @@
 
 from __future__ import annotations
 
+import builtins
 import functools
 import inspect
 import pathlib
 import symtable
 import textwrap
+import types
 from collections.abc import Callable, Iterator
 from dataclasses import dataclass
 from typing import Any, cast
@@ -21,14 +23,55 @@
 MISSING_FILENAME = "<string>"
 
 
+def _global_names_from_source(source: str) -> set[str]:
+    """
+    Names referenced as globals in the function defined by `source`, including its nested scopes.
+
+    The compiler's own scope analysis decides what is a global, so locals shadowing a
+    global name and comprehension targets are never reported.
+    """
+
+    def walk(table: symtable.SymbolTable) -> Iterator[str]:
+        if isinstance(table, symtable.Function):
+            yield from table.get_globals()
+        for child in table.get_children():
+            yield from walk(child)
+
+    # Analyzed as if under PEP 563 so that annotations contribute no names: Python evaluates
+    # parameter and return annotations in the enclosing scope and never evaluates local
+    # variable annotations, but only Python 3.14's symtable stops reporting the latter.
+    source = "from __future__ import annotations\n" + source
+    return set(walk(symtable.symtable(source, MISSING_FILENAME, "exec")))
+
+
+@functools.cache
+def _global_names_of_code(code: types.CodeType) -> frozenset[str]:
+    return frozenset(_global_names_from_source(make_source_definition_from_function(code).source))
+
+
 def get_closure_vars_from_function(function: Callable) -> dict[str, Any]:
-    (nonlocals, globals, builtins, _unbound) = inspect.getclosurevars(function)  # noqa: A001 [builtin-variable-shadowing]
+    # `inspect.getclosurevars` only sees the names of the function's own code object, which
+    # misses names referenced only inside a nested scope such as a generator expression.
+    # Free variables are unaffected (they are cells of the function itself), so only the
+    # global names are taken from the source instead.
+    nonlocals = inspect.getclosurevars(function).nonlocals
+    global_ns = function.__globals__
+    builtin_ns = global_ns.get("__builtins__", builtins.__dict__)
+    if inspect.ismodule(builtin_ns):
+        builtin_ns = builtin_ns.__dict__
+
+    closure_vars: dict[str, Any] = {}
+    for name in _global_names_of_code(function.__code__):
+        if name in global_ns:
+            closure_vars[name] = global_ns[name]
+        elif name in builtin_ns:
+            closure_vars[name] = builtin_ns[name]
 
     # nonlocals override globals, sorted for deterministic results
-    return dict(sorted({**builtins, **globals, **nonlocals}.items()))
+    return dict(sorted({**closure_vars, **nonlocals}.items()))
 
 
-def make_source_definition_from_function(func: Callable) -> SourceDefinition:
+def make_source_definition_from_function(func: Callable | types.CodeType) -> SourceDefinition:
     try:
         filename = str(pathlib.Path(inspect.getabsfile(func)).resolve())
         if not filename:

tests/next_tests/unit_tests/ffront_tests/test_source_utils.py

@@ -0,0 +1,110 @@
+# GT4Py - GridTools Framework
+#
+# Copyright (c) 2014-2024, ETH Zurich
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Closure variable collection across nested scopes.
+
+A generator expression, lambda or nested function compiles to its own code object,
+so names referenced only inside it appear in that object's `co_names` and never in
+the enclosing function's. Collecting closure variables from the enclosing function
+alone therefore misses them; `get_closure_vars_from_function` takes the global names
+from the compiler's symbol table of the source instead, which covers nested scopes.
+
+Free variables need no such treatment: the enclosing code object carries a cell for
+them.
+"""
+
+import gt4py.next as gtx
+from gt4py.next import Dims, Dimension, float64, neighbor_sum
+from gt4py.next.ffront import source_utils
+from gt4py.next.ffront.source_utils import get_closure_vars_from_function
+
+
+Cell = Dimension("Cell")
+Edge = Dimension("Edge")
+C2EDim = Dimension("C2E", kind=gtx.DimensionKind.LOCAL)
+C2E = gtx.FieldOffset("C2E", source=Edge, target=(Cell, C2EDim))
+
+CField = gtx.Field[Dims[Cell], float64]
+EField = gtx.Field[Dims[Edge], float64]
+
+
+@gtx.field_operator
+def scale(f: CField, factor: float64) -> CField:
+    return f * factor
+
+
+def _builtin_and_offset(tracers: tuple[EField, ...]) -> tuple[CField, ...]:
+    return tuple(neighbor_sum(t(C2E), axis=C2EDim) for t in tracers)
+
+
+def _module_level_operator(tracers: tuple[CField, ...], factor: float64) -> tuple[CField, ...]:
+    return tuple(scale(t, factor) for t in tracers)
+
+
+def test_names_are_collected_from_the_nested_code_object():
+    """The names live in the generator expression's code object, not the function's."""
+    assert "neighbor_sum" not in _builtin_and_offset.__code__.co_names
+    nested = [c for c in _builtin_and_offset.__code__.co_consts if hasattr(c, "co_names")]
+    assert any("neighbor_sum" in c.co_names for c in nested)
+
+    collected = get_closure_vars_from_function(_builtin_and_offset)
+    assert {"neighbor_sum", "C2E", "C2EDim"} <= set(collected)
+
+
+def test_comprehension_target_is_not_collected():
+    """The loop target is a local of the nested code object, not a global reference."""
+    assert "t" not in get_closure_vars_from_function(_module_level_operator)
+
+
+def test_free_variables_still_resolve():
+    """The path that already worked, kept as a guard."""
+
+    @gtx.field_operator
+    def local_scale(f: CField, factor: float64) -> CField:
+        return f * factor
+
+    def uses_freevar(tracers: tuple[CField, ...], factor: float64) -> tuple[CField, ...]:
+        return tuple(local_scale(t, factor) for t in tracers)
+
+    assert get_closure_vars_from_function(uses_freevar)["local_scale"] is local_scale
+
+
+def test_local_name_shadowing_a_global_is_not_collected_as_global():
+    """A comprehension referencing an enclosing local must bind the local, not the global."""
+
+    def shadows(tracers: tuple[CField, ...], factor: float64) -> tuple[CField, ...]:
+        scale = local_helper  # noqa: F841  shadows the module-level 'scale'
+        return tuple(scale(t, factor) for t in tracers)
+
+    def local_helper(t, factor):
+        return t
+
+    assert "scale" not in get_closure_vars_from_function(shadows)
+
+
+def test_source_is_analyzed_once_per_code_object(monkeypatch):
+    calls = []
+    read_source = source_utils.make_source_definition_from_function
+    monkeypatch.setattr(
+        source_utils,
+        "make_source_definition_from_function",
+        lambda f: calls.append(f) or read_source(f),
+    )
+    source_utils._global_names_of_code.cache_clear()
+
+    get_closure_vars_from_function(_module_level_operator)
+    get_closure_vars_from_function(_module_level_operator)
+
+    assert calls == [_module_level_operator.__code__]
+
+
+def test_values_are_looked_up_on_every_call(monkeypatch):
+    """Only the names are cached, so rebinding a global is seen by the next collection."""
+    assert get_closure_vars_from_function(_module_level_operator)["scale"] is scale
+    monkeypatch.setitem(globals(), "scale", 42)
+    assert get_closure_vars_from_function(_module_level_operator)["scale"] == 42