coretrace-python-analyzer 0.1.0__py3-none-any.whl

coretrace_python/ir/model.py

@@ -0,0 +1,481 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, fields, replace
+from typing import Any, TypeAlias, TypeVar
+
+from coretrace_python.cfg import BlockId
+from coretrace_python.semantic.symbols import SymbolId
+from coretrace_python.source import SourceSpan
+
+# Bumped whenever the instruction set or its meaning changes; part of the cache key.
+PYIR_SCHEMA_VERSION = 1
+
+
+@dataclass(frozen=True)
+class Value:
+    id: int
+
+
+class Operands:
+    """Generic operand access for instructions and terminators."""
+
+    def operands(self) -> tuple[Value, ...]:
+        found: list[Value] = []
+        for field_info in fields(self):  # type: ignore[arg-type]
+            if field_info.name != "result":
+                found.extend(_values_in(getattr(self, field_info.name)))
+        return tuple(found)
+
+
+def _values_in(value: object) -> list[Value]:
+    if isinstance(value, Value):
+        return [value]
+    if isinstance(value, tuple):
+        return [v for item in value for v in _values_in(item)]
+    return []
+
+
+def _map_values(value: Any, mapping: Callable[[Value], Value]) -> Any:
+    if isinstance(value, Value):
+        return mapping(value)
+    if isinstance(value, tuple):
+        return tuple(_map_values(item, mapping) for item in value)
+    return value
+
+
+N = TypeVar("N", bound=Operands)
+
+
+def substitute(node: N, mapping: Callable[[Value], Value], *, include_result: bool = False) -> N:
+    """Copy ``node`` with every operand (and optionally its result) passed through ``mapping``."""
+
+    changes: dict[str, Any] = {}
+    for field_info in fields(node):  # type: ignore[arg-type]
+        value = getattr(node, field_info.name)
+        if field_info.name == "result":
+            if include_result and isinstance(value, Value):
+                changes["result"] = mapping(value)
+        else:
+            changes[field_info.name] = _map_values(value, mapping)
+    return replace(node, **changes)  # type: ignore[type-var]
+
+
+@dataclass(frozen=True)
+class Constant(Operands):
+    result: Value
+    location: SourceSpan
+    value: object
+
+
+@dataclass(frozen=True)
+class Global(Operands):
+    result: Value
+    location: SourceSpan
+    name: str
+
+
+@dataclass(frozen=True)
+class Symbol(Operands):
+    result: Value
+    location: SourceSpan
+    symbol_id: SymbolId
+
+
+@dataclass(frozen=True)
+class BinaryOp(Operands):
+    result: Value
+    location: SourceSpan
+    operator: str
+    left: Value
+    right: Value
+
+
+@dataclass(frozen=True)
+class UnaryOp(Operands):
+    result: Value
+    location: SourceSpan
+    operator: str
+    operand: Value
+
+
+@dataclass(frozen=True)
+class Compare(Operands):
+    result: Value
+    location: SourceSpan
+    operator: str
+    left: Value
+    right: Value
+
+
+@dataclass(frozen=True)
+class Call(Operands):
+    """``starred`` are ``*iterable`` arguments whose positions are unknown; keywords
+    with a ``None`` name are ``**mapping`` arguments."""
+
+    result: Value
+    location: SourceSpan
+    callee: Value
+    arguments: tuple[Value, ...]
+    keywords: tuple[tuple[str | None, Value], ...] = ()
+    starred: tuple[Value, ...] = ()
+
+    def argument_values(self) -> tuple[Value, ...]:
+        return (*self.arguments, *self.starred, *(value for _, value in self.keywords))
+
+
+@dataclass(frozen=True)
+class BoolOp(Operands):
+    result: Value
+    location: SourceSpan
+    operator: str
+    values: tuple[Value, ...]
+
+
+@dataclass(frozen=True)
+class BuildList(Operands):
+    result: Value
+    location: SourceSpan
+    elements: tuple[Value, ...]
+    unpacked: tuple[Value, ...] = ()
+
+
+@dataclass(frozen=True)
+class BuildTuple(Operands):
+    result: Value
+    location: SourceSpan
+    elements: tuple[Value, ...]
+    unpacked: tuple[Value, ...] = ()
+
+
+@dataclass(frozen=True)
+class BuildDict(Operands):
+    result: Value
+    location: SourceSpan
+    items: tuple[tuple[Value, Value], ...]
+    unpacked: tuple[Value, ...] = ()
+
+
+@dataclass(frozen=True)
+class BuildSet(Operands):
+    result: Value
+    location: SourceSpan
+    elements: tuple[Value, ...]
+    unpacked: tuple[Value, ...] = ()
+
+
+@dataclass(frozen=True)
+class MakeFunction(Operands):
+    """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)
+class MakeClass(Operands):
+    """A class defined inside a function, bound to its name; its methods are analysed
+    as nested functions."""
+
+    result: Value
+    location: SourceSpan
+    name: str
+
+
+@dataclass(frozen=True)
+class BuildString(Operands):
+    """An f-string: its constant and formatted parts, in order."""
+
+    result: Value
+    location: SourceSpan
+    parts: tuple[Value, ...]
+
+
+@dataclass(frozen=True)
+class BuildSlice(Operands):
+    result: Value
+    location: SourceSpan
+    lower: Value | None
+    upper: Value | None
+    step: Value | None
+
+
+@dataclass(frozen=True)
+class WithEnter(Operands):
+    result: Value
+    location: SourceSpan
+    context: Value
+
+
+@dataclass(frozen=True)
+class Catch(Operands):
+    """The exception bound by ``except type as name`` at the start of a handler."""
+
+    result: Value
+    location: SourceSpan
+    type: Value | None
+
+
+@dataclass(frozen=True)
+class Await(Operands):
+    result: Value
+    location: SourceSpan
+    value: Value
+
+
+@dataclass(frozen=True)
+class Yield(Operands):
+    result: Value
+    location: SourceSpan
+    value: Value | None
+
+
+@dataclass(frozen=True)
+class GetAttr(Operands):
+    result: Value
+    location: SourceSpan
+    object: Value
+    attribute: str
+
+
+@dataclass(frozen=True)
+class GetItem(Operands):
+    result: Value
+    location: SourceSpan
+    object: Value
+    key: Value
+
+
+@dataclass(frozen=True)
+class GetIter(Operands):
+    result: Value
+    location: SourceSpan
+    iterable: Value
+
+
+@dataclass(frozen=True)
+class LoadLocal(Operands):
+    result: Value
+    location: SourceSpan
+    name: str
+
+
+@dataclass(frozen=True)
+class Phi(Operands):
+    """Merge of the values of local ``name`` arriving from each predecessor (SSA only)."""
+
+    result: Value
+    location: SourceSpan
+    name: str
+    incoming: tuple[tuple[Value, BlockId], ...]
+
+
+@dataclass(frozen=True)
+class Undefined(Operands):
+    """The value of local ``name`` on a path where it was never assigned (SSA only)."""
+
+    result: Value
+    location: SourceSpan
+    name: str
+
+
+@dataclass(frozen=True)
+class StoreLocal(Operands):
+    result: None
+    location: SourceSpan
+    name: str
+    value: Value
+
+
+@dataclass(frozen=True)
+class SetAttr(Operands):
+    result: None
+    location: SourceSpan
+    object: Value
+    attribute: str
+    value: Value
+
+
+@dataclass(frozen=True)
+class SetItem(Operands):
+    result: None
+    location: SourceSpan
+    object: Value
+    key: Value
+    value: Value
+
+
+@dataclass(frozen=True)
+class WithExit(Operands):
+    result: None
+    location: SourceSpan
+    context: Value
+
+
+@dataclass(frozen=True)
+class Assert(Operands):
+    result: None
+    location: SourceSpan
+    test: Value
+    message: Value | None
+
+
+@dataclass(frozen=True)
+class DelItem(Operands):
+    result: None
+    location: SourceSpan
+    object: Value
+    key: Value
+
+
+@dataclass(frozen=True)
+class DelAttr(Operands):
+    result: None
+    location: SourceSpan
+    object: Value
+    attribute: str
+
+
+@dataclass(frozen=True)
+class SetGlobal(Operands):
+    """Assignment to a name declared ``global``."""
+
+    result: None
+    location: SourceSpan
+    name: str
+    value: Value
+
+
+@dataclass(frozen=True)
+class Import(Operands):
+    """An import executed where it stands: ``module`` as written (relative dots
+    included), the canonical ``symbol_id`` bound and the local ``name`` it binds."""
+
+    result: None
+    location: SourceSpan
+    module: str
+    symbol_id: SymbolId
+    name: str
+
+
+ValueInstruction: TypeAlias = (
+    Constant
+    | Global
+    | Symbol
+    | BinaryOp
+    | UnaryOp
+    | Compare
+    | Call
+    | GetAttr
+    | GetItem
+    | GetIter
+    | LoadLocal
+    | Phi
+    | Undefined
+    | BoolOp
+    | BuildList
+    | BuildTuple
+    | BuildDict
+    | BuildSet
+    | BuildString
+    | BuildSlice
+    | MakeFunction
+    | MakeClass
+    | WithEnter
+    | Catch
+    | Await
+    | Yield
+)
+EffectInstruction: TypeAlias = (
+    StoreLocal | SetAttr | SetItem | WithExit | Assert | Import | SetGlobal | DelItem | DelAttr
+)
+Instruction: TypeAlias = ValueInstruction | EffectInstruction
+
+
+# --------------------------------------------------------------------------- terminators
+
+
+@dataclass(frozen=True)
+class Return(Operands):
+    location: SourceSpan
+    value: Value | None
+
+
+@dataclass(frozen=True)
+class Branch(Operands):
+    location: SourceSpan
+    condition: Value
+    then_block: BlockId
+    else_block: BlockId
+
+
+@dataclass(frozen=True)
+class Jump(Operands):
+    location: SourceSpan
+    target: BlockId
+
+
+@dataclass(frozen=True)
+class Raise(Operands):
+    location: SourceSpan
+    exception: Value | None
+    cause: Value | None = None
+
+
+@dataclass(frozen=True)
+class ForNext(Operands):
+    """Store the next item of ``iterator`` into local ``target`` and enter ``body``, or ``exit``."""
+
+    location: SourceSpan
+    iterator: Value
+    target: str
+    body: BlockId
+    exit: BlockId
+    result: Value | None = None
+
+
+Terminator: TypeAlias = Return | Branch | Jump | Raise | ForNext
+
+
+@dataclass(frozen=True)
+class BasicBlock:
+    id: BlockId
+    instructions: tuple[Instruction, ...]
+    terminator: Terminator
+    exception_targets: tuple[BlockId, ...] = ()
+
+
+def successors(block: BasicBlock) -> tuple[BlockId, ...]:
+    """Terminator targets followed by the exception edges of ``block``."""
+
+    found = list(_terminator_targets(block.terminator))
+    found.extend(t for t in block.exception_targets if t not in found)
+    return tuple(found)
+
+
+def _terminator_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, ForNext):
+        return (terminator.body, terminator.exit)
+    return ()
+
+
+@dataclass(frozen=True)
+class FunctionIR:
+    name: str
+    parameters: tuple[Value, ...]
+    entry: BlockId
+    blocks: tuple[BasicBlock, ...]
+    location: SourceSpan
+
+
+@dataclass(frozen=True)
+class ModuleIR:
+    functions: tuple[FunctionIR, ...]

coretrace_python/ir/printer.py

@@ -0,0 +1,201 @@
+from __future__ import annotations
+
+from coretrace_python.ir.model import (
+    Assert,
+    Await,
+    BinaryOp,
+    BoolOp,
+    Branch,
+    BuildDict,
+    BuildList,
+    BuildSet,
+    BuildSlice,
+    BuildString,
+    BuildTuple,
+    Call,
+    Catch,
+    Compare,
+    Constant,
+    DelAttr,
+    DelItem,
+    ForNext,
+    GetAttr,
+    GetItem,
+    GetIter,
+    Global,
+    Import,
+    Instruction,
+    Jump,
+    LoadLocal,
+    MakeClass,
+    MakeFunction,
+    ModuleIR,
+    Phi,
+    Raise,
+    Return,
+    SetAttr,
+    SetGlobal,
+    SetItem,
+    StoreLocal,
+    Symbol,
+    Terminator,
+    UnaryOp,
+    Undefined,
+    Value,
+    WithEnter,
+    WithExit,
+    Yield,
+)
+
+
+def _value(value: Value) -> str:
+    return f"%{value.id}"
+
+
+def _instruction(instruction: Instruction) -> str:
+    if isinstance(instruction, Constant):
+        return f"{_value(instruction.result)} = const {instruction.value!r}"
+    if isinstance(instruction, Global):
+        return f"{_value(instruction.result)} = global {instruction.name!r}"
+    if isinstance(instruction, Symbol):
+        return f"{_value(instruction.result)} = symbol @{instruction.symbol_id}"
+    if isinstance(instruction, BinaryOp):
+        return (
+            f"{_value(instruction.result)} = binary.{instruction.operator} "
+            f"{_value(instruction.left)}, {_value(instruction.right)}"
+        )
+    if isinstance(instruction, UnaryOp):
+        return (
+            f"{_value(instruction.result)} = unary.{instruction.operator} "
+            f"{_value(instruction.operand)}"
+        )
+    if isinstance(instruction, Compare):
+        return (
+            f"{_value(instruction.result)} = compare.{instruction.operator} "
+            f"{_value(instruction.left)}, {_value(instruction.right)}"
+        )
+    if isinstance(instruction, Call):
+        parts = [_value(argument) for argument in instruction.arguments]
+        parts.extend(f"*{_value(value)}" for value in instruction.starred)
+        for name, value in instruction.keywords:
+            parts.append(f"{name}={_value(value)}" if name is not None else f"**{_value(value)}")
+        return f"{_value(instruction.result)} = call {_value(instruction.callee)}({', '.join(parts)})"
+    if isinstance(instruction, BoolOp):
+        values = ", ".join(_value(v) for v in instruction.values)
+        return f"{_value(instruction.result)} = bool_op.{instruction.operator} {values}"
+    if isinstance(instruction, BuildList | BuildTuple):
+        kind = "build_list" if isinstance(instruction, BuildList) else "build_tuple"
+        parts = [_value(v) for v in instruction.elements] + [f"*{_value(v)}" for v in instruction.unpacked]
+        return f"{_value(instruction.result)} = {kind} {', '.join(parts)}".rstrip()
+    if isinstance(instruction, BuildDict):
+        parts = [f"{_value(k)}: {_value(v)}" for k, v in instruction.items]
+        parts.extend(f"**{_value(v)}" for v in instruction.unpacked)
+        return f"{_value(instruction.result)} = build_dict {', '.join(parts)}".rstrip()
+    if isinstance(instruction, BuildSet):
+        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):
+        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):
+        return f'del_attr {_value(instruction.object)}, "{instruction.attribute}"'
+    if isinstance(instruction, MakeClass):
+        return f'{_value(instruction.result)} = make_class "{instruction.name}"'
+    if isinstance(instruction, SetGlobal):
+        return f'set_global "{instruction.name}", {_value(instruction.value)}'
+    if isinstance(instruction, BuildString):
+        return f"{_value(instruction.result)} = build_string {', '.join(_value(v) for v in instruction.parts)}".rstrip()
+    if isinstance(instruction, BuildSlice):
+        bounds = ", ".join("none" if b is None else _value(b) for b in (instruction.lower, instruction.upper, instruction.step))
+        return f"{_value(instruction.result)} = build_slice {bounds}"
+    if isinstance(instruction, WithEnter):
+        return f"{_value(instruction.result)} = with_enter {_value(instruction.context)}"
+    if isinstance(instruction, WithExit):
+        return f"with_exit {_value(instruction.context)}"
+    if isinstance(instruction, Catch):
+        suffix = "" if instruction.type is None else f" {_value(instruction.type)}"
+        return f"{_value(instruction.result)} = catch{suffix}"
+    if isinstance(instruction, Await):
+        return f"{_value(instruction.result)} = await {_value(instruction.value)}"
+    if isinstance(instruction, Yield):
+        suffix = "" if instruction.value is None else f" {_value(instruction.value)}"
+        return f"{_value(instruction.result)} = yield{suffix}"
+    if isinstance(instruction, SetAttr):
+        return f"set_attr {_value(instruction.object)}, {instruction.attribute!r}, {_value(instruction.value)}"
+    if isinstance(instruction, SetItem):
+        return (
+            f"set_item {_value(instruction.object)}, {_value(instruction.key)}, "
+            f"{_value(instruction.value)}"
+        )
+    if isinstance(instruction, Import):
+        return f'import "{instruction.module}" as "{instruction.name}" @{instruction.symbol_id}'
+    if isinstance(instruction, Assert):
+        message = "" if instruction.message is None else f", {_value(instruction.message)}"
+        return f"assert {_value(instruction.test)}{message}"
+    if isinstance(instruction, GetAttr):
+        return (
+            f"{_value(instruction.result)} = get_attr {_value(instruction.object)}, "
+            f"{instruction.attribute!r}"
+        )
+    if isinstance(instruction, GetItem):
+        return (
+            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, Phi):
+        incoming = ", ".join(f"[{_value(v)}, {b}]" for v, b in instruction.incoming)
+        return f'{_value(instruction.result)} = phi "{instruction.name}", {incoming}'
+    if isinstance(instruction, Undefined):
+        return f'{_value(instruction.result)} = undefined "{instruction.name}"'
+    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):
+        if terminator.exception is None:
+            return "raise"
+        text = f"raise {_value(terminator.exception)}"
+        return text if terminator.cause is None else f"{text} from {_value(terminator.cause)}"
+    if isinstance(terminator, ForNext):
+        prefix = "" if terminator.result is None else f"{_value(terminator.result)} = "
+        return (
+            f'{prefix}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:
+            header = str(block.id)
+            if block.exception_targets:
+                header += f" [except: {', '.join(str(t) for t in block.exception_targets)}]"
+            lines.append(f"{header}:")
+            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)