feat: closure bodies

README.md

@@ -159,9 +159,10 @@ assignment, list, tuple and dict literals with `*` and `**` unpacking, f-strings
 `try`/`except`/`else`/`finally`, `await`, `yield`, `if`/`elif`/`else`, `while`, `for`
 with name or tuple targets, `break`, `continue` and `raise`, `from` clause included,
 conditional expressions and comprehensions, laid out as real branches and loops over a
-synthetic local, set literals, chained assignments, assignments to `global` names, and
-lambdas and nested function definitions as function values whose bodies are not analysed
-yet, `del`, loop `else` clauses, and `match` over literal, singleton, capture, wildcard
+synthetic local, set literals, chained assignments, assignments to `global` names,
+lambdas and nested function definitions, analysed as functions of their own whose
+captured variables are implicit parameters so taint flows through closures like through
+any call, `del`, loop `else` clauses, and `match` over literal, singleton, capture, wildcard
 and or-patterns with guards, laid out as an `if` chain. Not yet: sequence, mapping and
 class patterns, `nonlocal` assignments, classes defined inside functions, and a
 conditional expression or comprehension inside a `while` condition. Blocks inside a `try`

src/coretrace_python/interprocedural/callgraph.py

@@ -23,6 +23,7 @@
     GetAttr,
     GetItem,
     Global,
+    MakeFunction,
     Symbol,
     Value,
     WithEnter,
@@ -149,8 +150,10 @@ def resolve_targets(
     scopes: ScopeTable,
     known: frozenset[str],
     parameters: Mapping[Value, SymbolId] | None = None,
+    nested: frozenset[str] = frozenset(),
 ) -> tuple[dict[Value, Target], dict[Value, SymbolId]]:
-    """Map every callee value of ``function`` to its target, and every symbol value."""
+    """Map every callee value of ``function`` to its target, and every symbol value.
+    ``nested`` names the functions defined inside this one, ``outer.inner``."""
 
     module = scopes.module_scope
     symbols = derive_symbols(function, parameters)
@@ -159,6 +162,10 @@ def resolve_targets(
     }
     for block in function.blocks:
         for instruction in block.instructions:
+            if isinstance(instruction, MakeFunction):
+                qualified = f"{function.name}.{instruction.name}"
+                if qualified in nested:
+                    targets[instruction.result] = KnownFunction(qualified)
             if isinstance(instruction, Global):
                 binding = module.bindings.get(instruction.name)
                 if (
@@ -178,7 +185,8 @@ def _annotated(
     scope = scopes.scope_for(function)
     enclosing = scope.parent if scope.parent is not None else scope.id
     found: dict[Value, SymbolId] = {}
-    for value, parameter in zip(ssa.parameters, function.parameters, strict=True):
+    # Captured variables follow the explicit parameters and carry no annotation.
+    for value, parameter in zip(ssa.parameters, function.parameters, strict=False):
         if parameter.annotation is None:
             continue
         symbol = table.resolve_expression(enclosing, parameter.annotation)
@@ -218,7 +226,9 @@ def compute(cls, ctx: AnalysisContext) -> CallGraph:
                 unsupported.add(name)
                 sites[name] = ()
                 continue
-            targets, symbols[name] = resolve_targets(ssa, scopes, known, _annotated(function, ssa, scopes, table))
+            targets, symbols[name] = resolve_targets(
+                ssa, scopes, known, _annotated(function, ssa, scopes, table), frozenset(definitions)
+            )
             found: list[CallSite] = []
             for block in ssa.blocks:
                 for instruction in block.instructions:

src/coretrace_python/interprocedural/summaries.py

@@ -49,6 +49,7 @@
     Global,
     Instruction,
     Jump,
+    MakeFunction,
     Phi,
     Return,
     SetAttr,
@@ -310,6 +311,9 @@ def call(self, call: Call, state: dict[Key, Dep]) -> Dep:
             )
             return everything | Dep(externals=frozenset({target.symbol}))
         if isinstance(target, KnownFunction):
+            made = self.defs.get(call.callee)
+            if isinstance(made, MakeFunction):
+                arguments = (*arguments, *(self.deep(v, state) for v in made.captured))
             return self.known(self.table[target.name], arguments, keywords, everything, call, state)
         assert isinstance(target, UnknownTarget)
         return everything | state.get(call.callee, EMPTY)
@@ -342,10 +346,12 @@ def mapped(indices: Dependencies) -> Dep:
                 reached.location,
                 call.location,
             )
+        made = self.defs.get(call.callee)
+        values = (*call.arguments, *(made.captured if isinstance(made, MakeFunction) else ()))
         for mutation in callee.mutations:
-            if mutation.parameter < len(call.arguments):
+            if mutation.parameter < len(values):
                 deps = mapped(mutation.dependencies) | Dep(externals=mutation.externals)
-                self.store(state, call.arguments[mutation.parameter], mutation.field, deps)
+                self.store(state, values[mutation.parameter], mutation.field, deps)
         return mapped(callee.return_dependencies) | Dep(externals=callee.return_externals)
 
     def record(

src/coretrace_python/ir/lowering.py

@@ -134,6 +134,9 @@ def expression(self, node: nodes.Expression) -> Value:
         if isinstance(node, nodes.Name):
             resolution = self.resolve(node.identifier)
             if resolution.kind is ResolutionKind.FREE:
+                # A captured variable is an implicit parameter of the nested function.
+                if node.identifier in self.parameters:
+                    return self.parameters[node.identifier]
                 self.fail(node, "closures are not supported yet")
             if resolution.kind is not ResolutionKind.LOCAL:
                 return self.emit(Global(self.new_value(), node.span, node.identifier))
@@ -181,7 +184,9 @@ def expression(self, node: nodes.Expression) -> Value:
             elements, unpacked = self.spread(node.elements)
             return self.emit(BuildSet(self.new_value(), node.span, elements, unpacked))
         if isinstance(node, nodes.Lambda):
-            return self.emit(MakeFunction(self.new_value(), node.span, "<lambda>"))
+            synthesized = lambda_function(node)
+            captured = tuple(self.captured_value(name, node.span) for name in captured_names(synthesized, self.scopes))
+            return self.emit(MakeFunction(self.new_value(), node.span, synthesized.name, captured))
         if isinstance(node, nodes.Conditional | nodes.Comprehension):
             self.fail(node, "expression-level control flow must be laid out by the CFG builder")
         if isinstance(node, nodes.Await):
@@ -256,7 +261,8 @@ def statement(self, node: nodes.Statement) -> None:
             for parameter in node.parameters:
                 if parameter.default is not None:
                     self.expression(parameter.default)
-            made = self.emit(MakeFunction(self.new_value(), node.span, node.name))
+            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)
             return
         if isinstance(node, nodes.AugAssign):
@@ -345,20 +351,29 @@ def enter_loop(self, source: BlockId, target: BlockId) -> None:
             iterable = self.expression(header.iterable)
             self.iterators[target] = self.emit(GetIter(self.new_value(), header.span, iterable))
 
+    def captured_value(self, name: str, span: SourceSpan) -> Value:
+        """The current value of a variable a nested function captures."""
+
+        return self.expression(nodes.Name(name, span))
+
     def function(self, node: nodes.Function) -> FunctionIR:
-        parameter_values = tuple(self.new_value() for _ in node.parameters)
+        captured = captured_names(node, self.scopes)
+        parameter_values = tuple(self.new_value() for _ in (*node.parameters, *captured))
         reassigned = _reassigned_parameters(node)
         self.parameters = {
             parameter.name: value
-            for parameter, value in zip(node.parameters, parameter_values, strict=True)
+            for parameter, value in zip(node.parameters, parameter_values, strict=False)
             if parameter.name not in reassigned
         }
+        # Captured variables come after the explicit parameters; they are never
+        # reassigned, an assignment would have made them local.
+        self.parameters.update(zip(captured, parameter_values[len(node.parameters) :], strict=True))
         blocks: list[BasicBlock] = []
         for cfg_block in self.cfg.blocks.values():
             self.instructions = []
             if cfg_block.id == self.cfg.entry:
                 # Reassigned parameters live in locals so every block reads the same slot.
-                for parameter, value in zip(node.parameters, parameter_values, strict=True):
+                for parameter, value in zip(node.parameters, parameter_values, strict=False):
                     if parameter.name in reassigned:
                         self.emit_effect(StoreLocal(None, parameter.span, parameter.name, value))
             for statement in cfg_block.statements:
@@ -375,17 +390,51 @@ def function(self, node: nodes.Function) -> FunctionIR:
 
 
 def qualified_name(scopes: ScopeTable, function: nodes.Function) -> str:
-    """``Class.method`` for methods, the bare name for module-level functions."""
+    """``Class.method`` for methods, ``outer.inner`` for nested functions, the bare name
+    for module-level functions; lambda and comprehension scopes keep identifier names."""
 
     names = [function.name]
     scope = scopes.scope_for(function)
     parent = scopes.scope(scope.parent) if scope.parent else None
     while parent is not None and parent.kind is not ScopeKind.MODULE:
-        names.append(parent.name)
+        names.append(parent.name.strip("<>"))
         parent = scopes.scope(parent.parent) if parent.parent else None
     return ".".join(reversed(names))
 
 
+def lambda_function(node: nodes.Lambda) -> nodes.Function:
+    """A lambda as a function named after its position, returning its body."""
+
+    return nodes.Function(
+        f"lambda_{node.span.start_line}_{node.span.start_column}",
+        node.parameters,
+        (nodes.Return(node.body, node.span),),
+        False,
+        node.span,
+    )
+
+
+def captured_names(function: nodes.Function, scopes: ScopeTable) -> tuple[str, ...]:
+    """The enclosing-function variables ``function`` reads, sorted: its implicit
+    parameters after the explicit ones."""
+
+    scope = scopes.scope_for(function)
+    names: set[str] = set()
+
+    def walk(node: Node) -> None:
+        if isinstance(node, nodes.Name):
+            names.add(node.identifier)
+        for child in children(node):
+            walk(child)
+
+    for statement in function.body:
+        walk(statement)
+    for parameter in function.parameters:
+        if parameter.default is not None:
+            walk(parameter.default)
+    return tuple(sorted(n for n in names if scopes.resolve(scope.id, n).kind is ResolutionKind.FREE))
+
+
 def _reassigned_parameters(function: nodes.Function) -> frozenset[str]:
     """Parameters that the function body assigns to, outside nested scopes."""
 
@@ -434,17 +483,39 @@ def compute(cls, ctx: AnalysisContext, function: nodes.Function) -> FunctionIR:
 
 
 def analyzable_functions(module: nodes.Module) -> tuple[nodes.Function, ...]:
-    """Top-level functions and the methods of top-level classes, in source order."""
+    """Top-level functions, the methods of top-level classes, and the functions and
+    lambdas nested inside them, in source order."""
 
     functions: list[nodes.Function] = []
     for statement in module.body:
         if isinstance(statement, nodes.Function):
-            functions.append(statement)
+            _collect(statement, functions)
         elif isinstance(statement, nodes.Class):
-            functions.extend(s for s in statement.body if isinstance(s, nodes.Function))
+            for member in statement.body:
+                if isinstance(member, nodes.Function):
+                    _collect(member, functions)
     return tuple(functions)
 
 
+def _collect(function: nodes.Function, into: list[nodes.Function]) -> None:
+    into.append(function)
+
+    def walk(node: Node) -> None:
+        if isinstance(node, nodes.Function):
+            _collect(node, into)
+            return
+        if isinstance(node, nodes.Lambda):
+            _collect(lambda_function(node), into)
+            return
+        if isinstance(node, nodes.Class):
+            return
+        for child in children(node):
+            walk(child)
+
+    for statement in function.body:
+        walk(statement)
+
+
 class ModuleIRAnalysis(Analysis[ModuleIR]):
     """Assemble the PyIR of every top-level function of the module."""
 

src/coretrace_python/ir/model.py

@@ -166,11 +166,14 @@ class BuildSet(Operands):
 
 @dataclass(frozen=True)
 class MakeFunction(Operands):
-    """A function value: a lambda or a nested definition, whose body is not lowered."""
+    """A function value: a lambda or a nested definition. ``captured`` are the values of
+    the variables the body captures, at the definition point, in the order the nested
+    function takes them as implicit parameters after its explicit ones."""
 
     result: Value
     location: SourceSpan
     name: str
+    captured: tuple[Value, ...] = ()
 
 
 @dataclass(frozen=True)

src/coretrace_python/ir/printer.py

@@ -94,7 +94,10 @@ def _instruction(instruction: Instruction) -> str:
         parts = [_value(v) for v in instruction.elements] + [f"*{_value(v)}" for v in instruction.unpacked]
         return f"{_value(instruction.result)} = build_set {', '.join(parts)}".rstrip()
     if isinstance(instruction, MakeFunction):
-        return f'{_value(instruction.result)} = make_function "{instruction.name}"'
+        text = f'{_value(instruction.result)} = make_function "{instruction.name}"'
+        if instruction.captured:
+            text += f" [{', '.join(_value(v) for v in instruction.captured)}]"
+        return text
     if isinstance(instruction, DelItem):
         return f"del_item {_value(instruction.object)}, {_value(instruction.key)}"
     if isinstance(instruction, DelAttr):