pyir 0.1.0__py3-none-any.whl

pyir/__init__.py

@@ -0,0 +1,78 @@
+"""pyir - Python Intermediate Representation.
+
+Lower a small, statically-tractable subset of Python into a three-address-code
+IR organised in basic blocks, print it, optimise it and interpret it.
+
+Typical use::
+
+    import pyir
+
+    module = pyir.compile_source(source)
+    print(pyir.format_module(module))
+    pyir.optimize(module)
+    result = pyir.interpret(module, "gcd", 48, 18)
+"""
+
+from __future__ import annotations
+
+from .errors import IRRuntimeError, IRVerificationError, PyIRError, UnsupportedSyntaxError
+from .interp import Interpreter, interpret
+from .ir import (
+    BUILTINS,
+    BasicBlock,
+    BinOp,
+    Branch,
+    Call,
+    Const,
+    Copy,
+    Function,
+    Jump,
+    Module,
+    Return,
+    UnaryOp,
+    Var,
+)
+from .lower import compile_function, compile_source
+from .passes import DEFAULT_PASSES, constant_fold, eliminate_dead_code, optimize, simplify_cfg
+from .printer import format_block, format_function, format_module
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "__version__",
+    # front end
+    "compile_source",
+    "compile_function",
+    # IR
+    "Module",
+    "Function",
+    "BasicBlock",
+    "Const",
+    "Var",
+    "Copy",
+    "BinOp",
+    "UnaryOp",
+    "Call",
+    "Jump",
+    "Branch",
+    "Return",
+    "BUILTINS",
+    # printing
+    "format_module",
+    "format_function",
+    "format_block",
+    # optimisation
+    "optimize",
+    "constant_fold",
+    "eliminate_dead_code",
+    "simplify_cfg",
+    "DEFAULT_PASSES",
+    # execution
+    "interpret",
+    "Interpreter",
+    # errors
+    "PyIRError",
+    "UnsupportedSyntaxError",
+    "IRVerificationError",
+    "IRRuntimeError",
+]

pyir/__main__.py

@@ -0,0 +1,47 @@
+"""Command line interface: ``python -m pyir FILE [-O] [--run FUNC ARG...]``."""
+
+from __future__ import annotations
+
+import argparse
+import ast
+import sys
+from typing import Sequence
+
+from .errors import PyIRError
+from .interp import interpret
+from .lower import compile_source
+from .passes import optimize
+from .printer import format_module
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    """Entry point for ``python -m pyir`` and the ``pyir`` console script."""
+    parser = argparse.ArgumentParser(prog="pyir", description="Lower Python functions to pyir IR.")
+    parser.add_argument("file", help="Python source file containing only function definitions ('-' for stdin)")
+    parser.add_argument("-O", "--optimize", action="store_true", help="run the default optimisation passes")
+    parser.add_argument(
+        "--run",
+        nargs="+",
+        metavar=("FUNC", "ARG"),
+        help="interpret FUNC with literal arguments (e.g. --run gcd 48 18) instead of printing IR",
+    )
+    args = parser.parse_args(argv)
+
+    source = sys.stdin.read() if args.file == "-" else open(args.file, encoding="utf-8").read()
+    try:
+        module = compile_source(source)
+        if args.optimize:
+            optimize(module)
+        if args.run:
+            func, *raw = args.run
+            print(repr(interpret(module, func, *(ast.literal_eval(a) for a in raw))))
+        else:
+            print(format_module(module))
+    except (PyIRError, SyntaxError) as exc:
+        print(f"pyir: error: {exc}", file=sys.stderr)
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

pyir/errors.py

@@ -0,0 +1,38 @@
+"""Exception hierarchy for :mod:`pyir`."""
+
+from __future__ import annotations
+
+__all__ = [
+    "PyIRError",
+    "UnsupportedSyntaxError",
+    "IRVerificationError",
+    "IRRuntimeError",
+]
+
+
+class PyIRError(Exception):
+    """Base class for every error raised by pyir."""
+
+
+class UnsupportedSyntaxError(PyIRError):
+    """Raised when the lowering step meets Python syntax outside the supported subset."""
+
+    def __init__(self, message: str, lineno: int | None = None) -> None:
+        self.lineno = lineno
+        if lineno is not None:
+            message = f"line {lineno}: {message}"
+        super().__init__(message)
+
+
+class IRVerificationError(PyIRError):
+    """Raised by :meth:`pyir.ir.Function.verify` when a function is malformed."""
+
+
+class IRRuntimeError(PyIRError):
+    """Raised by the IR interpreter for IR-level failures.
+
+    Examples are reading an unassigned variable, calling an unknown function or
+    exceeding the step budget. Exceptions raised by the operations themselves
+    (``ZeroDivisionError``, ``TypeError``, ...) propagate unchanged so they match
+    what the original Python code would raise.
+    """

pyir/interp.py

@@ -0,0 +1,110 @@
+"""A straightforward interpreter for pyir IR, used to check semantics."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .errors import IRRuntimeError
+from .ir import (
+    BINARY_OPS,
+    BUILTINS,
+    UNARY_OPS,
+    BinOp,
+    Branch,
+    Call,
+    Const,
+    Copy,
+    Function,
+    Jump,
+    Module,
+    Operand,
+    Return,
+    UnaryOp,
+)
+
+__all__ = ["Interpreter", "interpret", "DEFAULT_MAX_STEPS"]
+
+DEFAULT_MAX_STEPS = 10_000_000
+"""Default budget of executed instructions (terminators included) per :meth:`Interpreter.run`."""
+
+
+class Interpreter:
+    """Executes functions of a :class:`~pyir.ir.Module`.
+
+    Attributes:
+        module: The module whose functions can be run and called.
+        max_steps: Maximum number of instructions (including terminators) a
+            single :meth:`run` may execute, across nested calls.
+        steps: Number of instructions executed by the most recent :meth:`run`.
+    """
+
+    def __init__(self, module: Module, max_steps: int = DEFAULT_MAX_STEPS) -> None:
+        self.module = module
+        self.max_steps = max_steps
+        self.steps = 0
+
+    def run(self, name: str, *args: Any) -> Any:
+        """Call function ``name`` with ``args`` and return its result.
+
+        Raises:
+            IRRuntimeError: on unknown functions, unassigned variables or when
+                the step budget is exhausted.
+            TypeError: when the argument count does not match the parameters.
+            Exception: anything the executed operations raise
+                (e.g. ``ZeroDivisionError``) propagates unchanged.
+        """
+        self.steps = 0
+        return self._call(name, list(args))
+
+    def _call(self, name: str, args: list[Any]) -> Any:
+        fn = self.module.functions.get(name)
+        if fn is not None:
+            return self._execute(fn, args)
+        builtin = BUILTINS.get(name)
+        if builtin is not None:
+            return builtin(*args)
+        raise IRRuntimeError(f"call to unknown function {name!r}")
+
+    def _execute(self, fn: Function, args: list[Any]) -> Any:
+        if len(args) != len(fn.params):
+            raise TypeError(f"{fn.name}() takes {len(fn.params)} positional arguments but {len(args)} were given")
+        env: dict[str, Any] = dict(zip(fn.params, args))
+
+        def read(op: Operand) -> Any:
+            if isinstance(op, Const):
+                return op.value
+            try:
+                return env[op.name]
+            except KeyError:
+                raise IRRuntimeError(f"{fn.name}: variable {op.name!r} read before assignment") from None
+
+        block = fn.entry
+        while True:
+            self.steps += len(block.instructions) + 1
+            if self.steps > self.max_steps:
+                raise IRRuntimeError(f"step budget of {self.max_steps} exhausted in {fn.name}")
+            for inst in block.instructions:
+                if isinstance(inst, Copy):
+                    env[inst.dest.name] = read(inst.src)
+                elif isinstance(inst, BinOp):
+                    env[inst.dest.name] = BINARY_OPS[inst.op](read(inst.left), read(inst.right))
+                elif isinstance(inst, UnaryOp):
+                    env[inst.dest.name] = UNARY_OPS[inst.op](read(inst.operand))
+                elif isinstance(inst, Call):
+                    env[inst.dest.name] = self._call(inst.func, [read(a) for a in inst.args])
+                else:  # pragma: no cover - guarded by the Instruction union
+                    raise IRRuntimeError(f"unknown instruction {inst!r}")
+            term = block.terminator
+            if isinstance(term, Return):
+                return read(term.value)
+            if isinstance(term, Jump):
+                block = fn.blocks[term.target]
+            elif isinstance(term, Branch):
+                block = fn.blocks[term.if_true if read(term.cond) else term.if_false]
+            else:
+                raise IRRuntimeError(f"{fn.name}: block {block.label!r} has no terminator")
+
+
+def interpret(module: Module, name: str, *args: Any, max_steps: int = DEFAULT_MAX_STEPS) -> Any:
+    """Convenience wrapper: ``Interpreter(module, max_steps).run(name, *args)``."""
+    return Interpreter(module, max_steps).run(name, *args)

pyir/ir.py

@@ -0,0 +1,412 @@
+"""Core IR data structures.
+
+The IR is a classic *three-address code* organised into basic blocks:
+
+* A :class:`Module` holds named :class:`Function` objects.
+* A :class:`Function` holds an ordered mapping of labels to :class:`BasicBlock`.
+  The first block is the entry block.
+* A :class:`BasicBlock` is a list of straight-line instructions followed by
+  exactly one terminator (:class:`Jump`, :class:`Branch` or :class:`Return`).
+* Operands are either a :class:`Const` or a :class:`Var`. Variables named
+  after Python locals keep their source name; compiler temporaries are
+  spelled ``%tN``.
+
+Variables may be assigned more than once (the IR is *not* SSA), which keeps
+lowering and interpretation simple; the optimisation passes use classic
+dataflow analyses that do not rely on single assignment.
+"""
+
+from __future__ import annotations
+
+import operator
+from dataclasses import dataclass, field
+from typing import Any, Callable, Iterator, Union
+
+from .errors import IRVerificationError
+
+__all__ = [
+    "Const",
+    "Var",
+    "Operand",
+    "BINARY_OPS",
+    "UNARY_OPS",
+    "BUILTINS",
+    "Copy",
+    "BinOp",
+    "UnaryOp",
+    "Call",
+    "Instruction",
+    "Jump",
+    "Branch",
+    "Return",
+    "Terminator",
+    "BasicBlock",
+    "Function",
+    "Module",
+]
+
+
+# --------------------------------------------------------------------------
+# Operands
+# --------------------------------------------------------------------------
+
+
+class Const:
+    """A literal constant operand (``int``, ``float``, ``bool``, ``str`` or ``None``).
+
+    Equality is type-aware so that ``Const(1)``, ``Const(True)`` and
+    ``Const(1.0)`` are all distinct, and ``Const(0.0) != Const(-0.0)``.
+    """
+
+    __slots__ = ("value",)
+
+    def __init__(self, value: Any) -> None:
+        self.value = value
+
+    def _key(self) -> tuple[type, str]:
+        return (type(self.value), repr(self.value))
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, Const) and self._key() == other._key()
+
+    def __hash__(self) -> int:
+        return hash(self._key())
+
+    def __repr__(self) -> str:
+        return f"Const({self.value!r})"
+
+    def __str__(self) -> str:
+        return repr(self.value)
+
+
+@dataclass(frozen=True)
+class Var:
+    """A named variable: a Python local (``x``) or a compiler temporary (``%t3``)."""
+
+    name: str
+
+    @property
+    def is_temp(self) -> bool:
+        """``True`` for compiler-generated temporaries."""
+        return self.name.startswith("%")
+
+    def __str__(self) -> str:
+        return self.name
+
+
+Operand = Union[Const, Var]
+
+
+# --------------------------------------------------------------------------
+# Operator tables (shared by the interpreter and constant folding)
+# --------------------------------------------------------------------------
+
+BINARY_OPS: dict[str, Callable[[Any, Any], Any]] = {
+    "+": operator.add,
+    "-": operator.sub,
+    "*": operator.mul,
+    "/": operator.truediv,
+    "//": operator.floordiv,
+    "%": operator.mod,
+    "**": operator.pow,
+    "<<": operator.lshift,
+    ">>": operator.rshift,
+    "&": operator.and_,
+    "|": operator.or_,
+    "^": operator.xor,
+    "==": operator.eq,
+    "!=": operator.ne,
+    "<": operator.lt,
+    "<=": operator.le,
+    ">": operator.gt,
+    ">=": operator.ge,
+    "is": operator.is_,
+    "is not": operator.is_not,
+}
+"""Binary operator spelling -> Python implementation."""
+
+UNARY_OPS: dict[str, Callable[[Any], Any]] = {
+    "-": operator.neg,
+    "+": operator.pos,
+    "~": operator.invert,
+    "not": operator.not_,
+}
+"""Unary operator spelling -> Python implementation."""
+
+BUILTINS: dict[str, Callable[..., Any]] = {
+    "abs": abs,
+    "min": min,
+    "max": max,
+    "int": int,
+    "float": float,
+    "bool": bool,
+    "str": str,
+    "round": round,
+    "print": print,
+}
+"""Python builtins callable from IR code (by name, via :class:`Call`)."""
+
+
+# --------------------------------------------------------------------------
+# Instructions
+# --------------------------------------------------------------------------
+
+
+@dataclass
+class Copy:
+    """``dest = src``"""
+
+    dest: Var
+    src: Operand
+
+    def uses(self) -> list[Operand]:
+        """Operands read by this instruction."""
+        return [self.src]
+
+    def __str__(self) -> str:
+        return f"{self.dest} = {self.src}"
+
+
+@dataclass
+class BinOp:
+    """``dest = left <op> right`` for any operator in :data:`BINARY_OPS`."""
+
+    dest: Var
+    op: str
+    left: Operand
+    right: Operand
+
+    def uses(self) -> list[Operand]:
+        """Operands read by this instruction."""
+        return [self.left, self.right]
+
+    def __str__(self) -> str:
+        return f"{self.dest} = {self.left} {self.op} {self.right}"
+
+
+@dataclass
+class UnaryOp:
+    """``dest = <op> operand`` for any operator in :data:`UNARY_OPS`."""
+
+    dest: Var
+    op: str
+    operand: Operand
+
+    def uses(self) -> list[Operand]:
+        """Operands read by this instruction."""
+        return [self.operand]
+
+    def __str__(self) -> str:
+        sep = " " if self.op == "not" else ""
+        return f"{self.dest} = {self.op}{sep}{self.operand}"
+
+
+@dataclass
+class Call:
+    """``dest = call func(args...)``.
+
+    ``func`` names either another function of the same :class:`Module` or an
+    entry of :data:`BUILTINS` (module functions take precedence).
+    """
+
+    dest: Var
+    func: str
+    args: list[Operand]
+
+    def uses(self) -> list[Operand]:
+        """Operands read by this instruction."""
+        return list(self.args)
+
+    def __str__(self) -> str:
+        args = ", ".join(str(a) for a in self.args)
+        return f"{self.dest} = call {self.func}({args})"
+
+
+Instruction = Union[Copy, BinOp, UnaryOp, Call]
+
+
+@dataclass
+class Jump:
+    """Unconditional transfer of control to ``target``."""
+
+    target: str
+
+    def uses(self) -> list[Operand]:
+        """Operands read by this terminator."""
+        return []
+
+    def successors(self) -> list[str]:
+        """Labels this terminator may transfer control to."""
+        return [self.target]
+
+    def __str__(self) -> str:
+        return f"jump {self.target}"
+
+
+@dataclass
+class Branch:
+    """Go to ``if_true`` when ``cond`` is truthy, otherwise to ``if_false``."""
+
+    cond: Operand
+    if_true: str
+    if_false: str
+
+    def uses(self) -> list[Operand]:
+        """Operands read by this terminator."""
+        return [self.cond]
+
+    def successors(self) -> list[str]:
+        """Labels this terminator may transfer control to."""
+        return [self.if_true, self.if_false]
+
+    def __str__(self) -> str:
+        return f"branch {self.cond}, {self.if_true}, {self.if_false}"
+
+
+@dataclass
+class Return:
+    """Return ``value`` from the current function."""
+
+    value: Operand
+
+    def uses(self) -> list[Operand]:
+        """Operands read by this terminator."""
+        return [self.value]
+
+    def successors(self) -> list[str]:
+        """Labels this terminator may transfer control to (none)."""
+        return []
+
+    def __str__(self) -> str:
+        return f"return {self.value}"
+
+
+Terminator = Union[Jump, Branch, Return]
+
+
+# --------------------------------------------------------------------------
+# Containers
+# --------------------------------------------------------------------------
+
+
+@dataclass
+class BasicBlock:
+    """A label, a list of instructions and a single terminator."""
+
+    label: str
+    instructions: list[Instruction] = field(default_factory=list)
+    terminator: Terminator | None = None
+
+    def successors(self) -> list[str]:
+        """Labels of the blocks control may flow to next."""
+        return self.terminator.successors() if self.terminator is not None else []
+
+
+@dataclass
+class Function:
+    """A function: parameter names plus an ordered label -> block mapping."""
+
+    name: str
+    params: list[str]
+    blocks: dict[str, BasicBlock] = field(default_factory=dict)
+
+    @property
+    def entry(self) -> BasicBlock:
+        """The entry block (the first block inserted)."""
+        return next(iter(self.blocks.values()))
+
+    def __iter__(self) -> Iterator[BasicBlock]:
+        return iter(self.blocks.values())
+
+    def predecessors(self) -> dict[str, list[str]]:
+        """Map every label to the labels of its predecessor blocks."""
+        preds: dict[str, list[str]] = {label: [] for label in self.blocks}
+        for block in self.blocks.values():
+            for succ in block.successors():
+                if succ in preds and block.label not in preds[succ]:
+                    preds[succ].append(block.label)
+        return preds
+
+    def reachable(self) -> list[str]:
+        """Labels reachable from the entry block, in reverse post-order."""
+        seen: set[str] = set()
+        order: list[str] = []
+
+        def visit(label: str) -> None:
+            # Iterative DFS to avoid recursion limits on large CFGs.
+            stack: list[tuple[str, Iterator[str]]] = [(label, iter(self.blocks[label].successors()))]
+            seen.add(label)
+            while stack:
+                current, succs = stack[-1]
+                for succ in succs:
+                    if succ not in seen and succ in self.blocks:
+                        seen.add(succ)
+                        stack.append((succ, iter(self.blocks[succ].successors())))
+                        break
+                else:
+                    stack.pop()
+                    order.append(current)
+
+        if self.blocks:
+            visit(self.entry.label)
+        order.reverse()
+        return order
+
+    def instruction_count(self) -> int:
+        """Total number of instructions plus terminators across all blocks."""
+        return sum(len(b.instructions) + (b.terminator is not None) for b in self)
+
+    def verify(self) -> None:
+        """Check structural invariants.
+
+        Raises:
+            IRVerificationError: if the function has no blocks, a block lacks a
+                terminator, a label is inconsistent, a jump targets an unknown
+                block, or an operator is unknown.
+        """
+        if not self.blocks:
+            raise IRVerificationError(f"{self.name}: function has no blocks")
+        for label, block in self.blocks.items():
+            if block.label != label:
+                raise IRVerificationError(f"{self.name}: block key {label!r} != label {block.label!r}")
+            if block.terminator is None:
+                raise IRVerificationError(f"{self.name}: block {label!r} has no terminator")
+            for succ in block.successors():
+                if succ not in self.blocks:
+                    raise IRVerificationError(f"{self.name}: block {label!r} jumps to unknown block {succ!r}")
+            for inst in block.instructions:
+                if isinstance(inst, BinOp) and inst.op not in BINARY_OPS:
+                    raise IRVerificationError(f"{self.name}: unknown binary operator {inst.op!r}")
+                if isinstance(inst, UnaryOp) and inst.op not in UNARY_OPS:
+                    raise IRVerificationError(f"{self.name}: unknown unary operator {inst.op!r}")
+
+    def __str__(self) -> str:
+        from .printer import format_function
+
+        return format_function(self)
+
+
+@dataclass
+class Module:
+    """A collection of functions that may call each other by name."""
+
+    functions: dict[str, Function] = field(default_factory=dict)
+
+    def __getitem__(self, name: str) -> Function:
+        return self.functions[name]
+
+    def __contains__(self, name: object) -> bool:
+        return name in self.functions
+
+    def __iter__(self) -> Iterator[Function]:
+        return iter(self.functions.values())
+
+    def verify(self) -> None:
+        """Verify every function (see :meth:`Function.verify`)."""
+        for fn in self:
+            fn.verify()
+
+    def __str__(self) -> str:
+        from .printer import format_module
+
+        return format_module(self)