feat: control-flow syntax pass

README.md

@@ -140,9 +140,12 @@ forms, decorators, assignments to names, attributes, items and unpacked tuples,
 assignment, list, tuple and dict literals with `*` and `**` unpacking, f-strings, slices,
 `and`/`or`, chained comparisons, keyword and starred arguments, `with`, `assert`,
 `try`/`except`/`else`/`finally`, `await`, `yield`, `if`/`elif`/`else`, `while`, `for`
-with name or tuple targets, `break`, `continue` and `raise`, `from` clause included. Not
-yet: conditional expressions, lambdas, comprehensions, nested function definitions and
-assignments to `global` names. Blocks inside a `try`
+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. Not yet: `match`, `del`, `nonlocal` assignments, classes defined inside functions, and
+a conditional expression or comprehension inside a `while` condition. Blocks inside a `try`
 body carry exception edges to the handlers; `finally` is modelled on the normal path only. Methods of module-level classes are analysed like functions; other module-level
 code is skipped. An import inside a function is shown where it runs, as an `import`
 instruction naming the module, the bound name and the canonical symbol. A function using syntax outside this subset is reported by `--check` as

src/coretrace_python/abstract/heap.py

@@ -23,6 +23,7 @@
     Await,
     BuildDict,
     BuildList,
+    BuildSet,
     BuildTuple,
     Call,
     Catch,
@@ -183,23 +184,24 @@ def instruction(self, instruction: Instruction) -> bool:
         result = instruction.result
         if result is None:
             return False
-        if isinstance(instruction, BuildList | BuildTuple | BuildDict | Call | WithEnter | Catch | Yield):
+        if isinstance(instruction, BuildList | BuildTuple | BuildDict | BuildSet | Call | WithEnter | Catch | Yield):
             kind = {
                 BuildList: "list",
                 BuildTuple: "tuple",
                 BuildDict: "dict",
+                BuildSet: "set",
                 Call: "call",
                 WithEnter: "context",
                 Catch: "exception",
                 Yield: "sent",
             }[type(instruction)]
             site = AbstractObject(AllocationSite(kind, instruction.location))
             changed = self.assign(result, {site})
-            if isinstance(instruction, BuildList | BuildTuple | BuildDict):
+            if isinstance(instruction, BuildList | BuildTuple | BuildDict | BuildSet):
                 values = (
-                    instruction.elements
-                    if isinstance(instruction, BuildList | BuildTuple)
-                    else tuple(v for _, v in instruction.items)
+                    tuple(v for _, v in instruction.items)
+                    if isinstance(instruction, BuildDict)
+                    else instruction.elements
                 )
                 for element in values:
                     changed |= self.store({site}, ELEMENTS, element)

src/coretrace_python/cfg/builder.py

@@ -2,7 +2,8 @@
 
 from __future__ import annotations
 
-from typing import ClassVar
+from dataclasses import fields, is_dataclass, replace
+from typing import Any, ClassVar
 
 from coretrace_python.analysis import AnalysisContext, FunctionAnalysis
 from coretrace_python.cfg.model import (
@@ -36,13 +37,137 @@ def __init__(self, function: nodes.Function) -> None:
         self.counters: dict[str, int] = {}
         self.loops: list[tuple[BlockId, BlockId]] = []
         self.handlers: list[tuple[BlockId, ...]] = []
+        self.synthetic: set[str] = set()
 
     def build(self) -> CFG:
         entry = BlockId("entry")
         end = self.sequence(self.function.body, _Open(entry), None, self.function.span)
         if end is not None:
             self.finish(end, Return(None, self.function.span))
-        return CFG(entry, self.blocks)
+        return CFG(entry, self.blocks, frozenset(self.synthetic))
+
+    # ------------------------------------------------------------------ expression-level control flow
+
+    def hidden(self, kind: str, span: SourceSpan) -> nodes.Name:
+        name = f"_coretrace_{kind}_{span.start_line}_{span.start_column}_{len(self.synthetic)}"
+        self.synthetic.add(name)
+        return nodes.Name(name, span)
+
+    def desugared(self, node: nodes.Statement, block: _Open) -> tuple[nodes.Statement, _Open]:
+        """``node`` with its conditional expressions and comprehensions replaced by reads
+        of synthetic locals, after laying out the statements that compute them."""
+
+        pending: list[nodes.Statement] = []
+        if isinstance(node, nodes.While):
+            if _has_control_flow(node.condition):
+                raise CFGError(
+                    f"{node.span.display()}: conditional expressions and comprehensions in a "
+                    "loop condition are not supported yet"
+                )
+            return node, block
+        if isinstance(node, nodes.Assign | nodes.AugAssign):
+            node = replace(node, target=self.hoist(node.target, pending), value=self.hoist(node.value, pending))
+        elif isinstance(node, nodes.ExpressionStatement):
+            node = replace(node, expression=self.hoist(node.expression, pending))
+        elif isinstance(node, nodes.Return | nodes.Raise) and node.__class__ is nodes.Return:
+            if node.value is not None:
+                node = replace(node, value=self.hoist(node.value, pending))
+        elif isinstance(node, nodes.Raise):
+            exception = self.hoist(node.exception, pending) if node.exception is not None else None
+            cause = self.hoist(node.cause, pending) if node.cause is not None else None
+            node = replace(node, exception=exception, cause=cause)
+        elif isinstance(node, nodes.Assert):
+            message = self.hoist(node.message, pending) if node.message is not None else None
+            node = replace(node, test=self.hoist(node.test, pending), message=message)
+        elif isinstance(node, nodes.If):
+            node = replace(node, condition=self.hoist(node.condition, pending))
+        elif isinstance(node, nodes.For):
+            node = replace(node, iterable=self.hoist(node.iterable, pending))
+        elif isinstance(node, nodes.With):
+            items = tuple(replace(item, context=self.hoist(item.context, pending)) for item in node.items)
+            node = replace(node, items=items)
+        if not pending:
+            return node, block
+        laid_out = self.sequence(tuple(pending), block, None, node.span)
+        assert laid_out is not None, "hoisted statements always fall through"
+        return node, laid_out
+
+    def hoist(self, node: Any, pending: list[nodes.Statement]) -> Any:
+        """Rewrite one expression tree, appending the statements it needs to ``pending``."""
+
+        if isinstance(node, nodes.Conditional):
+            test = self.hoist(node.test, pending)
+            result = self.hidden("cond", node.span)
+            pending.append(
+                nodes.If(
+                    test,
+                    (nodes.Assign(result, node.body, node.span),),
+                    (nodes.Assign(result, node.orelse, node.span),),
+                    node.span,
+                )
+            )
+            return result
+        if isinstance(node, nodes.Comprehension):
+            return self.comprehension(node, pending)
+        if isinstance(node, nodes.Lambda):
+            defaults = tuple(
+                replace(p, default=self.hoist(p.default, pending)) if p.default is not None else p
+                for p in node.parameters
+            )
+            return replace(node, parameters=defaults)
+        if isinstance(node, tuple):
+            return tuple(self.hoist(item, pending) for item in node)
+        if is_dataclass(node) and not isinstance(node, type):
+            changes = {
+                f.name: self.hoist(getattr(node, f.name), pending)
+                for f in fields(node)
+                if f.name != "span" and _may_hold_expressions(getattr(node, f.name))
+            }
+            return replace(node, **changes) if changes else node
+        return node
+
+    def comprehension(self, node: nodes.Comprehension, pending: list[nodes.Statement]) -> nodes.Name:
+        """Lay a comprehension out as loops filling a synthetic collection."""
+
+        span = node.span
+        result = self.hidden("comp", span)
+        renames: dict[str, str] = {}
+        for generator in node.generators:
+            for name in _bound_names(generator.target):
+                renames[name] = self.hidden(f"var_{name}", generator.target.span).identifier
+        if node.kind == "set":
+            initial: nodes.Expression = nodes.Call(nodes.Name("set", span), (), (), span)
+        elif node.kind == "dict":
+            initial = nodes.Dict((), span)
+        else:
+            initial = nodes.List((), span)
+        pending.append(nodes.Assign(result, initial, span))
+
+        element = _rename(node.element, renames)
+        if node.kind == "dict":
+            assert node.key is not None
+            innermost: nodes.Statement = nodes.Assign(
+                nodes.Subscript(result, _rename(node.key, renames), span), element, span
+            )
+        else:
+            method = "add" if node.kind == "set" else "append"
+            call = nodes.Call(nodes.Attribute(result, method, span), (element,), (), span)
+            innermost = nodes.ExpressionStatement(call, span)
+
+        body: tuple[nodes.Statement, ...] = (innermost,)
+        for index in range(len(node.generators) - 1, -1, -1):
+            generator = node.generators[index]
+            for condition in reversed(generator.conditions):
+                body = (nodes.If(_rename(condition, renames), body, (), condition.span),)
+            iterable = _rename(generator.iterable, renames) if index else self.hoist(generator.iterable, pending)
+            if isinstance(generator.target, nodes.Name):
+                target = nodes.Name(renames[generator.target.identifier], generator.target.span)
+            else:
+                target = self.hidden("item", generator.target.span)
+                body = (nodes.Assign(_rename_target(generator.target, renames), target, generator.target.span), *body)
+            body = (nodes.For(target, iterable, body, False, generator.span),)
+        pending.extend(body)
+        return result
 
     # ------------------------------------------------------------------ blocks
 
@@ -92,6 +217,7 @@ def sequence(
     def statement(
         self, node: nodes.Statement, block: _Open, continuation: BlockId | None
     ) -> _Open | None:
+        node, block = self.desugared(node, block)
         if isinstance(node, nodes.Return):
             self.finish(block, Return(node.value, node.span))
             return None
@@ -201,6 +327,72 @@ def loop(self, keyword: str, span: SourceSpan) -> tuple[BlockId, BlockId]:
         return self.loops[-1]
 
 
+def _may_hold_expressions(value: object) -> bool:
+    return isinstance(value, tuple) or (is_dataclass(value) and not isinstance(value, type))
+
+
+def _has_control_flow(node: object) -> bool:
+    if isinstance(node, nodes.Conditional | nodes.Comprehension):
+        return True
+    if isinstance(node, nodes.Lambda):
+        return False
+    if isinstance(node, tuple):
+        return any(_has_control_flow(item) for item in node)
+    if is_dataclass(node) and not isinstance(node, type):
+        return any(_has_control_flow(getattr(node, f.name)) for f in fields(node) if f.name != "span")
+    return False
+
+
+def _bound_names(target: nodes.Target) -> list[str]:
+    if isinstance(target, nodes.Name):
+        return [target.identifier]
+    if isinstance(target, nodes.Tuple):
+        return [name for element in target.elements for name in _bound_names(element)]  # type: ignore[arg-type]
+    return []
+
+
+def _rename(node: Any, renames: dict[str, str]) -> Any:
+    """``node`` with comprehension variables replaced by their synthetic locals, without
+    entering scopes that rebind them."""
+
+    if not renames:
+        return node
+    if isinstance(node, nodes.Name):
+        return nodes.Name(renames[node.identifier], node.span) if node.identifier in renames else node
+    if isinstance(node, nodes.Lambda):
+        inner = {k: v for k, v in renames.items() if k not in {p.name for p in node.parameters}}
+        return replace(node, body=_rename(node.body, inner))
+    if isinstance(node, nodes.Comprehension):
+        shadowed = {name for g in node.generators for name in _bound_names(g.target)}
+        inner = {k: v for k, v in renames.items() if k not in shadowed}
+        generators = tuple(
+            replace(
+                g,
+                iterable=_rename(g.iterable, renames if index == 0 else inner),
+                conditions=_rename(g.conditions, inner),
+            )
+            for index, g in enumerate(node.generators)
+        )
+        key = _rename(node.key, inner) if node.key is not None else None
+        return replace(node, element=_rename(node.element, inner), generators=generators, key=key)
+    if isinstance(node, tuple):
+        return tuple(_rename(item, renames) for item in node)
+    if is_dataclass(node) and not isinstance(node, type):
+        changes = {
+            f.name: _rename(getattr(node, f.name), renames)
+            for f in fields(node)
+            if f.name != "span" and _may_hold_expressions(getattr(node, f.name))
+        }
+        return replace(node, **changes) if changes else node
+    return node
+
+
+def _rename_target(target: nodes.Target, renames: dict[str, str]) -> nodes.Target:
+    renamed = _rename(target, renames)
+    assert isinstance(renamed, nodes.Name | nodes.Tuple | nodes.Attribute | nodes.Subscript)
+    return renamed
+
+
 def build_cfg(function: nodes.Function) -> CFG:
     return _Builder(function).build()
 

src/coretrace_python/cfg/model.py

@@ -92,6 +92,9 @@ class BasicBlock:
 class CFG:
     entry: BlockId
     blocks: Mapping[BlockId, BasicBlock]
+    # Locals the builder introduced while laying out expression-level control flow
+    # (conditional expressions, comprehensions); the lowering treats them as locals.
+    synthetic_locals: frozenset[str] = frozenset()
     _successors: Mapping[BlockId, tuple[BlockId, ...]] = field(
         init=False, repr=False, compare=False
     )

src/coretrace_python/frontend/ast_adapter.py

@@ -29,7 +29,7 @@ class HIRBuildError(Exception):
     ast.MatMult: "matmul",
 }
 _UNARY_OPERATORS = {ast.Invert: "invert", ast.Not: "not", ast.UAdd: "pos", ast.USub: "neg"}
-_COMPREHENSIONS = {ast.ListComp: "list", ast.SetComp: "set", ast.GeneratorExp: "generator"}
+_COMPREHENSIONS = {ast.ListComp: "list", ast.SetComp: "set", ast.GeneratorExp: "generator", ast.DictComp: "dict"}
 _COMPARE_OPERATORS = {
     ast.Eq: "eq",
     ast.NotEq: "not_eq",
@@ -118,6 +118,17 @@ def expression(self, node: ast.expr) -> nodes.Expression:
             )
         if isinstance(node, ast.Starred):
             return nodes.Starred(self.expression(node.value), span)
+        if isinstance(node, ast.Set):
+            return nodes.Set(self.elements(node.elts), span)
+        if isinstance(node, ast.IfExp):
+            return nodes.Conditional(
+                self.expression(node.test), self.expression(node.body), self.expression(node.orelse), span
+            )
+        if isinstance(node, ast.Lambda):
+            return nodes.Lambda(self.parameters(node.args), self.expression(node.body), span)
+        if isinstance(node, ast.DictComp):
+            generators = tuple(self.generator(generator) for generator in node.generators)
+            return nodes.Comprehension("dict", self.expression(node.value), generators, span, self.expression(node.key))
         if isinstance(node, ast.Attribute):
             return nodes.Attribute(self.expression(node.value), node.attr, span)
         if isinstance(node, ast.Subscript):
@@ -176,9 +187,7 @@ def target(self, node: ast.expr) -> nodes.Target:
     def generator(self, node: ast.comprehension) -> nodes.ComprehensionGenerator:
         if node.is_async:
             self.fail(node.target, "async comprehensions are not supported yet")
-        if not isinstance(node.target, ast.Name):
-            self.fail(node.target, "only a single name is supported as a comprehension target")
-        target = nodes.Name(node.target.id, self.span(node.target))
+        target = self.target(node.target)
         conditions = tuple(self.expression(condition) for condition in node.ifs)
         # ``ast.comprehension`` carries no location; span it from the target to the last clause.
         last = self.span(node.ifs[-1] if node.ifs else node.iter)
@@ -191,11 +200,21 @@ def generator(self, node: ast.comprehension) -> nodes.ComprehensionGenerator:
         )
         return nodes.ComprehensionGenerator(target, self.expression(node.iter), conditions, span)
 
+    def statements(self, node: ast.stmt) -> list[nodes.Statement]:
+        """The HIR statements of one source statement: several for ``a = b = value``,
+        which assigns a hidden local once and every target from it."""
+
+        span = self.span(node)
+        if isinstance(node, ast.Assign) and len(node.targets) > 1:
+            hidden = nodes.Name(f"_coretrace_chain_{span.start_line}_{span.start_column}", span)
+            found: list[nodes.Statement] = [nodes.Assign(hidden, self.expression(node.value), span)]
+            found.extend(nodes.Assign(self.target(target), hidden, self.span(target)) for target in node.targets)
+            return found
+        return [self.statement(node)]
+
     def statement(self, node: ast.stmt) -> nodes.Statement:
         span = self.span(node)
         if isinstance(node, ast.Assign):
-            if len(node.targets) != 1:
-                self.fail(node, "chained assignment is not supported yet")
             return nodes.Assign(self.target(node.targets[0]), self.expression(node.value), span)
         if isinstance(node, ast.AnnAssign):
             # The annotation never affects behaviour; a bare declaration does nothing.
@@ -309,18 +328,18 @@ def statement(self, node: ast.stmt) -> nodes.Statement:
         self.fail(node)
 
     def block(self, statements: list[ast.stmt]) -> tuple[nodes.Statement, ...]:
-        return tuple(self.statement(statement) for statement in statements)
+        return tuple(found for statement in statements for found in self.statements(statement))
 
     def class_definition(self, node: ast.ClassDef) -> nodes.Class:
         if node.keywords:
             self.fail(node, "class keyword arguments are not supported yet")
         bases = tuple(self.expression(base) for base in node.bases)
-        body = tuple(self.statement(statement) for statement in node.body)
+        body = self.block(node.body)
         decorators = tuple(self.expression(d) for d in node.decorator_list)
         return nodes.Class(node.name, bases, body, self.span(node), decorators)
 
     def function(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> nodes.Function:
-        body = tuple(self.statement(statement) for statement in node.body)
+        body = self.block(node.body)
         return nodes.Function(
             name=node.name,
             parameters=self.parameters(node.args),
@@ -357,7 +376,7 @@ def parameter(self, argument: ast.arg, default: ast.expr | None, kind: str) -> n
         return nodes.Parameter(argument.arg, self.span(argument), value, kind, annotation)
 
     def module(self, tree: ast.Module) -> nodes.Module:
-        body = tuple(self.statement(statement) for statement in tree.body)
+        body = self.block(tree.body)
         if tree.body:
             first_span = self.span(tree.body[0])
             last_span = self.span(tree.body[-1])