From 84ef895cb8c065c1a582600822b582437f893252 Mon Sep 17 00:00:00 2001 From: Hugo Date: Thu, 3 Sep 2026 07:05:24 +0200 Subject: [PATCH 1/2] test: define PyIR over the CFG --- tests/test_analysis_manager.py | 5 +- tests/test_emit_ir.py | 158 +++++++++++++++++++++++++++++++++ tests/test_ir.py | 141 +++++++++++++++++++++++++++-- 3 files changed, 293 insertions(+), 11 deletions(-) diff --git a/tests/test_analysis_manager.py b/tests/test_analysis_manager.py index 64d989a..8456950 100644 --- a/tests/test_analysis_manager.py +++ b/tests/test_analysis_manager.py @@ -21,6 +21,7 @@ import pytest +from coretrace_python.cfg import CFGAnalysis from coretrace_python.frontend import build_hir from coretrace_python.hir import nodes from coretrace_python.ir.lowering import lower_module @@ -318,7 +319,7 @@ def test_analyses_declare_names_and_versions() -> None: def test_pyir_is_computed_per_function_on_demand() -> None: - engine = manager(*SEMANTIC_ANALYSES, PyIRAnalysis) + engine = manager(*SEMANTIC_ANALYSES, CFGAnalysis, PyIRAnalysis) module_functions = functions(engine.module) first = engine.get(PyIRAnalysis, module_functions["first"]) @@ -331,7 +332,7 @@ def test_pyir_is_computed_per_function_on_demand() -> None: def test_module_ir_matches_lower_module() -> None: module = hir() engine = AnalysisManager(module) - engine.register(*SEMANTIC_ANALYSES, PyIRAnalysis, ModuleIRAnalysis) + engine.register(*SEMANTIC_ANALYSES, CFGAnalysis, PyIRAnalysis, ModuleIRAnalysis) module_ir = engine.get(ModuleIRAnalysis) diff --git a/tests/test_emit_ir.py b/tests/test_emit_ir.py index b3af862..c141f9b 100644 --- a/tests/test_emit_ir.py +++ b/tests/test_emit_ir.py @@ -64,3 +64,161 @@ def test_emit_ir_for_local_assignment(tmp_path, capsys) -> None: " return %3\n" "}\n" ) + + +# --------------------------------------------------------------------------- control flow +# ``--emit-ir`` prints one block per CFG block with explicit terminators +# (docs/architecture.md §5 and §6). Expected to remain red until lowering runs over the CFG. + + +def test_emit_ir_for_if_without_else(tmp_path, capsys) -> None: + # The example of §5: the sanitizer runs on one path only. + source = tmp_path / "guard.py" + source.write_text( + "def f(safe, x):\n" + " if safe:\n" + " x = sanitize(x)\n" + " sink(x)\n", + encoding="utf-8", + ) + + assert main(["--emit-ir", str(source)]) == 0 + assert capsys.readouterr().out == ( + "func @f(%0, %1) {\n" + "entry:\n" + ' store_local "x", %1\n' + " branch %0, then_1, merge_1\n" + "then_1:\n" + " %2 = global 'sanitize'\n" + ' %3 = load_local "x"\n' + " %4 = call %2(%3)\n" + ' store_local "x", %4\n' + " jump merge_1\n" + "merge_1:\n" + " %5 = global 'sink'\n" + ' %6 = load_local "x"\n' + " %7 = call %5(%6)\n" + " return\n" + "}\n" + ) + + +def test_emit_ir_for_while_loop(tmp_path, capsys) -> None: + source = tmp_path / "loop.py" + source.write_text( + "def f(n):\n" + " while n > 0:\n" + " n = n - 1\n" + " return n\n", + encoding="utf-8", + ) + + assert main(["--emit-ir", str(source)]) == 0 + assert capsys.readouterr().out == ( + "func @f(%0) {\n" + "entry:\n" + ' store_local "n", %0\n' + " jump loop_1\n" + "loop_1:\n" + ' %1 = load_local "n"\n' + " %2 = const 0\n" + " %3 = compare.gt %1, %2\n" + " branch %3, body_1, exit_1\n" + "body_1:\n" + ' %4 = load_local "n"\n' + " %5 = const 1\n" + " %6 = binary.sub %4, %5\n" + ' store_local "n", %6\n' + " jump loop_1\n" + "exit_1:\n" + ' %7 = load_local "n"\n' + " return %7\n" + "}\n" + ) + + +def test_emit_ir_for_for_loop(tmp_path, capsys) -> None: + source = tmp_path / "sum.py" + source.write_text( + "def f(items):\n" + " total = 0\n" + " for item in items:\n" + " total = total + item\n" + " return total\n", + encoding="utf-8", + ) + + assert main(["--emit-ir", str(source)]) == 0 + assert capsys.readouterr().out == ( + "func @f(%0) {\n" + "entry:\n" + " %1 = const 0\n" + ' store_local "total", %1\n' + " %2 = get_iter %0\n" + " jump loop_1\n" + "loop_1:\n" + ' for_next %2 -> "item", body_1, exit_1\n' + "body_1:\n" + ' %3 = load_local "total"\n' + ' %4 = load_local "item"\n' + " %5 = binary.add %3, %4\n" + ' store_local "total", %5\n' + " jump loop_1\n" + "exit_1:\n" + ' %6 = load_local "total"\n' + " return %6\n" + "}\n" + ) + + +def test_emit_ir_for_raise(tmp_path, capsys) -> None: + source = tmp_path / "raise.py" + source.write_text( + "def f(a):\n" + " if a:\n" + " raise ValueError(a)\n" + " return a\n\n" + "def g():\n" + " raise\n", + encoding="utf-8", + ) + + assert main(["--emit-ir", str(source)]) == 0 + output = capsys.readouterr().out + assert ( + "then_1:\n" + " %1 = symbol @python.builtins.ValueError\n" + " %2 = call %1(%0)\n" + " raise %2\n" + "merge_1:\n" + " return %0\n" + ) in output + assert "func @g() {\nentry:\n raise\n}" in output + + +def test_emit_ir_for_break_and_continue(tmp_path, capsys) -> None: + source = tmp_path / "control.py" + source.write_text( + "def f(items):\n" + " for item in items:\n" + " if item:\n" + " continue\n" + " if not item:\n" + " break\n" + " return 0\n", + encoding="utf-8", + ) + + assert main(["--emit-ir", str(source)]) == 0 + output = capsys.readouterr().out + assert "then_1:\n jump loop_1\n" in output + assert "then_2:\n jump exit_1\n" in output + + +def test_emit_ir_keeps_unreachable_code(tmp_path, capsys) -> None: + source = tmp_path / "dead.py" + source.write_text("def f():\n return 1\n x = 2\n", encoding="utf-8") + + assert main(["--emit-ir", str(source)]) == 0 + output = capsys.readouterr().out + assert "dead_1:\n %1 = const 2\n" in output diff --git a/tests/test_ir.py b/tests/test_ir.py index 7fd6319..714dcc8 100644 --- a/tests/test_ir.py +++ b/tests/test_ir.py @@ -1,14 +1,137 @@ +"""Acceptance tests for PyIR over the control-flow graph (``docs/architecture.md`` §6). + +PyIR mirrors the CFG: one immutable block per CFG block, in CFG order, each ending in +an explicit terminator (``Branch``, ``Jump``, ``Return``, ``Raise`` or ``ForNext``). +Values are numbered once per function, so a value defined in a dominating block can +be used in later blocks. Parameters that are never reassigned stay plain values; +reassigned ones are stored to a local at entry so every block reads the same slot. + +The multi-block cases are expected to remain red until lowering consumes +``CFGAnalysis`` and the IR model carries terminators. +""" + +from __future__ import annotations + +import dataclasses + +import pytest + +from coretrace_python.analysis import AnalysisManager +from coretrace_python.cfg import CFGAnalysis from coretrace_python.frontend import build_hir -from coretrace_python.ir.lowering import lower_module -from coretrace_python.ir.model import BinaryOp, Return +from coretrace_python.hir import nodes +from coretrace_python.ir.lowering import PyIRAnalysis, lower_module +from coretrace_python.ir.model import BasicBlock, BinaryOp, FunctionIR, LoadLocal, StoreLocal +from coretrace_python.semantic import SEMANTIC_ANALYSES from coretrace_python.source import SourceManager +try: + from coretrace_python.ir.model import Branch, ForNext, GetIter, Jump, Raise, Return +except ImportError as error: # pragma: no cover - red until terminators exist + MISSING = error +else: + MISSING = None + + +@pytest.fixture(autouse=True) +def require_terminators() -> None: + if MISSING is not None: + pytest.fail(f"PyIR terminators are not implemented yet: {MISSING}") + + +def lower(source_text: str) -> FunctionIR: + module = lower_module(build_hir(SourceManager().add_source("ir.py", source_text))) + return module.functions[0] + + +def block_ids(function: FunctionIR) -> list[str]: + return [str(block.id) for block in function.blocks] + def test_lowering_preserves_locations() -> None: - source = SourceManager().add_source("add.py", "def add(a, b):\n return a + b\n") - module = lower_module(build_hir(source)) - instructions = module.functions[0].blocks[0].instructions - assert isinstance(instructions[0], BinaryOp) - assert instructions[0].location.start_line == 2 - assert isinstance(instructions[1], Return) - assert instructions[1].location.source_id.value == "add.py" + function = lower("def add(a, b):\n return a + b\n") + entry = function.blocks[0] + + assert isinstance(entry.instructions[0], BinaryOp) + assert entry.instructions[0].location.start_line == 2 + assert isinstance(entry.terminator, Return) + assert entry.terminator.location.source_id.value == "ir.py" + + +def test_blocks_follow_the_cfg_and_are_immutable() -> None: + function = lower("def f(a):\n if a:\n return 1\n return 0\n") + + assert str(function.entry) == "entry" + assert block_ids(function) == ["entry", "then_1", "merge_1"] + entry = function.blocks[0] + assert isinstance(entry, BasicBlock) + assert isinstance(entry.instructions, tuple) + assert isinstance(entry.terminator, Branch) + assert str(entry.terminator.then_block) == "then_1" + assert str(entry.terminator.else_block) == "merge_1" + with pytest.raises(dataclasses.FrozenInstanceError): + entry.terminator = entry.terminator # type: ignore[misc] + + +def test_values_are_numbered_once_per_function() -> None: + function = lower("def f(a):\n x = a + 1\n if x:\n y = x + 1\n return x\n") + ids = [ + instruction.result.id + for block in function.blocks + for instruction in block.instructions + if instruction.result is not None + ] + + assert ids == sorted(ids) + assert len(ids) == len(set(ids)) + + +def test_reassigned_parameters_are_stored_at_entry() -> None: + function = lower("def f(n, limit):\n while n < limit:\n n = n + 1\n return n\n") + entry = function.blocks[0] + + assert isinstance(entry.instructions[0], StoreLocal) + assert entry.instructions[0].name == "n" + assert entry.instructions[0].value == function.parameters[0] + assert all( + not (isinstance(i, StoreLocal) and i.name == "limit") + for block in function.blocks + for i in block.instructions + ) + header = function.blocks[1] + assert isinstance(header.instructions[0], LoadLocal) + assert header.instructions[0].name == "n" + + +def test_for_loops_take_the_iterator_before_the_header() -> None: + function = lower("def f(items):\n for item in items:\n pass\n return item\n") + entry, header = function.blocks[0], function.blocks[1] + + assert isinstance(entry.instructions[-1], GetIter) + assert entry.instructions[-1].iterable == function.parameters[0] + assert isinstance(entry.terminator, Jump) + assert header.instructions == () + assert isinstance(header.terminator, ForNext) + assert header.terminator.iterator == entry.instructions[-1].result + assert header.terminator.target == "item" + assert str(header.terminator.body) == "body_1" + assert str(header.terminator.exit) == "exit_1" + + +def test_raise_lowers_to_a_terminator() -> None: + function = lower("def f(a):\n raise ValueError(a)\n") + entry = function.blocks[0] + + assert isinstance(entry.terminator, Raise) + assert entry.terminator.exception == entry.instructions[-1].result + + +def test_pyir_requires_the_cfg() -> None: + module = build_hir(SourceManager().add_source("ir.py", "def f():\n pass\n")) + manager = AnalysisManager(module) + manager.register(*SEMANTIC_ANALYSES, CFGAnalysis, PyIRAnalysis) + function = next(s for s in module.body if isinstance(s, nodes.Function)) + + assert CFGAnalysis in PyIRAnalysis.requires + manager.get(PyIRAnalysis, function) + assert manager.is_cached(CFGAnalysis, function) From a3bccc22262f6c75f361bf4849c3a2471c0d82b9 Mon Sep 17 00:00:00 2001 From: Hugo Date: Thu, 3 Sep 2026 07:05:24 +0200 Subject: [PATCH 2/2] feat: lower PyIR over the control-flow graph --- README.md | 5 +- src/coretrace_python/ir/lowering.py | 142 ++++++++++++++++++++++------ src/coretrace_python/ir/model.py | 63 ++++++++++-- src/coretrace_python/ir/printer.py | 33 ++++++- 4 files changed, 202 insertions(+), 41 deletions(-) diff --git a/README.md b/README.md index 4d14ff7..846941a 100644 --- a/README.md +++ b/README.md @@ -33,7 +33,10 @@ PYTHONPATH=src python -m coretrace_python --emit-ir example.py ``` Currently supported inside functions: arguments, assignments, names, constants, binary -operations, unary operations, comparisons, calls, attribute access, indexing, and returns. +operations, unary operations, comparisons, calls, attribute access, indexing, returns, +`if`/`elif`/`else`, `while`, `for`, `break`, `continue` and `raise`. Each function is +emitted as its control-flow graph: one block per basic block, ending in an explicit +terminator (`branch`, `jump`, `return`, `raise`, `for_next`). Unsupported syntax produces a source-located diagnostic and a non-zero exit status. Names are resolved to canonical symbols through lexical scopes, imports and builtins. diff --git a/src/coretrace_python/ir/lowering.py b/src/coretrace_python/ir/lowering.py index 0476e06..a01ebbd 100644 --- a/src/coretrace_python/ir/lowering.py +++ b/src/coretrace_python/ir/lowering.py @@ -1,10 +1,11 @@ -"""Lower parser-independent PyHIR into analysis-oriented PyIR.""" +"""Lower parser-independent PyHIR into analysis-oriented PyIR, one CFG block at a time.""" from __future__ import annotations from dataclasses import dataclass, field from typing import ClassVar, NoReturn +from coretrace_python import cfg as control_flow from coretrace_python.analysis import ( Analysis, AnalysisContext, @@ -12,22 +13,31 @@ AnyAnalysis, FunctionAnalysis, ) +from coretrace_python.cfg import CFG, BlockId, CFGAnalysis from coretrace_python.hir import nodes +from coretrace_python.hir.visitors import Node, children from coretrace_python.ir.model import ( BasicBlock, BinaryOp, + Branch, Call, Compare, Constant, + ForNext, FunctionIR, GetAttr, GetItem, + GetIter, Global, + Instruction, + Jump, LoadLocal, ModuleIR, + Raise, Return, StoreLocal, Symbol, + Terminator, UnaryOp, Value, ValueInstruction, @@ -52,10 +62,11 @@ class _FunctionLowerer: symbols: SymbolTable scopes: ScopeTable scope: Scope + cfg: CFG next_value_id: int = 0 - assigned: set[str] = field(default_factory=set) parameters: dict[str, Value] = field(default_factory=dict) - block: BasicBlock = field(default_factory=lambda: BasicBlock("entry")) + instructions: list[Instruction] = field(default_factory=list) + iterators: dict[BlockId, Value] = field(default_factory=dict) def fail( self, @@ -71,15 +82,17 @@ def new_value(self) -> Value: return value def emit(self, instruction: ValueInstruction) -> Value: - self.block.instructions.append(instruction) + self.instructions.append(instruction) return instruction.result - def emit_effect(self, instruction: StoreLocal | Return) -> None: - self.block.instructions.append(instruction) + def emit_effect(self, instruction: StoreLocal) -> None: + self.instructions.append(instruction) def resolve(self, name: str) -> Resolution: return self.scopes.resolve(self.scope.id, name) + # ------------------------------------------------------------------ expressions + def imported_symbol(self, node: nodes.Expression) -> SymbolId | None: if isinstance(node, nodes.Name): return self.symbols.resolve(self.scope.id, node.identifier) @@ -98,7 +111,7 @@ def expression(self, node: nodes.Expression) -> Value: 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)) - if node.identifier in self.parameters and node.identifier not in self.assigned: + if node.identifier in self.parameters: return self.parameters[node.identifier] return self.emit(LoadLocal(self.new_value(), node.span, node.identifier)) if isinstance(node, nodes.Constant): @@ -129,45 +142,120 @@ def expression(self, node: nodes.Expression) -> Value: return self.emit(GetItem(self.new_value(), node.span, object_value, key)) self.fail(node) + # ------------------------------------------------------------------ statements + + def store(self, target: nodes.Name, value: Value) -> None: + if self.resolve(target.identifier).kind is not ResolutionKind.LOCAL: + self.fail(target, "assignment to a global or nonlocal name is not supported yet") + self.emit_effect(StoreLocal(None, target.span, target.identifier, value)) + def statement(self, node: nodes.Statement) -> None: if isinstance(node, nodes.Assign): - if self.resolve(node.target.identifier).kind is not ResolutionKind.LOCAL: - self.fail(node, "assignment to a global or nonlocal name is not supported yet") - value = self.expression(node.value) - self.emit_effect(StoreLocal(None, node.span, node.target.identifier, value)) - self.assigned.add(node.target.identifier) - return - if isinstance(node, nodes.Return): - return_value = self.expression(node.value) if node.value is not None else None - self.emit_effect(Return(None, node.span, return_value)) + self.store(node.target, self.expression(node.value)) return if isinstance(node, nodes.ExpressionStatement): self.expression(node.expression) return - if isinstance(node, (nodes.Pass, nodes.Global, nodes.Import, nodes.ImportFrom)): + if isinstance(node, nodes.Pass | nodes.Global | nodes.Import | nodes.ImportFrom): # Declarations and imports are already applied by the semantic analyses. return self.fail(node) + # ------------------------------------------------------------------ blocks + + def terminator(self, block: control_flow.BasicBlock) -> Terminator: + terminator = block.terminator + if isinstance(terminator, control_flow.Return): + value = self.expression(terminator.value) if terminator.value is not None else None + return Return(terminator.span, value) + if isinstance(terminator, control_flow.Branch): + condition = self.expression(terminator.condition) + return Branch(terminator.span, condition, terminator.then_block, terminator.else_block) + if isinstance(terminator, control_flow.Raise): + exception = ( + self.expression(terminator.exception) if terminator.exception is not None else None + ) + return Raise(terminator.span, exception) + if isinstance(terminator, control_flow.Jump): + self.enter_loop(block.id, terminator.target) + return Jump(terminator.span, terminator.target) + target = terminator.target + if self.resolve(target.identifier).kind is not ResolutionKind.LOCAL: + self.fail(target, "assignment to a global or nonlocal name is not supported yet") + return ForNext( + terminator.span, + self.iterators[block.id], + target.identifier, + terminator.body, + terminator.exit, + ) + + def enter_loop(self, source: BlockId, target: BlockId) -> None: + """Take the iterator of a ``for`` loop in the block that enters its header.""" + + header = self.cfg.block(target).terminator + if isinstance(header, control_flow.ForEach) and (source, target) not in self.cfg.back_edges(): + iterable = self.expression(header.iterable) + self.iterators[target] = self.emit(GetIter(self.new_value(), header.span, iterable)) + def function(self, node: nodes.Function) -> FunctionIR: - parameters = tuple(self.new_value() for _ in node.parameters) - parameter_names = (parameter.name for parameter in node.parameters) - self.parameters.update(zip(parameter_names, parameters, strict=True)) - for statement in node.body: - self.statement(statement) - return FunctionIR(node.name, parameters, (self.block,), node.span) + parameter_values = tuple(self.new_value() for _ in node.parameters) + reassigned = _reassigned_parameters(node) + self.parameters = { + parameter.name: value + for parameter, value in zip(node.parameters, parameter_values, strict=True) + if parameter.name not in reassigned + } + 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): + if parameter.name in reassigned: + self.emit_effect(StoreLocal(None, parameter.span, parameter.name, value)) + for statement in cfg_block.statements: + self.statement(statement) + terminator = self.terminator(cfg_block) + blocks.append(BasicBlock(cfg_block.id, tuple(self.instructions), terminator)) + return FunctionIR(node.name, parameter_values, self.cfg.entry, tuple(blocks), node.span) + + +def _reassigned_parameters(function: nodes.Function) -> frozenset[str]: + """Parameters that the function body assigns to, outside nested scopes.""" + + assigned: set[str] = set() + + def walk(node: Node) -> None: + if isinstance(node, nodes.Assign | nodes.For): + assigned.add(node.target.identifier) + if isinstance(node, nodes.Function | nodes.Class | nodes.Comprehension): + return + for child in children(node): + walk(child) + + for statement in function.body: + walk(statement) + return frozenset(assigned & {parameter.name for parameter in function.parameters}) class PyIRAnalysis(FunctionAnalysis[FunctionIR]): - """Lower one function to PyIR on demand.""" + """Lower one function to PyIR on demand, following its control-flow graph.""" name: ClassVar[str] = "ir.pyir" - requires: ClassVar[frozenset[AnyAnalysis]] = frozenset({ScopeAnalysis, SymbolAnalysis}) + requires: ClassVar[frozenset[AnyAnalysis]] = frozenset( + {ScopeAnalysis, SymbolAnalysis, CFGAnalysis} + ) @classmethod def compute(cls, ctx: AnalysisContext, function: nodes.Function) -> FunctionIR: scopes = ctx.get(ScopeAnalysis) - lowerer = _FunctionLowerer(ctx.get(SymbolAnalysis), scopes, scopes.scope_for(function)) + lowerer = _FunctionLowerer( + ctx.get(SymbolAnalysis), + scopes, + scopes.scope_for(function), + ctx.get(CFGAnalysis, function), + ) return lowerer.function(function) @@ -204,5 +292,5 @@ def lower_module(module: nodes.Module) -> ModuleIR: """Lower a whole module through a fresh analysis manager.""" manager = AnalysisManager(module) - manager.register(*SEMANTIC_ANALYSES, PyIRAnalysis, ModuleIRAnalysis) + manager.register(*SEMANTIC_ANALYSES, CFGAnalysis, PyIRAnalysis, ModuleIRAnalysis) return manager.get(ModuleIRAnalysis) diff --git a/src/coretrace_python/ir/model.py b/src/coretrace_python/ir/model.py index 9ef738a..a129fa3 100644 --- a/src/coretrace_python/ir/model.py +++ b/src/coretrace_python/ir/model.py @@ -3,6 +3,7 @@ from dataclasses import dataclass from typing import TypeAlias +from coretrace_python.cfg import BlockId from coretrace_python.semantic.symbols import SymbolId from coretrace_python.source import SourceSpan @@ -84,10 +85,10 @@ class GetItem: @dataclass(frozen=True) -class Return: - result: None +class GetIter: + result: Value location: SourceSpan - value: Value | None + iterable: Value @dataclass(frozen=True) @@ -115,26 +116,68 @@ class StoreLocal: | Call | GetAttr | GetItem + | GetIter | LoadLocal ) -EffectInstruction: TypeAlias = StoreLocal | Return +EffectInstruction: TypeAlias = StoreLocal Instruction: TypeAlias = ValueInstruction | EffectInstruction +# --------------------------------------------------------------------------- terminators + + @dataclass(frozen=True) -class BasicBlock: - name: str - instructions: list[Instruction] +class Return: + location: SourceSpan + value: Value | None + + +@dataclass(frozen=True) +class Branch: + location: SourceSpan + condition: Value + then_block: BlockId + else_block: BlockId - def __init__(self, name: str, instructions: list[Instruction] | None = None) -> None: - object.__setattr__(self, "name", name) - object.__setattr__(self, "instructions", instructions or []) + +@dataclass(frozen=True) +class Jump: + location: SourceSpan + target: BlockId + + +@dataclass(frozen=True) +class Raise: + location: SourceSpan + exception: Value | None + + +@dataclass(frozen=True) +class ForNext: + """Store the next item of ``iterator`` into local ``target`` and enter ``body``, or ``exit``.""" + + location: SourceSpan + iterator: Value + target: str + body: BlockId + exit: BlockId + + +Terminator: TypeAlias = Return | Branch | Jump | Raise | ForNext + + +@dataclass(frozen=True) +class BasicBlock: + id: BlockId + instructions: tuple[Instruction, ...] + terminator: Terminator @dataclass(frozen=True) class FunctionIR: name: str parameters: tuple[Value, ...] + entry: BlockId blocks: tuple[BasicBlock, ...] location: SourceSpan diff --git a/src/coretrace_python/ir/printer.py b/src/coretrace_python/ir/printer.py index c0d24bd..a41f9f4 100644 --- a/src/coretrace_python/ir/printer.py +++ b/src/coretrace_python/ir/printer.py @@ -2,18 +2,24 @@ from coretrace_python.ir.model import ( BinaryOp, + Branch, Call, Compare, Constant, + ForNext, GetAttr, GetItem, + GetIter, Global, Instruction, + Jump, LoadLocal, ModuleIR, + Raise, Return, StoreLocal, Symbol, + Terminator, UnaryOp, Value, ) @@ -58,23 +64,44 @@ def _instruction(instruction: Instruction) -> str: f"{_value(instruction.result)} = get_item {_value(instruction.object)}, " f"{_value(instruction.key)}" ) + if isinstance(instruction, GetIter): + return f"{_value(instruction.result)} = get_iter {_value(instruction.iterable)}" if isinstance(instruction, LoadLocal): return f'{_value(instruction.result)} = load_local "{instruction.name}"' if isinstance(instruction, StoreLocal): return f'store_local "{instruction.name}", {_value(instruction.value)}' - if isinstance(instruction, Return): - return "return" if instruction.value is None else f"return {_value(instruction.value)}" raise TypeError(f"unknown instruction: {instruction!r}") +def _terminator(terminator: Terminator) -> str: + if isinstance(terminator, Return): + return "return" if terminator.value is None else f"return {_value(terminator.value)}" + if isinstance(terminator, Branch): + return ( + f"branch {_value(terminator.condition)}, " + f"{terminator.then_block}, {terminator.else_block}" + ) + if isinstance(terminator, Jump): + return f"jump {terminator.target}" + if isinstance(terminator, Raise): + return "raise" if terminator.exception is None else f"raise {_value(terminator.exception)}" + if isinstance(terminator, ForNext): + return ( + f'for_next {_value(terminator.iterator)} -> "{terminator.target}", ' + f"{terminator.body}, {terminator.exit}" + ) + raise TypeError(f"unknown terminator: {terminator!r}") + + def format_module(module: ModuleIR) -> str: functions: list[str] = [] for function in module.functions: parameters = ", ".join(_value(parameter) for parameter in function.parameters) lines = [f"func @{function.name}({parameters}) {{"] for block in function.blocks: - lines.append(f"{block.name}:") + lines.append(f"{block.id}:") lines.extend(f" {_instruction(instruction)}" for instruction in block.instructions) + lines.append(f" {_terminator(block.terminator)}") lines.append("}") functions.append("\n".join(lines)) return "\n\n".join(functions)