feat: data-flow framework and constant propagation

src/coretrace_python/abstract/__init__.py

@@ -0,0 +1,16 @@
+"""Abstract values and the analyses that compute them (architecture §18)."""
+
+from coretrace_python.abstract.constants import (
+    ConstantFacts,
+    ConstantPropagation,
+    propagate_constants,
+)
+from coretrace_python.abstract.values import AbstractValue, Truth
+
+__all__ = [
+    "AbstractValue",
+    "ConstantFacts",
+    "ConstantPropagation",
+    "Truth",
+    "propagate_constants",
+]

src/coretrace_python/abstract/constants.py

@@ -0,0 +1,225 @@
+"""Constant propagation: the reference client of the data-flow solver (§18, §38)."""
+
+from __future__ import annotations
+
+import operator
+from collections.abc import Callable, Mapping
+from types import MappingProxyType
+from typing import Any, ClassVar
+
+from coretrace_python.abstract.values import AbstractValue, Truth
+from coretrace_python.analysis import AnalysisContext, AnyAnalysis, FunctionAnalysis
+from coretrace_python.cfg import CFG, BlockId, CFGAnalysis
+from coretrace_python.dataflow import BOTTOM, TOP, DataflowProblem, Direction, solve
+from coretrace_python.hir import nodes
+from coretrace_python.ir.model import (
+    BasicBlock,
+    BinaryOp,
+    Branch,
+    Compare,
+    Constant,
+    ForNext,
+    FunctionIR,
+    Instruction,
+    Jump,
+    Phi,
+    UnaryOp,
+    Value,
+)
+from coretrace_python.ir.ssa import SSAAnalysis
+
+State = Mapping[Value, AbstractValue]
+
+_SAFE_TYPES = (int, float, bool, str, bytes, type(None))
+_NUMERIC = frozenset({"int", "float", "bool"})
+
+_BINARY: Mapping[str, Callable[..., Any]] = {
+    "add": operator.add,
+    "sub": operator.sub,
+    "mul": operator.mul,
+    "div": operator.truediv,
+    "floor_div": operator.floordiv,
+    "mod": operator.mod,
+    "bit_or": operator.or_,
+    "bit_xor": operator.xor,
+    "bit_and": operator.and_,
+}
+_COMPARE: Mapping[str, Callable[..., Any]] = {
+    "eq": operator.eq,
+    "not_eq": operator.ne,
+    "lt": operator.lt,
+    "lt_eq": operator.le,
+    "gt": operator.gt,
+    "gt_eq": operator.ge,
+    "is": operator.is_,
+    "is_not": operator.is_not,
+    "in": lambda a, b: a in b,
+    "not_in": lambda a, b: a not in b,
+}
+_UNARY: Mapping[str, Callable[..., Any]] = {
+    "neg": operator.neg,
+    "pos": operator.pos,
+    "invert": operator.invert,
+}
+
+
+class ConstantFacts:
+    def __init__(self, values: Mapping[Value, AbstractValue], reachable: frozenset[BlockId]):
+        self._values = MappingProxyType(dict(values))
+        self._reachable = reachable
+
+    def value(self, value: Value) -> AbstractValue:
+        return self._values.get(value, AbstractValue.bottom())
+
+    def reachable(self, block: BlockId) -> bool:
+        return block in self._reachable
+
+
+class _ConstantProblem(DataflowProblem[State]):
+    direction: ClassVar[Direction] = Direction.FORWARD
+
+    def __init__(self, function: FunctionIR) -> None:
+        self.function = function
+        self.blocks = {block.id: block for block in function.blocks}
+
+    def initial(self) -> State:
+        return MappingProxyType({p: AbstractValue.unknown() for p in self.function.parameters})
+
+    def join(self, a: State, b: State) -> State:
+        merged = dict(a)
+        for value, fact in b.items():
+            merged[value] = merged[value].join(fact) if value in merged else fact
+        return MappingProxyType(merged)
+
+    def evaluate(self, block: BasicBlock, incoming: Mapping[BlockId, State]) -> State:
+        """State after ``block`` given the states on its executable incoming edges."""
+
+        states = list(incoming.values())
+        state = dict(states[0])
+        for other in states[1:]:
+            state = dict(self.join(state, other))
+        for instruction in block.instructions:
+            if instruction.result is None:
+                continue
+            state[instruction.result] = (
+                self.phi(instruction, incoming)
+                if isinstance(instruction, Phi)
+                else self.instruction(instruction, state)
+            )
+        if isinstance(block.terminator, ForNext) and block.terminator.result is not None:
+            state[block.terminator.result] = AbstractValue.unknown()
+        return MappingProxyType(state)
+
+    def flow(self, cfg: CFG, block_id: BlockId, incoming: Mapping[BlockId, State]) -> Mapping[BlockId, State]:
+        block = self.blocks[block_id]
+        state = self.evaluate(block, incoming)
+        terminator = block.terminator
+        if isinstance(terminator, Branch):
+            truth = state[terminator.condition].truthiness
+            targets = {
+                Truth.TRUE: (terminator.then_block,),
+                Truth.FALSE: (terminator.else_block,),
+                Truth.UNKNOWN: (terminator.then_block, terminator.else_block),
+            }[truth]
+            return {target: state for target in targets}
+        if isinstance(terminator, Jump):
+            return {terminator.target: state}
+        if isinstance(terminator, ForNext):
+            return {terminator.body: state, terminator.exit: state}
+        return {}
+
+    # ------------------------------------------------------------------ transfer
+
+    def phi(self, phi: Phi, incoming: Mapping[BlockId, State]) -> AbstractValue:
+        result = AbstractValue.bottom()
+        for value, predecessor in phi.incoming:
+            if predecessor in incoming:
+                result = result.join(incoming[predecessor].get(value, AbstractValue.unknown()))
+        return result
+
+    def instruction(self, instruction: Instruction, state: Mapping[Value, AbstractValue]) -> AbstractValue:
+        if isinstance(instruction, Constant):
+            return AbstractValue.of(instruction.value)
+        if isinstance(instruction, BinaryOp):
+            return self.binary(instruction, state[instruction.left], state[instruction.right])
+        if isinstance(instruction, Compare):
+            return self.fold(_COMPARE[instruction.operator], state[instruction.left], state[instruction.right])
+        if isinstance(instruction, UnaryOp):
+            return self.unary(instruction, state[instruction.operand])
+        return AbstractValue.unknown()
+
+    def binary(self, op: BinaryOp, left: AbstractValue, right: AbstractValue) -> AbstractValue:
+        folder = _BINARY.get(op.operator)
+        if folder is not None:
+            numeric_only = op.operator == "mul"
+            folded = self.fold(folder, left, right, numeric_only=numeric_only)
+            if folded.constant is not TOP:
+                return folded
+        return AbstractValue.unknown(self.result_types(op.operator, left, right))
+
+    @staticmethod
+    def result_types(operator_name: str, left: AbstractValue, right: AbstractValue) -> frozenset[str] | None:
+        if left.types is None or right.types is None:
+            return None
+        if left.types <= _NUMERIC and right.types <= _NUMERIC:
+            if operator_name == "div" or "float" in left.types | right.types:
+                return frozenset({"float"})
+            if operator_name in _BINARY:
+                return frozenset({"int"})
+        if operator_name == "add" and left.types == right.types == frozenset({"str"}):
+            return frozenset({"str"})
+        return None
+
+    def unary(self, op: UnaryOp, operand: AbstractValue) -> AbstractValue:
+        if op.operator == "not":
+            if operand.truthiness is Truth.UNKNOWN:
+                return AbstractValue.unknown(frozenset({"bool"}))
+            return AbstractValue.of(operand.truthiness is Truth.FALSE)
+        folder = _UNARY.get(op.operator)
+        if folder is None or operand.constant is TOP or operand.constant is BOTTOM:
+            return AbstractValue.unknown()
+        if not isinstance(operand.constant, (int, float, bool)):
+            return AbstractValue.unknown()
+        try:
+            return AbstractValue.of(folder(operand.constant))
+        except (TypeError, ValueError, ArithmeticError):
+            return AbstractValue.unknown()
+
+    @staticmethod
+    def fold(
+        folder: Callable[..., Any],
+        left: AbstractValue,
+        right: AbstractValue,
+        *,
+        numeric_only: bool = False,
+    ) -> AbstractValue:
+        for side in (left.constant, right.constant):
+            if side is TOP or side is BOTTOM or not isinstance(side, _SAFE_TYPES):
+                return AbstractValue.unknown()
+        if numeric_only and not all(isinstance(c, (int, float, bool)) for c in (left.constant, right.constant)):
+            return AbstractValue.unknown()
+        try:
+            return AbstractValue.of(folder(left.constant, right.constant))
+        except (TypeError, ValueError, ArithmeticError):
+            return AbstractValue.unknown()
+
+
+def propagate_constants(function: FunctionIR, cfg: CFG) -> ConstantFacts:
+    problem = _ConstantProblem(function)
+    solution = solve(problem, cfg)
+    values: dict[Value, AbstractValue] = {}
+    reachable: set[BlockId] = set()
+    for block in function.blocks:
+        if solution.reached(block.id):
+            reachable.add(block.id)
+            values.update(problem.evaluate(block, solution.incoming(block.id)))
+    return ConstantFacts(values, frozenset(reachable))
+
+
+class ConstantPropagation(FunctionAnalysis[ConstantFacts]):
+    name: ClassVar[str] = "abstract.constants"
+    requires: ClassVar[frozenset[AnyAnalysis]] = frozenset({SSAAnalysis, CFGAnalysis})
+
+    @classmethod
+    def compute(cls, ctx: AnalysisContext, function: nodes.Function) -> ConstantFacts:
+        return propagate_constants(ctx.get(SSAAnalysis, function), ctx.get(CFGAnalysis, function))

src/coretrace_python/abstract/values.py

@@ -0,0 +1,56 @@
+"""Abstract values (architecture §18).
+
+An ``AbstractValue`` records what is known about one SSA value: its constant on a flat
+lattice, the set of Python types it may have (``None`` when unknown) and the truthiness
+that follows. Taints, ranges, string constraints and nullability join this record as
+their analyses arrive.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+
+from coretrace_python.dataflow import BOTTOM, TOP, Element, FlatLattice
+
+_CONSTANTS: FlatLattice[object] = FlatLattice()
+
+
+class Truth(Enum):
+    TRUE = "true"
+    FALSE = "false"
+    UNKNOWN = "unknown"
+
+
+@dataclass(frozen=True)
+class AbstractValue:
+    constant: Element[object]
+    types: frozenset[str] | None
+
+    @classmethod
+    def of(cls, value: object) -> AbstractValue:
+        return cls(value, frozenset({type(value).__name__}))
+
+    @classmethod
+    def unknown(cls, types: frozenset[str] | None = None) -> AbstractValue:
+        return cls(TOP, types)
+
+    @classmethod
+    def bottom(cls) -> AbstractValue:
+        return cls(BOTTOM, frozenset())
+
+    @property
+    def truthiness(self) -> Truth:
+        if self.constant is TOP or self.constant is BOTTOM:
+            if self.types == frozenset({"NoneType"}):
+                return Truth.FALSE
+            return Truth.UNKNOWN
+        return Truth.TRUE if self.constant else Truth.FALSE
+
+    def join(self, other: AbstractValue) -> AbstractValue:
+        if self.constant is BOTTOM:
+            return other
+        if other.constant is BOTTOM:
+            return self
+        types = None if self.types is None or other.types is None else self.types | other.types
+        return AbstractValue(_CONSTANTS.join(self.constant, other.constant), types)

src/coretrace_python/dataflow/__init__.py

@@ -0,0 +1,29 @@
+"""Generic data-flow framework: lattices and a bidirectional worklist solver."""
+
+from coretrace_python.dataflow.lattice import (
+    BOTTOM,
+    TOP,
+    Bottom,
+    Element,
+    FlatLattice,
+    Lattice,
+    Top,
+    same_element,
+)
+from coretrace_python.dataflow.solver import ENTRY, DataflowProblem, Direction, Solution, solve
+
+__all__ = [
+    "BOTTOM",
+    "ENTRY",
+    "TOP",
+    "Bottom",
+    "DataflowProblem",
+    "Direction",
+    "Element",
+    "FlatLattice",
+    "Lattice",
+    "Solution",
+    "Top",
+    "same_element",
+    "solve",
+]