feat: control-flow graphs

src/coretrace_python/cfg/__init__.py

@@ -0,0 +1,32 @@
+"""Per-function control-flow graphs over PyHIR (architecture §5)."""
+
+from coretrace_python.cfg.builder import CFGAnalysis, build_cfg
+from coretrace_python.cfg.model import (
+    CFG,
+    BasicBlock,
+    BlockId,
+    Branch,
+    CFGError,
+    ForEach,
+    Jump,
+    Raise,
+    Return,
+    Terminator,
+    targets,
+)
+
+__all__ = [
+    "CFG",
+    "BasicBlock",
+    "BlockId",
+    "Branch",
+    "CFGAnalysis",
+    "CFGError",
+    "ForEach",
+    "Jump",
+    "Raise",
+    "Return",
+    "Terminator",
+    "build_cfg",
+    "targets",
+]

src/coretrace_python/cfg/builder.py

@@ -0,0 +1,153 @@
+"""Build a control-flow graph from a PyHIR function."""
+
+from __future__ import annotations
+
+from typing import ClassVar
+
+from coretrace_python.analysis import AnalysisContext, FunctionAnalysis
+from coretrace_python.cfg.model import (
+    CFG,
+    BasicBlock,
+    BlockId,
+    Branch,
+    CFGError,
+    ForEach,
+    Jump,
+    Raise,
+    Return,
+    Terminator,
+)
+from coretrace_python.hir import nodes
+from coretrace_python.source import SourceSpan
+
+
+class _Open:
+    """A block that is still collecting statements and has no terminator yet."""
+
+    def __init__(self, block_id: BlockId) -> None:
+        self.id = block_id
+        self.statements: list[nodes.Statement] = []
+
+
+class _Builder:
+    def __init__(self, function: nodes.Function) -> None:
+        self.function = function
+        self.blocks: dict[BlockId, BasicBlock] = {}
+        self.counters: dict[str, int] = {}
+        self.loops: list[tuple[BlockId, BlockId]] = []
+
+    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)
+
+    # ------------------------------------------------------------------ blocks
+
+    def new_id(self, kind: str) -> BlockId:
+        self.counters[kind] = self.counters.get(kind, 0) + 1
+        return BlockId(f"{kind}_{self.counters[kind]}")
+
+    def finish(self, block: _Open, terminator: Terminator) -> None:
+        self.blocks[block.id] = BasicBlock(block.id, tuple(block.statements), terminator)
+
+    def header(self, block_id: BlockId, terminator: Terminator) -> None:
+        self.blocks[block_id] = BasicBlock(block_id, (), terminator)
+
+    # ------------------------------------------------------------------ statements
+
+    def sequence(
+        self,
+        statements: tuple[nodes.Statement, ...],
+        block: _Open,
+        continuation: BlockId | None,
+        join_span: SourceSpan,
+    ) -> _Open | None:
+        """Lay ``statements`` out from ``block``.
+
+        Returns the block left open at the end, or ``None`` when control left the
+        sequence. With a ``continuation``, an open end jumps there instead.
+        """
+
+        current: _Open | None = block
+        last_index = len(statements) - 1
+        for index, statement in enumerate(statements):
+            assert current is not None
+            is_last = index == last_index
+            current = self.statement(statement, current, continuation if is_last else None)
+            if current is None and not is_last:
+                current = _Open(self.new_id("dead"))
+        if current is not None and continuation is not None:
+            self.finish(current, Jump(continuation, join_span))
+            return None
+        return current
+
+    def statement(
+        self, node: nodes.Statement, block: _Open, continuation: BlockId | None
+    ) -> _Open | None:
+        if isinstance(node, nodes.Return):
+            self.finish(block, Return(node.value, node.span))
+            return None
+        if isinstance(node, nodes.Raise):
+            self.finish(block, Raise(node.exception, node.span))
+            return None
+        if isinstance(node, nodes.Break):
+            self.finish(block, Jump(self.loop("break", node.span)[1], node.span))
+            return None
+        if isinstance(node, nodes.Continue):
+            self.finish(block, Jump(self.loop("continue", node.span)[0], node.span))
+            return None
+        if isinstance(node, nodes.If):
+            return self.conditional(node, block, continuation)
+        if isinstance(node, nodes.While | nodes.For):
+            return self.loop_statement(node, block, continuation)
+        block.statements.append(node)
+        return block
+
+    def conditional(
+        self, node: nodes.If, block: _Open, continuation: BlockId | None
+    ) -> _Open | None:
+        merge = continuation if continuation is not None else self.new_id("merge")
+        then_id = self.new_id("then")
+        else_id = self.new_id("else") if node.orelse else merge
+        self.finish(block, Branch(node.condition, then_id, else_id, node.span))
+        self.sequence(node.body, _Open(then_id), merge, node.span)
+        if node.orelse:
+            self.sequence(node.orelse, _Open(else_id), merge, node.span)
+        return None if continuation is not None else _Open(merge)
+
+    def loop_statement(
+        self, node: nodes.While | nodes.For, block: _Open, continuation: BlockId | None
+    ) -> _Open | None:
+        header_id = self.new_id("loop")
+        body_id = self.new_id("body")
+        exit_id = continuation if continuation is not None else self.new_id("exit")
+        self.finish(block, Jump(header_id, node.span))
+        if isinstance(node, nodes.While):
+            self.header(header_id, Branch(node.condition, body_id, exit_id, node.span))
+        else:
+            self.header(header_id, ForEach(node.target, node.iterable, body_id, exit_id, node.span))
+        self.loops.append((header_id, exit_id))
+        try:
+            self.sequence(node.body, _Open(body_id), header_id, node.span)
+        finally:
+            self.loops.pop()
+        return None if continuation is not None else _Open(exit_id)
+
+    def loop(self, keyword: str, span: SourceSpan) -> tuple[BlockId, BlockId]:
+        if not self.loops:
+            raise CFGError(f"{span.display()}: '{keyword}' outside loop")
+        return self.loops[-1]
+
+
+def build_cfg(function: nodes.Function) -> CFG:
+    return _Builder(function).build()
+
+
+class CFGAnalysis(FunctionAnalysis[CFG]):
+    name: ClassVar[str] = "cfg.function"
+
+    @classmethod
+    def compute(cls, ctx: AnalysisContext, function: nodes.Function) -> CFG:
+        return build_cfg(function)

src/coretrace_python/cfg/model.py

@@ -0,0 +1,158 @@
+"""Control-flow graph model (architecture §5).
+
+A block holds straight-line PyHIR statements and ends in exactly one terminator.
+Python's iteration protocol stays explicit as ``ForEach`` instead of being encoded as
+a generic switch; the CFG must remain semantically faithful to Python.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from types import MappingProxyType
+from typing import TypeAlias
+
+from coretrace_python.hir import nodes
+from coretrace_python.source import SourceSpan
+
+
+class CFGError(Exception):
+    """A malformed graph or a control statement used outside its construct."""
+
+
+@dataclass(frozen=True, order=True)
+class BlockId:
+    value: str
+
+    def __str__(self) -> str:
+        return self.value
+
+
+@dataclass(frozen=True)
+class Branch:
+    condition: nodes.Expression
+    then_block: BlockId
+    else_block: BlockId
+    span: SourceSpan
+
+
+@dataclass(frozen=True)
+class Jump:
+    target: BlockId
+    span: SourceSpan
+
+
+@dataclass(frozen=True)
+class Return:
+    value: nodes.Expression | None
+    span: SourceSpan
+
+
+@dataclass(frozen=True)
+class Raise:
+    exception: nodes.Expression | None
+    span: SourceSpan
+
+
+@dataclass(frozen=True)
+class ForEach:
+    """Bind the next item of ``iterable`` to ``target`` and enter ``body``, or ``exit``."""
+
+    target: nodes.Name
+    iterable: nodes.Expression
+    body: BlockId
+    exit: BlockId
+    span: SourceSpan
+
+
+Terminator: TypeAlias = Branch | Jump | Return | Raise | ForEach
+
+
+def targets(terminator: Terminator) -> tuple[BlockId, ...]:
+    if isinstance(terminator, Branch):
+        return (terminator.then_block, terminator.else_block)
+    if isinstance(terminator, Jump):
+        return (terminator.target,)
+    if isinstance(terminator, ForEach):
+        return (terminator.body, terminator.exit)
+    return ()
+
+
+@dataclass(frozen=True)
+class BasicBlock:
+    id: BlockId
+    statements: tuple[nodes.Statement, ...]
+    terminator: Terminator
+
+
+@dataclass(frozen=True)
+class CFG:
+    entry: BlockId
+    blocks: Mapping[BlockId, BasicBlock]
+    _successors: Mapping[BlockId, tuple[BlockId, ...]] = field(
+        init=False, repr=False, compare=False
+    )
+    _predecessors: Mapping[BlockId, tuple[BlockId, ...]] = field(
+        init=False, repr=False, compare=False
+    )
+
+    def __post_init__(self) -> None:
+        blocks = MappingProxyType(dict(self.blocks))
+        if self.entry not in blocks:
+            raise CFGError(f"entry block {self.entry} does not exist")
+        predecessors: dict[BlockId, list[BlockId]] = {block_id: [] for block_id in blocks}
+        successors: dict[BlockId, tuple[BlockId, ...]] = {}
+        for block_id, block in blocks.items():
+            if block.id != block_id:
+                raise CFGError(f"block {block.id} is stored under {block_id}")
+            successors[block_id] = targets(block.terminator)
+            for target in successors[block_id]:
+                if target not in blocks:
+                    raise CFGError(f"block {block_id} targets unknown block {target}")
+                predecessors[target].append(block_id)
+        object.__setattr__(self, "blocks", blocks)
+        object.__setattr__(self, "_successors", MappingProxyType(successors))
+        object.__setattr__(
+            self,
+            "_predecessors",
+            MappingProxyType({k: tuple(v) for k, v in predecessors.items()}),
+        )
+
+    def block(self, block_id: BlockId) -> BasicBlock:
+        return self.blocks[block_id]
+
+    def successors(self, block_id: BlockId) -> tuple[BlockId, ...]:
+        return self._successors[block_id]
+
+    def predecessors(self, block_id: BlockId) -> tuple[BlockId, ...]:
+        return self._predecessors[block_id]
+
+    def reachable(self) -> frozenset[BlockId]:
+        seen: set[BlockId] = set()
+        pending = [self.entry]
+        while pending:
+            block_id = pending.pop()
+            if block_id not in seen:
+                seen.add(block_id)
+                pending.extend(self.successors(block_id))
+        return frozenset(seen)
+
+    def back_edges(self) -> frozenset[tuple[BlockId, BlockId]]:
+        """Edges whose target is on the depth-first path from the entry to their source."""
+
+        found: set[tuple[BlockId, BlockId]] = set()
+        on_path: set[BlockId] = set()
+        finished: set[BlockId] = set()
+
+        def visit(block_id: BlockId) -> None:
+            on_path.add(block_id)
+            for successor in self.successors(block_id):
+                if successor in on_path:
+                    found.add((block_id, successor))
+                elif successor not in finished:
+                    visit(successor)
+            on_path.remove(block_id)
+            finished.add(block_id)
+
+        visit(self.entry)
+        return frozenset(found)

src/coretrace_python/frontend/ast_adapter.py

@@ -154,6 +154,36 @@ def statement(self, node: ast.stmt) -> nodes.Statement:
                 nodes.ImportAlias(alias.name, alias.asname, self.span(alias)) for alias in node.names
             )
             return nodes.ImportFrom(node.module, aliases, node.level, span)
+        if isinstance(node, ast.If):
+            return nodes.If(
+                self.expression(node.test), self.block(node.body), self.block(node.orelse), span
+            )
+        if isinstance(node, ast.While):
+            if node.orelse:
+                self.fail(node.orelse[0], "loop else clauses are not supported yet")
+            return nodes.While(self.expression(node.test), self.block(node.body), span)
+        if isinstance(node, (ast.For, ast.AsyncFor)):
+            if node.orelse:
+                self.fail(node.orelse[0], "loop else clauses are not supported yet")
+            if not isinstance(node.target, ast.Name):
+                self.fail(node.target, "only a single name is supported as a loop target")
+            target = nodes.Name(node.target.id, self.span(node.target))
+            return nodes.For(
+                target,
+                self.expression(node.iter),
+                self.block(node.body),
+                isinstance(node, ast.AsyncFor),
+                span,
+            )
+        if isinstance(node, ast.Break):
+            return nodes.Break(span)
+        if isinstance(node, ast.Continue):
+            return nodes.Continue(span)
+        if isinstance(node, ast.Raise):
+            if node.cause is not None:
+                self.fail(node.cause, "raise from is not supported yet")
+            exception = self.expression(node.exc) if node.exc is not None else None
+            return nodes.Raise(exception, span)
         if isinstance(node, ast.Global):
             return nodes.Global(tuple(node.names), span)
         if isinstance(node, ast.Nonlocal):
@@ -164,6 +194,9 @@ def statement(self, node: ast.stmt) -> nodes.Statement:
             return self.class_definition(node)
         self.fail(node)
 
+    def block(self, statements: list[ast.stmt]) -> tuple[nodes.Statement, ...]:
+        return tuple(self.statement(statement) for statement in statements)
+
     def class_definition(self, node: ast.ClassDef) -> nodes.Class:
         if node.decorator_list:
             self.fail(node, "decorated classes are not supported yet")