Flow nonlocal writes back to the enclosing function

docs/usage.md

@@ -366,8 +366,10 @@ boolean operators, chained comparisons, keyword and starred arguments, `with`, `
 comprehensions, lambdas, nested functions and classes, `global`, `nonlocal`, `del`,
 assignment expressions (`(y := f(x))`, laid out as an assignment before the statement),
 `yield from`, starred assignment targets (`first, *rest = items`) and `match` with every
-pattern kind. `finally` is modelled on the normal path only and a
-`nonlocal` write stays inside the nested function. Methods of module-level classes are
+pattern kind. `finally` is modelled on the normal path only. A
+`nonlocal` write made by a nested function flows back to the enclosing variable when the
+enclosing function calls that nested function directly by name; a nested function passed
+around and called elsewhere keeps its writes to itself. Methods of module-level classes are
 analysed like functions.
 
 ## Looking at the intermediate representation

src/coretrace_python/cache.py

@@ -29,14 +29,15 @@
     KnownFunction,
     ModuleGraph,
     Mutation,
+    NonlocalWrite,
     SummaryIndex,
     Target,
     UnknownTarget,
 )
 from coretrace_python.semantic.symbols import SymbolId
 from coretrace_python.source import SourceId, SourceSpan
 
-CACHE_FORMAT = 2
+CACHE_FORMAT = 3
 
 
 @dataclass(frozen=True)
@@ -247,6 +248,10 @@ def _encode_summary(summary: FunctionSummary) -> dict[str, Any]:
             for m in summary.mutations
         ],
         "side_effects": sorted(summary.side_effects),
+        "nonlocal_writes": [
+            {"name": w.name, "dependencies": sorted(w.dependencies), "externals": sorted(str(s) for s in w.externals)}
+            for w in summary.nonlocal_writes
+        ],
     }
 
 
@@ -268,6 +273,12 @@ def _decode_summary(data: Mapping[str, Any]) -> FunctionSummary:
             for m in data["mutations"]
         ),
         frozenset(_string(name) for name in data["side_effects"]),
+        tuple(
+            NonlocalWrite(
+                _string(w["name"]), _indices(w["dependencies"]), frozenset(SymbolId(_string(s)) for s in w["externals"])
+            )
+            for w in data["nonlocal_writes"]
+        ),
     )
 
 

src/coretrace_python/interprocedural/__init__.py

@@ -20,6 +20,7 @@
     ExternalCall,
     FunctionSummary,
     Mutation,
+    NonlocalWrite,
     ProjectSummaries,
     SummaryAnalysis,
     SummaryIndex,
@@ -36,6 +37,7 @@
     "KnownFunction",
     "ModuleGraph",
     "Mutation",
+    "NonlocalWrite",
     "ProjectSummaries",
     "SummaryAnalysis",
     "SummaryIndex",

src/coretrace_python/interprocedural/summaries.py

@@ -50,10 +50,12 @@
     Instruction,
     Jump,
     MakeFunction,
+    NonlocalResult,
     Phi,
     Return,
     SetAttr,
     SetItem,
+    SetNonlocal,
     Symbol,
     Undefined,
     Value,
@@ -102,6 +104,16 @@ class Mutation:
     externals: frozenset[SymbolId]
 
 
+@dataclass(frozen=True)
+class NonlocalWrite:
+    """What the function stores into its ``nonlocal`` ``name``: data of its parameters
+    and results of external symbols, written back to the caller's variable."""
+
+    name: str
+    dependencies: Dependencies
+    externals: frozenset[SymbolId]
+
+
 @dataclass(frozen=True)
 class FunctionSummary:
     name: str
@@ -112,6 +124,7 @@ class FunctionSummary:
     return_externals: frozenset[SymbolId] = frozenset()
     mutations: tuple[Mutation, ...] = ()
     side_effects: frozenset[str] = frozenset()
+    nonlocal_writes: tuple[NonlocalWrite, ...] = ()
 
 
 class SummaryIndex:
@@ -184,6 +197,7 @@ def __init__(
         self.returns: Dep = EMPTY
         self.stores: dict[HeapLocation, Dep] = {}
         self.mutated_globals: set[str] = set()
+        self.nonlocal_stores: dict[str, Dep] = {}
 
     # ------------------------------------------------------------------ heap
 
@@ -224,6 +238,9 @@ def evaluate(self, block: BasicBlock, incoming: Mapping[BlockId, State]) -> Stat
                 self.store(state, instruction.object, ATTRIBUTES, state.get(instruction.value, EMPTY))
             elif isinstance(instruction, SetItem):
                 self.store(state, instruction.object, ELEMENTS, state.get(instruction.value, EMPTY))
+            elif isinstance(instruction, SetNonlocal):
+                previous = self.nonlocal_stores.get(instruction.name, EMPTY)
+                self.nonlocal_stores[instruction.name] = previous | state.get(instruction.value, EMPTY)
             if instruction.result is None:
                 continue
             if isinstance(instruction, Phi):
@@ -233,6 +250,8 @@ def evaluate(self, block: BasicBlock, incoming: Mapping[BlockId, State]) -> Stat
                         deps |= incoming[predecessor].get(value, EMPTY)
             elif isinstance(instruction, Call):
                 deps = self.call(instruction, state)
+            elif isinstance(instruction, NonlocalResult):
+                deps = self.nonlocal_result(instruction, state)
             else:
                 deps = self.instruction(instruction, state)
                 if isinstance(instruction, GetAttr | GetItem | GetIter):
@@ -324,6 +343,34 @@ def call(self, call: Call, state: dict[Key, Dep]) -> Dep:
         assert isinstance(target, UnknownTarget)
         return everything | state.get(call.callee, EMPTY)
 
+    def nonlocal_result(self, instruction: NonlocalResult, state: dict[Key, Dep]) -> Dep:
+        """What the nested callee left in its nonlocal ``name``, in the caller's terms."""
+
+        call = self.defs.get(instruction.call)
+        if not isinstance(call, Call):
+            return EMPTY
+        target = self.graph.target_at(self.name, call.location)
+        if not isinstance(target, KnownFunction) or target.name not in self.table:
+            return EMPTY
+        callee = self.table[target.name]
+        arguments = tuple(self.deep(a, state) for a in call.arguments)
+        made = self.defs.get(call.callee)
+        if isinstance(made, MakeFunction):
+            arguments = (*arguments, *(self.deep(v, state) for v in made.captured))
+        bound = self.receiver(call, target.name)
+        arguments = (*(self.deep(r, state) for r in bound), *arguments)
+        keywords = EMPTY
+        for value in (*call.starred, *(v for _, v in call.keywords)):
+            keywords |= self.deep(value, state)
+        deps = EMPTY
+        for write in callee.nonlocal_writes:
+            if write.name != instruction.name:
+                continue
+            for index in write.dependencies:
+                deps |= arguments[index] if index < len(arguments) else keywords
+            deps |= Dep(externals=write.externals)
+        return deps
+
     def receiver(self, call: Call, name: str) -> tuple[Value, ...]:
         """The object a method call runs on: the receiver of ``obj.method(...)`` or the
         new object of ``Class(...)``, the method's implicit first parameter."""
@@ -407,6 +454,7 @@ def summarize(
     problem.returns = EMPTY
     problem.stores = {}
     problem.mutated_globals = set()
+    problem.nonlocal_stores = {}
     for block in function.blocks:
         if solution.reached(block.id):
             problem.evaluate(block, solution.incoming(block.id))
@@ -426,6 +474,11 @@ def summarize(
         return_externals=problem.returns.externals,
         mutations=tuple(mutations),
         side_effects=frozenset(problem.mutated_globals),
+        nonlocal_writes=tuple(
+            NonlocalWrite(name, deps.parameters, deps.externals)
+            for name, deps in sorted(problem.nonlocal_stores.items())
+            if deps.parameters or deps.externals
+        ),
     )
 
 

src/coretrace_python/ir/lowering.py

@@ -50,11 +50,13 @@
     MakeClass,
     MakeFunction,
     ModuleIR,
+    NonlocalResult,
     Raise,
     Return,
     SetAttr,
     SetGlobal,
     SetItem,
+    SetNonlocal,
     StoreLocal,
     Symbol,
     Terminator,
@@ -166,7 +168,9 @@ def expression(self, node: nodes.Expression) -> Value:
             callee = self.expression(node.callee)
             arguments, starred = self.spread(node.arguments)
             keywords = tuple((k.name, self.expression(k.value)) for k in node.keywords)
-            return self.emit(Call(self.new_value(), node.span, callee, arguments, keywords, starred))
+            result = self.emit(Call(self.new_value(), node.span, callee, arguments, keywords, starred))
+            self.write_back_nonlocals(node, result)
+            return result
         if isinstance(node, nodes.BoolOp):
             values = tuple(self.expression(value) for value in node.values)
             return self.emit(BoolOp(self.new_value(), node.span, node.operator, values))
@@ -221,6 +225,20 @@ def spread(self, elements: tuple[nodes.Expression, ...]) -> tuple[tuple[Value, .
                 plain.append(self.expression(element))
         return tuple(plain), tuple(unpacked)
 
+    def write_back_nonlocals(self, node: nodes.Call, call: Value) -> None:
+        """After a direct call to a nested function defined here, store what it left in
+        each of its ``nonlocal`` names back into this function's variable."""
+
+        if not isinstance(node.callee, nodes.Name):
+            return
+        nested = self.nested_defs.get(node.callee.identifier)
+        if nested is None:
+            return
+        bindings = self.scopes.scope_for(nested).bindings
+        for name in sorted(n for n, b in bindings.items() if b.kind is BindingKind.NONLOCAL):
+            written = self.emit(NonlocalResult(self.new_value(), node.span, call, name))
+            self.store(nodes.Name(name, node.span), written)
+
     # ------------------------------------------------------------------ statements
 
     def store(self, target: nodes.Target, value: Value) -> None:
@@ -232,6 +250,8 @@ def store(self, target: nodes.Target, value: Value) -> None:
             if resolution.kind is not ResolutionKind.LOCAL and target.identifier not in self.nonlocals:
                 self.fail(target, "assignment to a free variable that is not declared nonlocal")
             self.emit_effect(StoreLocal(None, target.span, target.identifier, value))
+            if target.identifier in self.nonlocals:
+                self.emit_effect(SetNonlocal(None, target.span, target.identifier, value))
         elif isinstance(target, nodes.Attribute):
             object_value = self.expression(target.value)
             self.emit_effect(SetAttr(None, target.span, object_value, target.name, value))
@@ -300,6 +320,7 @@ def statement(self, node: nodes.Statement) -> None:
             captured = tuple(self.captured_value(name, node.span) for name in captured_names(node, self.scopes))
             made = self.emit(MakeFunction(self.new_value(), node.span, node.name, captured))
             self.store(nodes.Name(node.name, node.span), made)
+            self.nested_defs[node.name] = node
             return
         if isinstance(node, nodes.AugAssign):
             current = self.expression(node.target)
@@ -405,6 +426,9 @@ def function(self, node: nodes.Function) -> FunctionIR:
         # assigned in this body, so it lives in a local slot initialised from the
         # captured value; the others are never reassigned.
         bindings = self.scopes.scope_for(node).bindings
+        # Nested definitions bound in this body, so a direct call to one can write its
+        # ``nonlocal`` names back into this function's variables.
+        self.nested_defs: dict[str, nodes.Function] = {}
         self.nonlocals = frozenset(
             name for name in captured if name in bindings and bindings[name].kind is BindingKind.NONLOCAL
         )
@@ -429,7 +453,11 @@ def function(self, node: nodes.Function) -> FunctionIR:
                 )
             )
         return FunctionIR(
-            qualified_name(self.scopes, node), parameter_values, self.cfg.entry, tuple(blocks), node.span
+            qualified_name(self.scopes, node),
+            parameter_values,
+            self.cfg.entry,
+            tuple(blocks),
+            node.span,
         )
 
 

src/coretrace_python/ir/model.py

@@ -234,6 +234,17 @@ class Yield(Operands):
     value: Value | None
 
 
+@dataclass(frozen=True)
+class NonlocalResult(Operands):
+    """The value a nested function left in its ``nonlocal`` ``name`` when ``call``
+    returned; the caller stores it back into its own variable."""
+
+    result: Value
+    location: SourceSpan
+    call: Value
+    name: str
+
+
 @dataclass(frozen=True)
 class GetAttr(Operands):
     result: Value
@@ -340,6 +351,17 @@ class DelAttr(Operands):
     attribute: str
 
 
+@dataclass(frozen=True)
+class SetNonlocal(Operands):
+    """A store to a name this function declares ``nonlocal``, kept through SSA so the
+    summary can tell the enclosing function what was written back."""
+
+    result: None
+    location: SourceSpan
+    name: str
+    value: Value
+
+
 @dataclass(frozen=True)
 class SetGlobal(Operands):
     """Assignment to a name declared ``global``."""
@@ -389,9 +411,11 @@ class Import(Operands):
     | Catch
     | Await
     | Yield
+    | NonlocalResult
 )
 EffectInstruction: TypeAlias = (
     StoreLocal | SetAttr | SetItem | WithExit | Assert | Import | SetGlobal | DelItem | DelAttr
+    | SetNonlocal
 )
 Instruction: TypeAlias = ValueInstruction | EffectInstruction
 

src/coretrace_python/ir/printer.py

@@ -30,12 +30,14 @@
     MakeClass,
     MakeFunction,
     ModuleIR,
+    NonlocalResult,
     Phi,
     Raise,
     Return,
     SetAttr,
     SetGlobal,
     SetItem,
+    SetNonlocal,
     StoreLocal,
     Symbol,
     Terminator,
@@ -107,6 +109,8 @@ def _instruction(instruction: Instruction) -> str:
         return f'{_value(instruction.result)} = make_class "{instruction.name}"'
     if isinstance(instruction, SetGlobal):
         return f'set_global "{instruction.name}", {_value(instruction.value)}'
+    if isinstance(instruction, SetNonlocal):
+        return f'set_nonlocal "{instruction.name}", {_value(instruction.value)}'
     if isinstance(instruction, BuildString):
         return f"{_value(instruction.result)} = build_string {', '.join(_value(v) for v in instruction.parts)}".rstrip()
     if isinstance(instruction, BuildSlice):
@@ -124,6 +128,8 @@ def _instruction(instruction: Instruction) -> str:
     if isinstance(instruction, Yield):
         suffix = "" if instruction.value is None else f" {_value(instruction.value)}"
         return f"{_value(instruction.result)} = yield{suffix}"
+    if isinstance(instruction, NonlocalResult):
+        return f"{_value(instruction.result)} = nonlocal_result {_value(instruction.call)}, {instruction.name!r}"
     if isinstance(instruction, SetAttr):
         return f"set_attr {_value(instruction.object)}, {instruction.attribute!r}, {_value(instruction.value)}"
     if isinstance(instruction, SetItem):