diff --git a/src/coretrace_python/frontend/lowering.py b/src/coretrace_python/frontend/lowering.py index b2be307..155384d 100644 --- a/src/coretrace_python/frontend/lowering.py +++ b/src/coretrace_python/frontend/lowering.py @@ -13,9 +13,11 @@ GetAttr, GetItem, Global, + LoadLocal, ModuleIR, Return, SourceLocation, + StoreLocal, UnaryOp, Value, ) @@ -61,6 +63,7 @@ class _FunctionLowerer: next_value_id: int = 0 locals: dict[str, Value] = field(default_factory=dict) block: BasicBlock = field(default_factory=lambda: BasicBlock("entry")) + parameters: dict[str, Value] = field(default_factory=dict) def location(self, node: ast.AST) -> SourceLocation: return SourceLocation( @@ -91,7 +94,9 @@ def expression(self, node: ast.expr) -> Value: location = self.location(node) if isinstance(node, ast.Name): if node.id in self.locals: - return self.locals[node.id] + return self.emit(LoadLocal(self.new_value(), location, node.id)) + if node.id in self.parameters: + return self.parameters[node.id] return self.emit(Global(self.new_value(), location, node.id)) if isinstance(node, ast.Constant): return self.emit(Constant(self.new_value(), location, node.value)) @@ -152,7 +157,10 @@ def statement(self, node: ast.stmt) -> None: if isinstance(node, ast.Assign): if len(node.targets) != 1 or not isinstance(node.targets[0], ast.Name): self.fail(node, "only assignment to one local name is supported") - self.locals[node.targets[0].id] = self.expression(node.value) + name = node.targets[0].id + value = self.expression(node.value) + self.emit(StoreLocal(None, self.location(node), name, value)) + self.locals[name] = value return if isinstance(node, ast.Return): value = self.expression(node.value) if node.value is not None else None @@ -175,7 +183,7 @@ def function(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> FunctionIR: self.fail(node.args, "only positional arguments are supported") parameters = tuple(self.new_value() for _ in node.args.args) argument_names = (argument.arg for argument in node.args.args) - self.locals.update(zip(argument_names, parameters, strict=True)) + self.parameters.update(zip(argument_names, parameters, strict=True)) for statement in node.body: self.statement(statement) return FunctionIR(node.name, parameters, (self.block,), self.location(node)) diff --git a/src/coretrace_python/ir/model.py b/src/coretrace_python/ir/model.py index 87213fa..7ccaa22 100644 --- a/src/coretrace_python/ir/model.py +++ b/src/coretrace_python/ir/model.py @@ -94,8 +94,33 @@ class Return: value: Value | None +@dataclass(frozen=True) +class LoadLocal: + result: Value + location: SourceLocation + name: str + + +@dataclass(frozen=True) +class StoreLocal: + result: None + location: SourceLocation + name: str + value: Value + + Instruction: TypeAlias = ( - Constant | Global | BinaryOp | UnaryOp | Compare | Call | GetAttr | GetItem | Return + Constant + | Global + | BinaryOp + | UnaryOp + | Compare + | Call + | GetAttr + | GetItem + | LoadLocal + | StoreLocal + | Return ) diff --git a/src/coretrace_python/ir/printer.py b/src/coretrace_python/ir/printer.py index cfda6c2..db5a385 100644 --- a/src/coretrace_python/ir/printer.py +++ b/src/coretrace_python/ir/printer.py @@ -9,8 +9,10 @@ GetItem, Global, Instruction, + LoadLocal, ModuleIR, Return, + StoreLocal, UnaryOp, Value, ) @@ -53,6 +55,10 @@ def _instruction(instruction: Instruction) -> str: f"{_value(instruction.result)} = get_item {_value(instruction.object)}, " f"{_value(instruction.key)}" ) + 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}") diff --git a/tests/test_emit_ir.py b/tests/test_emit_ir.py index a0ff65b..3fb9776 100644 --- a/tests/test_emit_ir.py +++ b/tests/test_emit_ir.py @@ -10,7 +10,9 @@ def test_emit_ir_golden(tmp_path, capsys) -> None: "func @add(%0, %1) {\n" "entry:\n" " %2 = binary.add %0, %1\n" - " return %2\n" + ' store_local "x", %2\n' + ' %3 = load_local "x"\n' + " return %3\n" "}\n" ) @@ -41,3 +43,24 @@ def test_emit_ir_for_call(tmp_path, capsys) -> None: output = capsys.readouterr().out assert "%1 = global 'print'" in output assert "%2 = call %1(%0)" in output + + +def test_emit_ir_for_local_assignment(tmp_path, capsys) -> None: + source = tmp_path / "locals.py" + source.write_text( + "def calculate(a, b):\n" + " result = a + b\n" + " return result\n", + encoding="utf-8", + ) + + assert main(["--emit-ir", str(source)]) == 0 + assert capsys.readouterr().out == ( + "func @calculate(%0, %1) {\n" + "entry:\n" + " %2 = binary.add %0, %1\n" + ' store_local "result", %2\n' + ' %3 = load_local "result"\n' + " return %3\n" + "}\n" + )