diff --git a/docs/source/introduction.ipynb b/docs/source/introduction.ipynb index 400e5104d..93ed18e13 100644 --- a/docs/source/introduction.ipynb +++ b/docs/source/introduction.ipynb @@ -21,7 +21,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "id": "5278fd54", "metadata": {}, "outputs": [], @@ -29,7 +29,7 @@ "import functools\n", "\n", "from effectful.ops.semantics import coproduct, fwd, handler\n", - "from effectful.ops.syntax import defdata, defop\n", + "from effectful.ops.syntax import defdata, defop, syntactic_eq\n", "from effectful.ops.types import NotHandled, Operation, Term\n", "\n", "\n", @@ -65,7 +65,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "id": "3c575e02", "metadata": { "lines_to_next_cell": 2 @@ -78,7 +78,7 @@ "\n", "\n", "def test_adding_two_numbers():\n", - " assert add(1, 3) == add(2, 2)\n", + " assert syntactic_eq(add(1, 3), add(2, 2))\n", "\n", "\n", "assert isinstance(add, Operation)\n", @@ -104,7 +104,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 3, "id": "6e293b33", "metadata": {}, "outputs": [], @@ -183,7 +183,7 @@ "\n", "\n", "def test_adding_a_variable():\n", - " assert add(x, 2) == add(x, 2)\n", + " assert syntactic_eq(add(x, 2), add(x, 2))\n", "\n", "\n", "fails(test_adding_a_variable)\n", @@ -374,7 +374,7 @@ "\n", "\n", "def test_mixed_together():\n", - " assert add(add(x, 2), add(2, y)) == add(x, add(add(3, y), 1))\n", + " assert syntactic_eq(add(add(x, 2), add(2, y)), add(x, add(add(3, y), 1)))\n", "\n", "\n", "fails(test_mixed_together)\n", @@ -512,7 +512,7 @@ "notebook_metadata_filter": "-all" }, "kernelspec": { - "display_name": "base", + "display_name": "effectful (3.12.9.final.0)", "language": "python", "name": "python3" }, @@ -526,7 +526,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.13" + "version": "3.12.9" } }, "nbformat": 4, diff --git a/effectful/handlers/jax/_handlers.py b/effectful/handlers/jax/_handlers.py index 9c933af43..ae70f4993 100644 --- a/effectful/handlers/jax/_handlers.py +++ b/effectful/handlers/jax/_handlers.py @@ -1,17 +1,21 @@ import functools import itertools import typing -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Iterable, Mapping, Sequence from types import EllipsisType from typing import Annotated +from opt_einsum import get_symbol +from opt_einsum.parser import parse_einsum_input + try: import jax + import jax.core import jax.numpy as jnp except ImportError: raise ImportError("JAX is required to use effectful.handlers.jax") -from effectful.ops.semantics import apply, evaluate, fvsof, typeof +from effectful.ops.semantics import apply, evaluate, fvsof, fwd, typeof from effectful.ops.syntax import ( ConstructorOperation, PureInterpretation, @@ -22,6 +26,7 @@ deffn, defop, syntactic_eq, + syntactic_hash, ) from effectful.ops.types import Expr, NotHandled, Operation, Term @@ -29,6 +34,9 @@ IndexElement = None | int | slice | Sequence[int] | EllipsisType | jax.Array +class JaxOperation[**P, T](Operation[P, T]): ... + + def is_eager_array(x): return isinstance(x, jax.Array) or ( isinstance(x, Term) @@ -82,11 +90,17 @@ def _getitem(arr, index): if not is_eager_array(term_arr): return functools.reduce(_merge, arg_sizes, {}) - sizes = ( - {k.op: term_arr.shape[i]} - for i, k in enumerate(term_index) - if isinstance(k, Term) and len(k.args) == 0 and len(k.kwargs) == 0 - ) + sizes = [] + i = 0 + for k in term_index: + if isinstance(k, Term) and len(k.args) == 0 and len(k.kwargs) == 0: + sizes.append({k.op: term_arr.shape[i]}) + + # None in the index expression creates a new dimension, which does + # not appear in the shape of the indexed tensor + if k is not None: + i += 1 + return functools.reduce(_merge, itertools.chain(arg_sizes, sizes), {}) return ( @@ -137,6 +151,13 @@ def _partial_eval(t: Expr[jax.Array]) -> Expr[jax.Array]: if not sized_fvs: return t + # if any dimension is zero sized, the result is empty + if any(size == 0 for size in sized_fvs.values()): + ops = tuple(sized_fvs.keys()) + key = tuple(k() for k in ops) + shape = tuple(sized_fvs[k] for k in ops) + return jax_getitem(jnp.empty(shape), key) + def _is_eager(t): return not isinstance(t, Term) or t.op in sized_fvs or is_eager_array(t) @@ -176,7 +197,7 @@ def _register_jax_op[**P, T](jax_fn: Callable[P, T]): if getattr(jax_fn, "__name__", None) == "__getitem__": return jax_getitem - @defop + @JaxOperation.define def _jax_op(*args, **kwargs) -> jax.Array: tm = defdata(_jax_op, *args, **kwargs) sized_fvs = sizesof(tm) @@ -186,9 +207,15 @@ def _jax_op(*args, **kwargs) -> jax.Array: and not isinstance(args[0], Term) and sized_fvs and args[1] - and all(isinstance(k, Term) and k.op in sized_fvs for k in args[1]) + and all( + (isinstance(k, Term) and k.op in sized_fvs) + or (isinstance(k, slice) and k == slice(None)) + for k in args[1] + ) ): raise NotHandled + elif _jax_op is jax_getitem and not args[1]: + return args[0] elif sized_fvs and set(sized_fvs.keys()) == fvsof(tm) - {jax_getitem, _jax_op}: # note: this cast is a lie. partial_eval can return non-arrays, as # can jax_fn. for example, some jax functions return tuples, @@ -226,6 +253,93 @@ def _jax_op(*args, **kwargs) -> jax.Array: return _jax_op +def _named_dims(term: Expr[jax.Array]) -> tuple[Operation, ...]: + if not (isinstance(term, Term) and term.op == jax_getitem): + return () + index = term.args[1] + assert isinstance(index, Iterable) + return tuple(i.op for i in index if isinstance(i, Term) and not i.args) + + +def _reduce_named(array, axis=None, **kwargs) -> jax.Array: + if axis is None: + return fwd() + + named_dims = _named_dims(array) + if not named_dims: + return fwd() + + bound_arr = bind_dims(array, *named_dims) + + if isinstance(axis, int): + axis = (axis,) + shifted_axis = tuple(a + len(named_dims) if a >= 0 else a for a in axis) + + reduced = fwd(bound_arr, axis=shifted_axis, **kwargs) + return unbind_dims(reduced, *named_dims) + + +def _einsum_named(subscripts, *operands, **kwargs) -> jax.Array: + # only the string-subscripts form is handled; forward the interleaved form + if not isinstance(subscripts, str): + if any(isinstance(x, Term) for x in (subscripts, *operands)): + raise ValueError("Interleaved einsum is not implemented with named tensors") + return jax.numpy.einsum(subscripts, *operands, **kwargs) + + # forward if any operand has a symbolic (Term) shape + if any(isinstance(arr, Term) and not is_eager_array(arr) for arr in operands): + raise NotHandled + + named = [_named_dims(op) for op in operands] + + # normalize: expand ellipses and make the output explicit, using the + # positional shapes (shapes=True avoids materializing the operands) + shapes = [op.shape for op in operands] + in_part, out_part, _ = parse_einsum_input([subscripts, *shapes], shapes=True) + in_specs = in_part.split(",") + assert len(in_specs) == len(operands) + + # fresh symbols for named dims, avoiding every symbol already in use; + # get_symbol gives an effectively unlimited supply (spills into unicode) + used = {c for c in (in_part + out_part) if c not in ",->"} + counter = itertools.count() + + def next_symbol(): + while True: + s = get_symbol(next(counter)) + if s not in used: + used.add(s) + return s + + # assign a letter per unique named dim; shared names reuse the same letter + # so einsum aligns them as batch dims rather than contracting + letter_of, order = {}, [] + for dims in named: + for d in dims: + if d not in letter_of: + letter_of[d] = next_symbol() + order.append(d) + + # bind named dims to leading positional axes and prepend their letters + bound, new_in_specs = [], [] + for op, dims, spec in zip(operands, named, in_specs): + bound.append(bind_dims(op, *dims) if dims else op) + new_in_specs.append("".join(letter_of[d] for d in dims) + spec) + + # add every named dim to the front of the output as passthrough + out_prefix = "".join(letter_of[d] for d in order) + new_subscripts = ",".join(new_in_specs) + "->" + out_prefix + out_part + + result = jax.numpy.einsum(new_subscripts, *bound, **kwargs) + + # unbind: leading axes correspond to `order`, reindex them back to named + reindexed = jax_getitem( + result, + tuple(d() for d in order) + tuple(slice(None) for _ in range(len(out_part))), + ) + return reindexed + + @_register_jax_op def jax_getitem(x: jax.Array, key: tuple[IndexElement, ...]) -> jax.Array: """Operation for indexing an array. Unlike the standard __getitem__ method, @@ -259,11 +373,20 @@ def bind_dims[T, A, B]( >>> bind_dims(t, b, a).shape (3, 2) """ + if not names: + return value + if isinstance(value, Term) and value.op == bind_dims: + return bind_dims(value.args[0], *(names + tuple(value.args[1:]))) if jax.tree_util.treedef_is_leaf(jax.tree.structure(value)): return __dispatch(typeof(value))(value, *names) return jax.tree.map(lambda v: bind_dims(v, *names), value) +@bind_dims.register(object) # type: ignore[attr-defined] +def _(*args, **kwargs): + raise NotHandled + + @defop @_CustomSingleDispatchCallable def unbind_dims[T, A, B]( @@ -272,11 +395,18 @@ def unbind_dims[T, A, B]( *names: Annotated[Operation[[], jax.Array], Scoped[B]], ) -> Annotated[T, Scoped[A | B]]: """Convert positional dimensions to named dimensions.""" + if not names: + return value if jax.tree_util.treedef_is_leaf(jax.tree.structure(value)): return __dispatch(typeof(value))(value, *names) return jax.tree.map(lambda v: unbind_dims(v, *names), value) +@unbind_dims.register(object) # type: ignore[attr-defined] +def _(*args, **kwargs): + raise NotHandled + + def jit(f, *args, **kwargs): f_noindex, f_reindex = _indexed_func_wrapper(f, jax_getitem, sizesof) f_noindex_jitted = jax.jit(f_noindex, *args, **kwargs) @@ -328,3 +458,20 @@ def _(x: jax.Array, other) -> bool: and x.shape == other.shape and bool((jnp.asarray(x) == jnp.asarray(other)).all()) ) + + +@syntactic_eq.register +def _(x: jax.core.Tracer, other) -> bool: + return isinstance(other, jax.core.Tracer) and id(x) == id(other) + + +@syntactic_hash.register +def _(x: jax.Array) -> int: + # Concrete arrays aren't hashable; hash by shape, dtype, and bytes. + return hash(("jax.Array", x.shape, str(x.dtype), x.tobytes())) + + +@syntactic_hash.register +def _(x: jax.core.Tracer) -> int: + # Concrete arrays aren't hashable; hash by shape, dtype, and bytes. + return hash(("jax.core.Tracer", id(x))) diff --git a/effectful/handlers/jax/_terms.py b/effectful/handlers/jax/_terms.py index 812062931..f29c36811 100644 --- a/effectful/handlers/jax/_terms.py +++ b/effectful/handlers/jax/_terms.py @@ -1,5 +1,6 @@ import functools import operator +import typing from collections.abc import Sequence from typing import Any, cast @@ -8,17 +9,24 @@ import effectful.handlers.jax.numpy as jnp from effectful.handlers.jax._handlers import ( IndexElement, - _partial_eval, _register_jax_op, bind_dims, + is_eager_array, jax_getitem, unbind_dims, ) from effectful.internals.tensor_utils import _desugar_tensor_index +from effectful.internals.unification import Box, nested_type from effectful.ops.syntax import defdata from effectful.ops.types import Expr, NotHandled, Operation, Term +@nested_type.register(jax.Array) +@nested_type.register(jax._src.core.Tracer) +def _(value): + return Box(jax.Array) + + class _IndexUpdateHelper: """Helper class to implement array-style .at[index].set() updates for effectful arrays.""" @@ -451,54 +459,76 @@ def _bind_dims_array(t: jax.Array, *args: Operation[[], jax.Array]) -> jax.Array >>> bind_dims(t, b, a).shape (3, 2) """ - - def _evaluate(expr): - if isinstance(expr, Term): - (args, kwargs) = jax.tree.map(_evaluate, (expr.args, expr.kwargs)) - return _partial_eval(expr) - if not jax.tree_util.treedef_is_leaf(jax.tree.structure(expr)): - return jax.tree.map(_evaluate, expr) - return expr - if not isinstance(t, Term): return t - result = _evaluate(t) - if not isinstance(result, Term) or not args: - return result - # ensure that the result is a jax_getitem with an array as the first argument - if not (result.op is jax_getitem and isinstance(result.args[0], jax.Array)): + if not (t.op is jax_getitem and is_eager_array(t)): raise NotHandled - array = result.args[0] - dims = result.args[1] + array = typing.cast(jax.Array, t.args[0]) + dims = t.args[1] assert isinstance(dims, Sequence) + ndim = len(array.shape) # ensure that the order is a subset of the named dimensions order_set = set(args) if not order_set <= set(a.op for a in dims if isinstance(a, Term)): raise NotHandled - # permute the inner array so that the leading dimensions are in the order - # specified and the trailing dimensions are the remaining named dimensions - # (or slices) - reindex_dims = [ - i - for i, o in enumerate(dims) - if not isinstance(o, Term) or o.op not in order_set - ] - dim_ops = [a.op if isinstance(a, Term) else None for a in dims] - perm = ( - [dim_ops.index(o) for o in args] - + reindex_dims - + list(range(len(dims), len(array.shape))) - ) - array = jnp.transpose(array, perm) - reindexed = jax_getitem( - array, (slice(None),) * len(args) + tuple(dims[i] for i in reindex_dims) + def axis_op(ax: int) -> Operation | None: + """The named op of a bare index term at axis ``ax``, else ``None``.""" + if ax < len(dims): + d = dims[ax] + if isinstance(d, Term) and not d.args and not d.kwargs: + return d.op + return None + + # Assign an einsum id to every axis of ``array``. Axes that share a named op + # get the *same* id — a repeated op (e.g. ``arr[i(), i()]``) ties its axes + # together, which einsum reads as a diagonal. Every other axis (slices, ints, + # fancy indices, compound terms, and trailing positional axes) gets a unique + # id, so einsum simply carries it through to be reindexed below. + op_ids: dict[Operation, int] = {} + in_ids: list[int] = [] + next_id = 0 + for ax in range(ndim): + op = axis_op(ax) + if op is not None: + if op not in op_ids: + op_ids[op] = next_id + next_id += 1 + in_ids.append(op_ids[op]) + else: + in_ids.append(next_id) + next_id += 1 + + # Output order: bound args that actually appear (in the requested order, + # deduplicated by the diagonal merge), then every remaining axis in + # first-appearance order. einsum does the permutation and the diagonals; with + # all distinct ids retained in the output it performs no reduction. + present_arg_ids = [op_ids[o] for o in args if o in op_ids] + seen = set(present_arg_ids) + rest_ids: list[int] = [] + for i in in_ids: + if i not in seen: + seen.add(i) + rest_ids.append(i) + + array = jnp.einsum(array, in_ids, present_arg_ids + rest_ids) + + # Re-apply the original index for each carried axis and re-name unbound op + # axes. Trailing positional axes (first appearance beyond ``dims``) are left + # for jax_getitem to carry implicitly. + first_pos: dict[int, int] = {} + for ax, i in enumerate(in_ids): + first_pos.setdefault(i, ax) + + index_expr = (slice(None),) * len(present_arg_ids) + tuple( + dims[first_pos[i]] if first_pos[i] < len(dims) else slice(None) + for i in rest_ids ) - return reindexed + return jax_getitem(array, index_expr) @unbind_dims.register(jax.Array) # type: ignore diff --git a/effectful/handlers/jax/lax/__init__.py b/effectful/handlers/jax/lax/__init__.py new file mode 100644 index 000000000..8b2ea5c77 --- /dev/null +++ b/effectful/handlers/jax/lax/__init__.py @@ -0,0 +1,19 @@ +from typing import TYPE_CHECKING + +import jax.lax + +from effectful.handlers.jax._handlers import _register_jax_op + +for name, op in jax.lax.__dict__.items(): + wrapped_value = None + if callable(op): + wrapped_value = _register_jax_op(op) + else: + continue + + globals()[name] = wrapped_value + + +# Tell mypy about our wrapped functions. +if TYPE_CHECKING: + from jax.lax import * # noqa: F403 diff --git a/effectful/handlers/jax/monoid.py b/effectful/handlers/jax/monoid.py new file mode 100644 index 000000000..2e99d6859 --- /dev/null +++ b/effectful/handlers/jax/monoid.py @@ -0,0 +1,819 @@ +import functools +import logging +import typing +from collections import defaultdict +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import Protocol + +import jax +import jax.core +import opt_einsum +from opt_einsum import get_symbol + +import effectful.handlers.jax.lax as lax +import effectful.handlers.jax.numpy as jnp +from effectful.handlers.jax import bind_dims, jax_getitem, unbind_dims +from effectful.handlers.jax._handlers import is_eager_array +from effectful.handlers.jax.scipy.special import logsumexp +from effectful.ops.monoid import ( + And, + CartesianProduct, + EvaluateIntp, + Max, + Min, + Monoid, + NormalizeIntp, + Or, + Product, + Streams, + Sum, + _conjuncts, + _is_monoid_mask, + _is_monoid_plus, + _is_simple_range, + complement, + distributes_over, + is_equality, +) +from effectful.ops.monoid import Union as UnionM +from effectful.ops.semantics import evaluate, fvsof, fwd, handler, typeof +from effectful.ops.syntax import ( + ObjectInterpretation, + as_dict, + deffn, + implements, + ite, +) +from effectful.ops.types import Expr, Interpretation, Operation, Term + +logger = logging.getLogger(__name__) + + +LogSumExp = Monoid(name="LogSumExp", identity=jnp.asarray(float("-inf"))) + +# ``Sum`` in log space is multiplication, which distributes over ``LogSumExp``: +# a + logsumexp(b, c) = logsumexp(a + b, a + c) +distributes_over.register(Sum, LogSumExp) + +is_equality.register(jnp.equal) +for a, b in { + (jnp.less, jnp.greater), + (jnp.less_equal, jnp.greater_equal), + (jnp.equal, jnp.not_equal), +}: + assert isinstance(a, Operation) + assert isinstance(b, Operation) + complement.register(a, b) + + +def _jax_args(args): + """True iff ``args`` is non-empty, every arg is a concrete + :class:`jax.typing.ArrayLike` or named tensor, and at least one of them is + an array rather than a Python scalar. + + :class:`jax.typing.ArrayLike` is a union that includes ``bool``, ``int``, + ``float`` and ``complex``, so admitting it alone would claim pure-Python + scalar arithmetic. These handlers extend ``EvaluateIntp`` after the scalar + implementations and so take precedence over them, which would silently + narrow a Python float to a ``float32`` array and leave downstream rules + treating a scalar body as array-valued. + """ + return ( + args + and all(isinstance(a, jax.typing.ArrayLike) or is_eager_array(a) for a in args) + and any(not isinstance(a, bool | int | float | complex) for a in args) + ) + + +class PlusCastArray(ObjectInterpretation): + @implements(Monoid.plus) + def plus(self, monoid, *args): + arg_types = [typeof(a) for a in args] + + def _is_jax(t): + return issubclass(t, jax.Array | jax.core.Tracer) + + # exists array valued and non-array-valued args + if any(_is_jax(t) for t in arg_types) and any( + not _is_jax(t) for t in arg_types + ): + return monoid.plus( + *( + a if _is_jax(t) else jnp.asarray(a) + for (a, t) in zip(args, arg_types, strict=True) + ) + ) + + return fwd() + + +class SumPlusJax(ObjectInterpretation): + @implements(Sum.plus) + def plus(self, *args): + if not _jax_args(args): + return fwd() + return functools.reduce(jnp.add, args) + + +class ProductPlusJax(ObjectInterpretation): + @implements(Product.plus) + def plus(self, *args): + if not _jax_args(args): + return fwd() + return functools.reduce(jnp.multiply, args) + + +class MinPlusJax(ObjectInterpretation): + @implements(Min.plus) + def plus(self, *args): + if not _jax_args(args): + return fwd() + return functools.reduce(jnp.minimum, args) + + +class MaxPlusJax(ObjectInterpretation): + @implements(Max.plus) + def plus(self, *args): + if not _jax_args(args): + return fwd() + return functools.reduce(jnp.maximum, args) + + +class LogSumExpPlusJax(ObjectInterpretation): + @implements(LogSumExp.plus) + def plus(self, *args): + if not _jax_args(args): + return fwd() + return functools.reduce(jnp.logaddexp, args) + + +class AndPlusJax(ObjectInterpretation): + @implements(And.plus) + def plus(self, *args): + if not _jax_args(args): + return fwd() + return functools.reduce(jnp.logical_and, args) + + +class OrPlusJax(ObjectInterpretation): + @implements(Or.plus) + def plus(self, *args): + if not _jax_args(args): + return fwd() + return functools.reduce(jnp.logical_or, args) + + +class IteJax(ObjectInterpretation): + @implements(ite) + def ite(self, cond, then, else_): + if not _jax_args((cond,)): + return fwd() + return jnp.where(cond, then, else_) + + +class MaskJax(ObjectInterpretation): + @implements(Monoid.mask) + def mask(self, monoid, value, mask): + if not ( + (is_eager_array(value) or not isinstance(value, Term)) + and (is_eager_array(mask) or not isinstance(mask, Term)) + ): + return fwd() + return jnp.where(mask, value, monoid.identity) + + +class ReduceArrayGather(ObjectInterpretation): + """Split an array-valued stream into an index range and a length-1 stream: + + M.reduce(body, {k: a} ∪ S) ≡ M.reduce(body, {i: range(a.shape[0]), k: (a[i()],)} ∪ S) + + where ``i`` is fresh and ``a[i()] = unbind_dims(a, i)``. The length-1 stream + ``{k: (a[i()],)}`` is then eliminated by + :class:`~effectful.ops.monoid.EliminateSingletonStreams`, which substitutes + ``k := a[i()]`` into the body and the remaining streams. Together the two + steps perform the gather + ``M.reduce(body[k := a[i()]], {i: range(a.shape[0])} ∪ S)``. + """ + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + if typeof(body) is not jax.Array: + return fwd() + + if isinstance(body, Term) and body.op is monoid.delta: + return fwd() + + body_fvs = fvsof(body) + stream_keys = set(streams) + + new_streams: dict = {} + progress = False + for k, v in streams.items(): + if is_eager_array(v) and k in body_fvs and not (fvsof(v) & stream_keys): + index = Operation.define(k) + new_streams[index] = range(v.shape[0]) + new_streams[k] = (unbind_dims(v, index),) + progress = True + else: + new_streams[k] = v + + if not progress: + return fwd() + + return monoid.reduce(body, new_streams) + + +class Reductor(Protocol): + def __call__( + self, arr: jax.Array, axis: int | tuple[int, ...] | None = None + ) -> jax.Array: ... + + +ARRAY_REDUCTORS: dict[Monoid, Reductor] = {} +for monoid, func in [ + (Sum, jnp.sum), + (Product, jnp.prod), + (Min, jnp.min), + (Max, jnp.max), +]: + assert isinstance(monoid, Monoid) + assert callable(func) + ARRAY_REDUCTORS[monoid] = functools.partial(func, initial=monoid.identity) + +ARRAY_REDUCTORS[LogSumExp] = logsumexp + + +class ReduceArray(ObjectInterpretation): + """Reduce an array body over range streams.""" + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + if not is_eager_array(body): + return fwd() + + reductor = ARRAY_REDUCTORS.get(monoid, None) + if reductor is None: + return fwd() + + body_fvs = fvsof(body) + used = {k for k, v in streams.items() if k in body_fvs and isinstance(v, range)} + if not used: + return fwd() + + index = tuple(k() for k in streams if k in used) + arr = monoid.reduce(monoid.delta(index, body), streams) + reduced_body = reductor(arr, axis=tuple(range(len(used)))) + return reduced_body + + +class ReduceDisequalityMask(ObjectInterpretation): + """M.reduce(M.mask(v, And.plus(a != b, c)), S) + ≡ M.reduce(M.plus(M.mask(v, And.plus(a < b, c)), M.mask(v, And.plus(a > b, c)))) + """ + + @implements(Monoid.reduce) + def _(self, monoid, body, streams: Streams): + match body: + case Term(mask_op, (value, mask), {}) if ( + _is_monoid_mask(mask_op) and mask_op.__self__ == monoid + ): + pass + + case _: + return fwd() + + mask_elems = _conjuncts(mask) + + def _neq_to_plus(args, tail_mask_elems): + match args: + case (Term(stream_op, (), {}), index) if _is_simple_range( + streams.get(stream_op, None) + ): + pass + case _: + return None + return monoid.reduce( + monoid.plus( + monoid.mask(value, And.plus(stream_op() < index, *tail_mask_elems)), + monoid.mask(value, And.plus(stream_op() > index, *tail_mask_elems)), + ), + streams, + ) + + for i, elem in enumerate(mask_elems): + match elem: + case Term(jnp.not_equal, args, {}): + tail_mask_elems = [e for (j, e) in enumerate(mask_elems) if i != j] + ret = _neq_to_plus(args, tail_mask_elems) or _neq_to_plus( + tuple(reversed(args)), tail_mask_elems + ) + if ret: + return ret + + return fwd() + + +class ReduceArrayScan(ObjectInterpretation): + @implements(Monoid.reduce) + def _(self, monoid, body, streams: Streams): + match body: + case Term(mask_op, (value, mask), {}) if ( + _is_monoid_mask(mask_op) and mask_op.__self__ == monoid + ): + pass + + case _: + return fwd() + + def _inequality_to_scan(cmp_op, args, tail_mask_elems): + match args: + case (Term(stream_op, (), {}), index): + pass + case _: + return None + + stream = streams.get(stream_op, None) + if not _is_simple_range(stream): + return None + assert isinstance(stream, range) + + reverse = cmp_op in (jnp.greater_equal, jnp.greater) + n = len(stream) + pos_value = monoid.reduce( + monoid.delta((stream_op(),), value), {stream_op: stream} + ) + # inclusive prefix (or suffix, when reverse) scan over the stream + inclusive = lax.associative_scan(monoid.plus, pos_value, reverse=reverse) + # The strict comparisons (`<`, `>`) need an *exclusive* scan: pad + # an identity element on the appropriate side and shift the result + # so that scan_val[k] aggregates only the positions satisfying the + # comparison against k. The non-strict comparisons (`<=`, `>=`) use + # the inclusive scan directly. + match cmp_op: + case jnp.less: + scan_val = jnp.pad( + inclusive, + (1, 0), + mode="constant", + constant_values=monoid.identity, + )[:n] + case jnp.greater: + scan_val = jnp.pad( + inclusive, + (0, 1), + mode="constant", + constant_values=monoid.identity, + )[1:] + case _: + scan_val = inclusive + + tail_body = monoid.mask( + jax_getitem(scan_val, (index,)), And.plus(*tail_mask_elems) + ) + tail_streams = {k: v for (k, v) in streams.items() if k != stream_op} + return monoid.reduce(tail_body, tail_streams) + + mask_elems = _conjuncts(mask) + for i, elem in enumerate(mask_elems): + tail_mask_elems = [e for (j, e) in enumerate(mask_elems) if i != j] + match elem: + case Term( + ( + jnp.less_equal | jnp.greater_equal | jnp.less | jnp.greater + ) as cmp_op, + args, + {}, + ): + ret = _inequality_to_scan( + cmp_op, args, tail_mask_elems + ) or _inequality_to_scan( + complement.of(cmp_op), tuple(reversed(args)), tail_mask_elems + ) + if ret is not None: + return ret + return fwd() + + +class ReduceDeltaSimpleRange(ObjectInterpretation): + """Eliminate a Delta that has independent, dense index arguments. + + + reduce(M, streams ∪ {v: range(N)}, delta((v(),) ++ idx', body)) + ═══════════════════════════════════════════════════════════════════════════ + bind_dims(reduce(M, streams, delta(idx', body[v() := unbind_dims(streams[v], fv)])), fv) + """ + + @implements(Monoid.reduce) + def _(self, monoid: Monoid, body, streams: Streams): + if not (isinstance(body, Term) and body.op == monoid.delta): + return fwd() + + index, weight = body.args + assert isinstance(index, tuple) + + if not index: + return fwd() + + head_index, tail_index = index[0], index[1:] + if not (isinstance(head_index, Term) and head_index.op in streams): + return fwd() + + head_op: Operation = head_index.op + head_stream = streams[head_op] + if not ( + isinstance(head_stream, range) + and head_stream.start == 0 + and head_stream.step == 1 + ): + return fwd() + + tail_streams = {k: v for (k, v) in streams.items() if k != head_op} + + # peel the head index: substitute it into the weight (slicing direct + # uses, materializing the rest) along a fresh named dim, but bind that + # dim only *after* the surrounding reduce -- see the class docstring. + + fresh_op = Operation.define(head_op) + + def _jax_getitem(arr, index): + inner_index, outer_index = [], [] + progress = False + for i in index: + if isinstance(i, Term) and i.op == head_op: + inner_index.append( + slice(head_stream.start, head_stream.stop, head_stream.step) + ) + outer_index.append(fresh_op()) + progress = True + else: + inner_index.append(slice(None)) + outer_index.append(i) + if progress: + return jax_getitem(jax_getitem(arr, inner_index), outer_index) + return fwd(arr, index) + + slice_subst = typing.cast(Interpretation, {jax_getitem: _jax_getitem}) + sliced_weight = handler(slice_subst)(evaluate)(weight) + sliced_streams = handler(slice_subst)(evaluate)(tail_streams) + + gather_subst = typing.cast( + Interpretation, + { + head_op: deffn( + unbind_dims( + jnp.arange( + head_stream.start, head_stream.stop, head_stream.step + ), + fresh_op, + ) + ) + }, + ) + gathered_weight = handler(gather_subst)(evaluate)(sliced_weight) + gathered_streams = handler(gather_subst)(evaluate)(sliced_streams) + + gathered_body = ( + monoid.delta(tail_index, gathered_weight) if tail_index else gathered_weight + ) + inner = ( + monoid.reduce(gathered_body, gathered_streams) + if gathered_streams + else gathered_weight + ) + return bind_dims(inner, fresh_op) + + +class ReduceSumProductContraction(ObjectInterpretation): + """Fast-path a sum-of-products contraction.""" + + @implements(Sum.reduce) + def _(self, body, streams: Streams): + if not ( + isinstance(body, Term) + and _is_monoid_plus(body.op) + and body.op.__self__ is Product + ): + return fwd() + + factors = body.args + if len(factors) != 2 or not all( + issubclass(typeof(f), jax.Array) for f in factors + ): + return fwd() + + (lhs, rhs) = factors + stream_vars = set(streams.keys()) + + # a fully factored reduce only has streams that are used by all factors + shared = fvsof(lhs) & fvsof(rhs) & stream_vars + if shared != stream_vars: + return fwd() + + if not all(isinstance(v, range) for v in streams.values()): + return fwd() + + # create leading reduction dimensions + index = tuple(k() for k in streams) + pos_lhs = Sum.reduce(Sum.delta(index, lhs) if index else lhs, streams) + pos_rhs = Sum.reduce(Sum.delta(index, rhs) if index else rhs, streams) + + dims = "".join(get_symbol(i) for i in range(len(streams))) + contraction = jnp.einsum(f"{dims}...,{dims}...->...", pos_lhs, pos_rhs) + return contraction + + +@dataclass +class Node: + ordinal: frozenset[str] + children: list["Node"] + factors: list[int] + + +class _EinsumBuilder: + in_specs: Sequence[str] + out_spec: str + plates: frozenset[str] + operands: Sequence[jax.Array] + + def __init__( + self, subscripts: str, /, *operands: jax.Array, plates: str | None = None + ): + if not operands: + raise ValueError("einsum requires at least one operand") + + in_spec, out_spec, _ = opt_einsum.parser.parse_einsum_input( + [subscripts, *(op.shape for op in operands)], shapes=True + ) + in_specs = in_spec.split(",") + + self.in_specs = in_specs + self.out_spec = out_spec + self.plates = frozenset(plates or "") + self.operands = operands + + # check that the output spec preserves required plates + out_spec_set = frozenset(self.out_spec) + out_plates = out_spec_set & self.plates + for c in out_spec_set - self.plates: + missing_plates = self.ordinal[c] - out_plates + if missing_plates: + raise ValueError( + "It is nonsensical to preserve a plated dim without preserving " + f"all of that dim's plates, but found {c!r} without " + f"{','.join(sorted(missing_plates))!r}" + ) + + @functools.cached_property + def sizes(self) -> Mapping[str, int]: + """`sizes[d]` is the length of a dimension `d`.""" + sizes: dict[str, int] = {} + for spec, op in zip(self.in_specs, self.operands, strict=True): + for l, s in zip(spec, op.shape, strict=True): + if l in sizes and sizes[l] != s: + raise ValueError(f"Dimension {l} given sizes {s} and {sizes[l]}") + else: + sizes[l] = s + for c in self.out_spec: + if c not in sizes: + raise ValueError(f"einsum: output index {c!r} not present in any input") + return sizes + + @functools.cached_property + def ordinal(self) -> Mapping[str, frozenset[str]]: + """`ordinal[d]` is the plate context of a dimension `d`.""" + ordinal: dict[str, frozenset[str]] = defaultdict(lambda: self.plates) + for spec in self.in_specs: + spec_set = frozenset(spec) + for c in spec_set - self.plates: + ordinal[c] &= spec_set + return ordinal + + @functools.cached_property + def plate_tree(self) -> Node: + # 1. ordinal (plate context) of each factor + factor_ordinal = { + k: frozenset(spec) & self.plates for k, spec in enumerate(self.in_specs) + } + + # 2. node set: every factor ordinal, the global root ∅, and every pairwise + # intersection. The intersection of two ordinals is the deepest context + # that contains both, i.e. their common ancestor -- a shared plate MUST + # have a node to live at, or containment can't be a tree. Closing under + # ∩ materializes those frame nodes (finite lattice, so it terminates). + ordinals: set[frozenset[str]] = set(factor_ordinal.values()) | {frozenset()} + changed = True + while changed: + changed = False + for A in list(ordinals): + for B in list(ordinals): + if (A & B) not in ordinals: + ordinals.add(A & B) + changed = True + + nodes = {o: Node(ordinal=o, children=[], factors=[]) for o in ordinals} + + # 3. parent of o = the largest ordinal strictly inside o ("next context out"). + # Validity: that maximum must dominate EVERY other strict subset, i.e. the + # strict subsets form a chain. Two incomparable maximal subsets means o is a + # join over two separate branches -- not a tree -> raise. + for o in ordinals: + if not o: # root ∅ has no parent + continue + subs = [p for p in ordinals if p < o] # strict subsets, present as nodes + parent = max(subs, key=len) + offenders = [s for s in subs if not (s <= parent)] + if offenders: + raise ValueError( + f"factors do not form a plate tree: context {set(o)} sits above " + f"incomparable sub-plates {[set(s) for s in offenders]} (a join, not a nest)" + ) + nodes[parent].children.append(nodes[o]) + + # 4. hang each factor on its ordinal's node + for k, o in factor_ordinal.items(): + nodes[o].factors.append(k) + + return nodes[frozenset()] # the root + + def dim_type(self, dim: str) -> type: + return ( + int + if dim in self.plates + else Mapping[tuple, int] + if self.ordinal[dim] + else int + ) + + @functools.cached_property + def dim_op(self) -> Mapping[str, Operation]: + return { + dim: Operation.define(self.dim_type(dim), name=dim) + for dim in set("".join(self.in_specs)) + } + + @functools.cached_property + def out_vars(self) -> Mapping[str, Operation]: + return {c: Operation.define(jax.Array, name=f"out_{c}") for c in self.out_spec} + + @functools.cached_property + def global_enums(self) -> frozenset[str]: + return frozenset(c for c, o in self.ordinal.items() if not o) + + @functools.cached_property + def arrays(self) -> list[Term]: + return [Operation.define(jax.Array)() for _ in self.operands] + + def dim_index(self, dim: str) -> Expr: + out_globals = self.global_enums & set(self.out_spec) + + return ( + self.out_vars[dim]() + if dim in out_globals + else self.dim_op[dim]() + if dim in self.plates or not self.ordinal[dim] + else self.dim_op[dim]()[ + tuple(self.dim_op[p]() for p in sorted(self.ordinal[dim])) + ] + ) + + def out_mask(self, vars): + eqs = tuple(self.dim_index(p) == self.out_vars[p]() for p in vars) + return And.plus(*eqs) + + def _build_plate_reductions( + self, plate_tree: Node, parent_plates: frozenset[str] = frozenset() + ) -> Expr[jax.Array]: + + masked_factors = [] + for factor_idx in plate_tree.factors: + spec = self.in_specs[factor_idx] + factor = jax_getitem( + self.arrays[factor_idx], tuple(self.dim_index(d) for d in spec) + ) + # Only plated output dims are delta'd per factor. Global output dims are + # delta'd once by the outer ``out_mask`` in ``_einsum_expr``; masking + # them here too would check the same equality (e.g. ``out_b == b``) in + # every factor that mentions the dim as well as the outer mask. + factor_out_dims = frozenset( + d + for d in set(spec) & set(self.out_vars) - self.plates + if self.ordinal[d] + ) + + if factor_out_dims: + # Preserving a plated output dim requires preserving all its plates, + # as checked in ``__init__``. + factor_out_plates = plate_tree.ordinal & set(self.out_vars) + masked_factors.append( + ite( + self.out_mask(factor_out_plates), + Sum.mask(factor, self.out_mask(factor_out_dims)), + factor, + ) + ) + else: + masked_factors.append(factor) + + child_reductions = ( + self._build_plate_reductions(n, parent_plates | plate_tree.ordinal) + for n in plate_tree.children + ) + product = Product.plus(*masked_factors, *child_reductions) + if plate_tree.ordinal: + plate_streams = { + self.dim_index(p).op: range(self.sizes[p]) + for p in sorted(plate_tree.ordinal - parent_plates) + } + return Product.reduce(product, plate_streams) + return product + + @functools.cached_property + def term(self) -> Expr[Callable]: + # one stream of per-plate-assignment rows for each plated sum dim + rows = {} + plated_enums = {c: o for c, o in self.ordinal.items() if o} + for c, o in plated_enums.items(): + ps = [(p, self.dim_op[p]) for p in sorted(o)] + delta_idx = tuple(op() for (_, op) in ps) + c_streams = {op: range(self.sizes[p]) for (p, op) in ps} + v = Operation.define(int, name=f"{c}_v") + rows[c] = CartesianProduct.reduce( + UnionM.reduce([as_dict((delta_idx, v()))], {v: range(self.sizes[c])}), + c_streams, + ) + + streams: Streams = { + self.dim_op[c]: range(self.sizes[c]) + for c in self.global_enums + if c not in self.out_spec + } | {self.dim_op[c]: r for c, r in rows.items()} + + dims = [(self.out_vars[c], self.sizes[c]) for c in self.out_spec] + + reductions = self._build_plate_reductions(self.plate_tree) + reduction = Sum.reduce(reductions, streams) + return deffn( + bind_dims( + deffn(reduction, *(self.out_vars[c] for c in self.out_spec))( + *(unbind_dims(jnp.arange(d), v) for (v, d) in dims) + ), + *(v for (v, _) in dims), + ), + *(a.op for a in self.arrays), + ) + + +@jax.jit(static_argnums=(0,), static_argnames=("plates",)) +def einsum( + subscripts: str, /, *operands: jax.Array, plates: str | None = None +) -> jax.Array: + """Evaluate an einsum expression using monoid reductions. + + Generalizes :func:`jax.numpy.einsum` with plated dimensions in the style of + :func:`pyro.ops.contract.einsum`: indices in ``plates`` are plate + dimensions, and reductions along plates are product reductions. A sum + dimension that always appears together with a plate denotes a distinct + variable for each slice of that plate; its plate context (ordinal) is the + intersection of the plate sets of the inputs that mention it. When such a + dimension appears in the output it must be accompanied by all of its + plates. + + The expression is represented naively: each plated sum dimension ranges + over a :data:`CartesianProduct` stream of per-plate-assignment rows, and + each plated input is a :data:`Product` reduction over its plates. + """ + expr = _EinsumBuilder(subscripts, *operands, plates=plates).term + + with handler(NormalizeIntp): + norm_expr = evaluate(expr) + + assert CartesianProduct.reduce not in fvsof(norm_expr), ( + "failed to eliminate cartesian products" + ) + + with handler(EvaluateIntp), handler(NormalizeIntp): + assert callable(norm_expr) + result = norm_expr(*operands) + assert isinstance(result, jax.Array), "failed to fully evaluate" + return result + + +EvaluateIntp.extend( + SumPlusJax(), + ProductPlusJax(), + MinPlusJax(), + MaxPlusJax(), + LogSumExpPlusJax(), + AndPlusJax(), + OrPlusJax(), + IteJax(), + MaskJax(), + ReduceSumProductContraction(), + ReduceArray(), + ReduceDeltaSimpleRange(), + ReduceArrayScan(), + PlusCastArray(), +) + +NormalizeIntp.extend(ReduceArrayGather()) diff --git a/effectful/handlers/jax/numpy/__init__.py b/effectful/handlers/jax/numpy/__init__.py index cc20d7498..34ec5a604 100644 --- a/effectful/handlers/jax/numpy/__init__.py +++ b/effectful/handlers/jax/numpy/__init__.py @@ -1,11 +1,20 @@ import types +import typing from typing import TYPE_CHECKING import jax.numpy -from .._handlers import _register_jax_op, _register_jax_op_no_partial_eval +from effectful.handlers.jax._handlers import ( + _einsum_named, + _reduce_named, + _register_jax_op, + _register_jax_op_no_partial_eval, +) +from effectful.ops.semantics import handler +from effectful.ops.types import Operation -_no_overload = ["array", "asarray"] +_NO_OVERLOAD = ["array", "asarray"] +_REDUCTION = ["sum", "prod", "min", "max", "any", "all", "mean", "argmax"] for name, op in jax.numpy.__dict__.items(): if isinstance(op, types.ModuleType): @@ -19,18 +28,41 @@ if name == "__getattr__": continue - elif name in _no_overload: + elif name in _NO_OVERLOAD: globals()[name] = _register_jax_op_no_partial_eval(op) else: globals()[name] = _register_jax_op(op) jax_op = ( _register_jax_op_no_partial_eval(op) - if name in _no_overload + if name in _NO_OVERLOAD else _register_jax_op(op) ) globals()[name] = jax_op +for name in _REDUCTION: + op = globals()[name] + globals()[name] = handler({op: _reduce_named})(op) + + +einsum = Operation.define(_einsum_named) + + +@Operation.define +def asarray(a, **kwargs) -> jax.Array: + import jax.core + + from effectful.ops.semantics import typeof + from effectful.ops.types import NotHandled, Term + + if isinstance(a, Term): + if issubclass(typeof(a), jax.Array | jax.core.Tracer) and not kwargs: + return typing.cast(jax.Array, a) + else: + raise NotHandled + return jax.numpy.asarray(a, **kwargs) + + # Tell mypy about our wrapped functions. if TYPE_CHECKING: - from jax.numpy import * # noqa: F403 + from jax.numpy import * # type: ignore[assignment] # noqa: F403 diff --git a/effectful/handlers/jax/scipy/special.py b/effectful/handlers/jax/scipy/special.py index afe1334b8..67b99621a 100644 --- a/effectful/handlers/jax/scipy/special.py +++ b/effectful/handlers/jax/scipy/special.py @@ -2,9 +2,11 @@ import jax.scipy.special -from effectful.handlers.jax._handlers import _register_jax_op +from effectful.handlers.jax._handlers import _reduce_named, _register_jax_op +from effectful.ops.semantics import handler logsumexp = _register_jax_op(jax.scipy.special.logsumexp) +logsumexp = handler({logsumexp: _reduce_named})(logsumexp) # Tell mypy about our wrapped functions. if TYPE_CHECKING: diff --git a/effectful/internals/disassembly.py b/effectful/internals/disassembly.py new file mode 100644 index 000000000..d2132b193 --- /dev/null +++ b/effectful/internals/disassembly.py @@ -0,0 +1,3590 @@ +""" +Generator expression bytecode reconstruction module. + +This module provides functionality to reconstruct AST representations from compiled +generator expressions by analyzing their bytecode. The primary use case is to recover +the original structure of generator comprehensions from their compiled form. + +The only public-facing interface is the `disassemble()` function, which takes a +generator object and returns an AST node representing the original comprehension. +All other functions and classes in this module are internal implementation details. + +Example: + >>> g = (x * 2 for x in range(10) if x % 2 == 0) + >>> ast_node = disassemble(g) + >>> # ast_node is now an ast.Expression representing the original expression + +What is recovered, and what is not: + + What comes back is the comprehension's *syntax*, not the generator's + suspended state, so evaluating the reconstruction runs every expression in + it a second time. For a comprehension over pure expressions the two agree + element for element. Where they can part company: + + * A **stateful filter or element expression** -- one that mutates something, + or whose value depends on state that has moved on -- is re-run, and can + answer differently the second time. `(x for x in xs if next(flags))` + reconstructs faithfully as source, but iterating the reconstruction draws + from `flags` where it now stands, not from where it stood when the + original generator was built. The side effects happen twice, once per + iteration of each generator. + + * The **outermost iterable** is not part of the comprehension's bytecode -- + it is an object the generator already holds -- so it is recovered from + that object rather than from source. What lands in the tree is a snapshot + of the elements not yet consumed: `iter([1, 2, 3])` advanced once becomes + the literal `[2, 3]`. The expression that produced it is gone, and so is + any laziness it had. + + * A **free variable** is likewise recovered by value, not by name: the value + in the closure cell is written into the tree. A captured object with no + AST spelling raises `TypeError` rather than reconstructing to a name that + would resolve against the evaluating namespace instead. +""" + +import ast +import builtins +import collections +import collections.abc +import copy +import dis +import enum +import functools +import inspect +import itertools +import sys +import types +import typing +from collections.abc import Callable, Generator, Iterator +from dataclasses import dataclass, field, replace + +CompExp = ast.GeneratorExp | ast.ListComp | ast.SetComp | ast.DictComp + + +class Placeholder(ast.Name): + """Placeholder for AST nodes that are not yet resolved.""" + + def __init__( + self, + id: typing.Literal[".PLACEHOLDER"] = ".PLACEHOLDER", + ctx: ast.Load = ast.Load(), + ): + super().__init__(id=id, ctx=ctx) + + +class DummyIterName(ast.Name): + """Dummy name for the iterator variable in generator expressions.""" + + def __init__(self, id: typing.Literal[".0"] = ".0", ctx: ast.Load = ast.Load()): + super().__init__(id=id, ctx=ctx) + + +class Skipped(ast.Name): + """Placeholder for skipped branches in if-expressions. + + ``id`` is defaulted so that ``copy.deepcopy`` can reconstruct the node: on + Python 3.12 ``ast.AST.__reduce__`` supplies no positional arguments (3.13+ + supplies one per field), so the constructor must be callable with none. + """ + + def __init__(self, id: str = "", ctx: ast.Load = ast.Load()): + super().__init__(id=id, ctx=ctx) + + +class CommonConstant(ast.Name): + """A constant pushed by 3.14's LOAD_COMMON_CONSTANT. + + 3.14 can inline `any()`/`all()` over a generator, guarding the fast path + with `loaded_name is `. Marking the builtin lets that guard be + recognised so the generic call path is followed instead, which still spells + out the original call. + """ + + def __init__(self, id: str = "", ctx: ast.Load = ast.Load()): + super().__init__(id=id, ctx=ctx) + + +class TargetHole(ast.Name): + """Placeholder for a comprehension loop target that is not yet named. + + ``FOR_ITER`` knows that a loop target exists but not what it is called; the + name only arrives with the ``STORE_FAST``/``UNPACK_SEQUENCE`` instructions + that follow. Each hole carries a unique ``id`` so that it can be located + again inside a comprehension after the surrounding state has been copied, + which matters for unpacking targets where several holes are live at once. + """ + + _counter: typing.ClassVar[Iterator[int]] = itertools.count() + + def __init__(self, id: str = "", ctx: ast.Store = ast.Store()): + super().__init__(id=id or f".TARGET_{next(TargetHole._counter)}", ctx=ctx) + + +class ReplaceTargetHole(ast.NodeTransformer): + """Replace the uniquely-identified :class:`TargetHole` ``id`` with ``replacement``.""" + + id: str + replacement: ast.expr + + def __init__(self, id: str, replacement: ast.expr): + self.id = id + self.replacement = replacement + super().__init__() + + def visit_TargetHole(self, node: TargetHole) -> ast.expr: + return self.replacement if node.id == self.id else node + + +def _bind_target_hole( + stack: list[ast.expr], hole: ast.expr, replacement: ast.expr +) -> list[ast.expr]: + """Fill the loop-target hole ``hole`` of the innermost matching comprehension. + + Returns a new stack; the hole itself is left in place for the caller to pop. + """ + assert isinstance(hole, TargetHole), f"Expected a loop target hole, got {hole}" + for pos, item in zip(reversed(range(len(stack))), reversed(stack)): + if not isinstance(item, CompExp) or not item.generators: + continue + if not any( + isinstance(n, TargetHole) and n.id == hole.id + for n in ast.walk(item.generators[-1].target) + ): + continue + new_comp = ReplaceTargetHole(hole.id, replacement).visit(copy.deepcopy(item)) + return stack[:pos] + [new_comp] + stack[pos + 1 :] + + raise TypeError(f"No comprehension found with loop target hole {hole.id}") + + +class Null(ast.Constant): + """Placeholder for NULL values generated in bytecode.""" + + def __init__(self, value: None = None): + super().__init__(value=value) + + +class ConvertedValue(ast.expr): + """Wrapper for values that have been converted with CONVERT_VALUE.""" + + value: ast.expr + conversion: int + ast_conversion: int + + def __init__(self, value: ast.expr, conversion: int): + self.value = value + self.conversion = conversion + # Map CONVERT_VALUE args to ast.FormattedValue conversion values + # CONVERT_VALUE: 0=None, 1=str, 2=repr, 3=ascii + # ast.FormattedValue: -1=none, 115=str, 114=repr, 97=ascii + conversion_map = {0: -1, 1: 115, 2: 114, 3: 97} + self.ast_conversion = conversion_map.get(conversion, -1) + + +class CompLambda(ast.Lambda): + """Placeholder AST node representing a lambda function used in comprehensions.""" + + def __init__(self, body: CompExp): + assert isinstance(body, CompExp) + assert sum(1 for x in ast.walk(body) if isinstance(x, DummyIterName)) == 1 + assert len(body.generators) > 0 + assert isinstance(body.generators[0].iter, DummyIterName) + args = ast.arguments( + posonlyargs=[ast.arg(DummyIterName().id)], + args=[], + kwonlyargs=[], + kw_defaults=[], + defaults=[], + ) + super().__init__(args=args, body=body) + + def __copy__(self): + """Support copy.copy operation.""" + assert isinstance(self.body, CompExp) + return CompLambda(self.body) + + def __deepcopy__(self, memo): + """Support copy.deepcopy operation.""" + assert isinstance(self.body, CompExp) + return CompLambda(copy.deepcopy(self.body, memo)) + + def inline(self, iterator: ast.expr) -> CompExp: + assert isinstance(self.body, CompExp) + res: CompExp = copy.deepcopy(self.body) + res.generators[0].iter = iterator + return res + + +class ReplacePlaceholder(ast.NodeTransformer): + value: ast.expr + _done: bool + + def __init__(self, value: ast.expr): + self.value = value + self._done = False + super().__init__() + + def visit(self, node): + if isinstance(node, Placeholder) and not self._done: + self._done = True + return self.value + else: + return self.generic_visit(node) + + +class ReplaceSkipped(ast.NodeTransformer): + id: str + replacement: ast.expr + + def __init__(self, id: str, replacement: ast.expr): + self.id = id + self.replacement = copy.deepcopy(replacement) + super().__init__() + + def visit_IfExp(self, node: ast.IfExp): + if isinstance(node.body, Skipped) and node.body.id == self.id: + return ast.IfExp(test=node.test, body=self.replacement, orelse=node.orelse) + elif isinstance(node.orelse, Skipped) and node.orelse.id == self.id: + return ast.IfExp(test=node.test, body=node.body, orelse=self.replacement) + else: + return self.generic_visit(node) + + +class BranchState(typing.NamedTuple): + testval: bool + value: ast.expr + + +class BranchIdentifier(ast.NodeVisitor): + branching: collections.abc.MutableMapping[str, BranchState] + filter_lengths: list[int] + + def __init__(self): + self.branching = {} + self.filter_lengths = [] + super().__init__() + + def visit_IfExp(self, node: ast.IfExp): + if isinstance(node.body, Skipped): + self.branching[node.body.id] = BranchState( + testval=False, value=copy.deepcopy(node.orelse) + ) + elif isinstance(node.orelse, Skipped): + self.branching[node.orelse.id] = BranchState( + testval=True, value=copy.deepcopy(node.body) + ) + return self.generic_visit(node) + + def visit_comprehension(self, node: ast.comprehension): + self.filter_lengths.append(len(node.ifs)) + return self.generic_visit(node) + + +@functools.cache +def _instructions( + code: types.CodeType, +) -> collections.abc.Mapping[int, dis.Instruction]: + """Decode a code object once; every state derived from it shares the result.""" + return collections.OrderedDict( + (instr.offset, instr) for instr in dis.get_instructions(code) + ) + + +@functools.cache +def _next_instructions( + code: types.CodeType, +) -> collections.abc.Mapping[int, dis.Instruction]: + """Map each instruction offset to the instruction that follows it.""" + ordered = list(_instructions(code).values()) + return {before.offset: after for before, after in zip(ordered[:-1], ordered[1:])} + + +@dataclass(frozen=True) +class ReconstructionState: + """State maintained during AST reconstruction from bytecode. + + This class tracks all the information needed while processing bytecode + instructions to reconstruct the original comprehension's AST. It acts + as the working memory during the reconstruction process, maintaining + both the evaluation stack state and the high-level comprehension structure + being built. + + The reconstruction process works by simulating the Python VM's execution + of the bytecode, but instead of executing operations, it builds AST nodes + that represent those operations. + + Attributes: + code: The compiled code object from which the bytecode is being processed. + This is typically obtained from a generator function or comprehension. + + stack: Simulates the Python VM's value stack. Contains AST nodes or + values that would be on the stack during execution. Operations + like LOAD_FAST push to this stack, while operations like + BINARY_ADD pop operands and push results. + """ + + code: types.CodeType + instruction: dis.Instruction + + stack: list[ast.expr] = field(default_factory=list) + result: ast.expr = field(default_factory=Placeholder) + + # How many times each FOR_ITER has been entered on this path. + loops: dict[int, int] = field(default_factory=dict) + finished: bool = field(default=False) + + # Which edge each already-resolved conditional jump took on this path. + branches: "dict[int, BranchEdge]" = field(default_factory=dict) + + # Locals bound to a known expression rather than to a loop target. 3.14 + # unrolls a single-iteration loop over a literal, storing its targets + # directly, so those names have to be substituted back at their uses. + bindings: dict[str, ast.expr] = field(default_factory=dict) + + # Set by KW_NAMES (Python 3.12 only) and consumed by the following CALL. + # KW_NAMES has no stack effect, so the names cannot live on `stack`. + kw_names: tuple[str, ...] | None = field(default=None) + + # The value captured in each closure cell the code reads, by name. A free + # variable is looked up in a cell, not in globals, so reconstructing it as a + # bare name would resolve to whatever the evaluating namespace happens to + # bind; the captured value has to be written into the tree instead. + freevars: dict[str, ast.expr] = field(default_factory=dict) + + @property + def instructions(self) -> collections.abc.Mapping[int, dis.Instruction]: + """The bytecode instructions of the current code object, by offset.""" + return _instructions(self.code) + + @property + def next_instructions(self) -> collections.abc.Mapping[int, dis.Instruction]: + return _next_instructions(self.code) + + +# Python version enum for version-specific handling +class PythonVersion(enum.IntEnum): + PY_312 = 12 + PY_313 = 13 + PY_314 = 14 + + +def current_version() -> PythonVersion: + """The bytecode dialect of the running interpreter. + + Raises on a Python this module has not been taught, rather than guessing + that the previous release's opcodes still mean what they used to. + """ + try: + return PythonVersion(sys.version_info.minor) + except ValueError as e: + supported = ", ".join(f"3.{v.value}" for v in PythonVersion) + raise NotImplementedError( + f"effectful.internals.disassembly supports {supported}, " + f"not 3.{sys.version_info.minor}" + ) from e + + +# Global handler registry +OpHandler = Callable[[ReconstructionState, dis.Instruction], ReconstructionState] + +OP_HANDLERS: dict[str, OpHandler] = {} + + +@typing.overload +def register_handler( + opname: str, *, version: PythonVersion +) -> Callable[[OpHandler], OpHandler]: ... + + +@typing.overload +def register_handler( + opname: str, + handler: OpHandler, + *, + version: PythonVersion, +) -> OpHandler: ... + + +def register_handler( + opname: str, + handler=None, + *, + version: PythonVersion, +): + """Register a handler for one opcode in one Python bytecode dialect. + + Every dialect a handler applies to is named explicitly. Opcodes are not + assumed to carry forward: a release can keep an opcode's name while changing + what it does, so applicability to a new Python is a decision to make per + opcode rather than a default. + """ + if handler is None: + return functools.partial(register_handler, opname, version=version) + + # Skip registration if version doesn't match current version + if version != current_version(): + return handler + + # Only check opmap if the version matches (or no version specified) + assert opname in dis.opmap, f"Invalid operation name: '{opname}'" + + if opname in OP_HANDLERS: + raise ValueError(f"Handler for '{opname}' (version {version}) already exists.") + + if dis.opmap[opname] in dis.hasjrel: + assert opname in LOOP_OPS | BRANCH_OPS | JUMP_OPS + else: + assert opname not in LOOP_OPS | BRANCH_OPS | JUMP_OPS + + @functools.wraps(handler) + def _wrapper( + state: ReconstructionState, + instr: dis.Instruction, + ) -> ReconstructionState: + assert instr.opname == opname, ( + f"Handler for '{opname}' called with wrong instruction" + ) + assert not state.finished, "Cannot process instruction on finished state" + + new_state = handler(state, instr) + + jump: bool | None # argument to dis.stack_effect + if instr.opname in LOOP_OPS: + if state.loops.get(instr.offset, 0) > 0: + new_state = replace( + new_state, instruction=state.instructions[instr.argval] + ) + jump = True + else: + # Copy rather than mutate: continuations forked from this state + # share the mapping and must not see each other's loop counts. + new_state = replace( + new_state, + instruction=state.next_instructions[instr.offset], + loops={ + **state.loops, + instr.offset: state.loops.get(instr.offset, 0) + 1, + }, + ) + jump = False + elif instr.opname in BRANCH_OPS: + if new_state.branches.get(instr.offset) == BranchEdge.FALL_THROUGH: + new_state = replace( + new_state, instruction=state.next_instructions[instr.offset] + ) + jump = False + else: + new_state = replace( + new_state, instruction=state.instructions[instr.argval] + ) + jump = True + elif instr.opname in JUMP_OPS: + new_state = replace(new_state, instruction=state.instructions[instr.argval]) + jump = True + elif instr.opname not in RETURN_OPS and instr.offset in state.next_instructions: + new_state = replace( + new_state, instruction=state.next_instructions[instr.offset] + ) + jump = None + else: + new_state = replace(new_state, finished=True) + jump = None + + # post-condition: check stack effect + expected_stack_effect = dis.stack_effect(instr.opcode, instr.arg, jump=jump) + actual_stack_effect = len(new_state.stack) - len(state.stack) + assert len(state.stack) + expected_stack_effect >= 0, ( + f"Handler for '{opname}' would result in negative stack size" + ) + assert actual_stack_effect == expected_stack_effect, ( + f"Handler for '{opname}' has incorrect stack effect: " + f"expected {expected_stack_effect}, got {actual_stack_effect}" + ) + + return new_state + + OP_HANDLERS[opname] = _wrapper + return handler # return the original handler for multiple decorator usage + + +LOOP_OPS: set[typing.Literal["FOR_ITER"]] = {"FOR_ITER"} + +BRANCH_OPS: set[ + typing.Literal[ + "POP_JUMP_IF_TRUE", + "POP_JUMP_IF_FALSE", + "POP_JUMP_IF_NOT_NONE", + "POP_JUMP_IF_NONE", + ] +] = { + "POP_JUMP_IF_TRUE", + "POP_JUMP_IF_FALSE", + "POP_JUMP_IF_NOT_NONE", + "POP_JUMP_IF_NONE", +} + +RETURN_OPS: set[typing.Literal["RETURN_VALUE", "RETURN_CONST"]] = { + "RETURN_VALUE", + "RETURN_CONST", +} + +JUMP_OPS = {dis.opname[d] for d in dis.hasjrel} - LOOP_OPS - BRANCH_OPS - RETURN_OPS + + +# Instructions that emit an element of the comprehension being built. Reaching +# one of these means the current iteration was *not* filtered out. +PRODUCE_OPS = {"YIELD_VALUE", "LIST_APPEND", "SET_ADD", "MAP_ADD"} + + +def _successor_offsets(state: ReconstructionState, instr: dis.Instruction) -> list[int]: + """Offsets control can transfer to from ``instr``, ignoring exception edges.""" + following = state.next_instructions.get(instr.offset) + if instr.opname in BRANCH_OPS | LOOP_OPS: + return [instr.argval] + ([following.offset] if following else []) + elif instr.opname in JUMP_OPS: + return [instr.argval] + elif instr.opname in RETURN_OPS: + return [] + else: + return [following.offset] if following else [] + + +def _reachable_outcomes(state: ReconstructionState, start: int) -> tuple[bool, bool]: + """From ``start``, can the iteration be skipped, and can an element be produced? + + Returns ``(can_skip, can_produce)``. "Skip" means reaching the loop + back-edge without emitting an element, i.e. being filtered out. + """ + seen: set[int] = set() + pending = [start] + can_skip = can_produce = False + + while pending: + offset = pending.pop() + if offset in seen or offset not in state.instructions: + continue + seen.add(offset) + + instr = state.instructions[offset] + if instr.opname in PRODUCE_OPS: + can_produce = True + continue + if ( + instr.opname == "JUMP_BACKWARD" + and state.instructions[instr.argval].opname in LOOP_OPS + ): + can_skip = True + continue + + pending.extend(_successor_offsets(state, instr)) + + return can_skip, can_produce + + +class BranchEdge(enum.IntEnum): + """Which way a conditional jump was resolved on the path being explored.""" + + TAKE_JUMP = 1 + FALL_THROUGH = 2 + + +class BranchKind(enum.Enum): + """What role a conditional jump plays in a comprehension. + + TERNARY + A conditional expression: the arms reconverge having each pushed a + value, and both are spliced back together into an ``ast.IfExp``. + FILTER + Part of a filter's condition. The condition is consumed rather than + producing a value, so each surviving path records the conjunction of + tests that got it to the element, and the filter as a whole is the + disjunction of those conjunctions. + """ + + TERNARY = enum.auto() + FILTER = enum.auto() + + +@functools.cache +def _stack_depths(code: types.CodeType) -> collections.abc.Mapping[int, int]: + """VM stack depth on entry to each reachable instruction. + + Depths are relative to the start of the code object, which is all that is + needed to tell a value-producing branch from a control-flow one. + """ + instructions = collections.OrderedDict( + (i.offset, i) for i in dis.get_instructions(code) + ) + ordered = list(instructions.values()) + following = {a.offset: b.offset for a, b in zip(ordered[:-1], ordered[1:])} + + depths: dict[int, int] = {ordered[0].offset: 0} + pending = collections.deque([ordered[0].offset]) + while pending: + offset = pending.popleft() + instr, depth = instructions[offset], depths[offset] + if instr.opname in RETURN_OPS: + continue + + edges: list[tuple[int, bool | None]] = [] + if instr.opname in BRANCH_OPS | LOOP_OPS: + edges = [(instr.argval, True)] + if offset in following: + edges.append((following[offset], False)) + elif instr.opname in JUMP_OPS: + edges = [(instr.argval, True)] + elif offset in following: + edges = [(following[offset], None)] + + for target, jump in edges: + if target in instructions and target not in depths: + depths[target] = depth + dis.stack_effect( + instr.opcode, instr.arg, jump=jump + ) + pending.append(target) + + return depths + + +def _forward_reachable(state: ReconstructionState, start: int) -> set[int]: + """Offsets reachable from ``start`` without producing or looping back.""" + seen: set[int] = set() + pending = [start] + while pending: + offset = pending.pop() + if offset in seen or offset not in state.instructions: + continue + seen.add(offset) + + instr = state.instructions[offset] + if instr.opname in PRODUCE_OPS: + continue + if ( + instr.opname == "JUMP_BACKWARD" + and state.instructions[instr.argval].opname in LOOP_OPS + ): + continue + + pending.extend(_successor_offsets(state, instr)) + + return seen + + +def _is_conditional_expression( + state: ReconstructionState, instr: dis.Instruction +) -> bool: + """Do the two edges of ``instr`` reconverge one stack slot deeper? + + That is the signature of a conditional expression: each arm leaves a value + behind and control rejoins to consume it. A filter's condition is consumed + by the jump itself, so wherever its edges meet -- if they meet at all -- the + stack is no deeper than it was. + """ + following = state.next_instructions.get(instr.offset) + if following is None: + return False + + common = _forward_reachable(state, instr.argval) & _forward_reachable( + state, following.offset + ) + if not common: + return False # the edges never rejoin, so nothing was left on the stack + + depths = _stack_depths(state.code) + join = min(common) # arms are laid out contiguously, so the join comes first + if join not in depths or following.offset not in depths: + return False + return depths[join] == depths[following.offset] + 1 + + +def _classify_branch( + state: ReconstructionState, instr: dis.Instruction +) -> tuple[BranchKind, list[BranchEdge]]: + """Classify a conditional jump and list the edges worth exploring.""" + both = [BranchEdge.TAKE_JUMP, BranchEdge.FALL_THROUGH] + following = state.next_instructions.get(instr.offset) + if following is None: + return BranchKind.TERNARY, both + + jump_skip, _ = _reachable_outcomes(state, instr.argval) + fall_skip, _ = _reachable_outcomes(state, following.offset) + + # Neither edge can drop the current iteration -- because there is no loop at + # all (a lambda body) or because the element is produced regardless (a + # conditional in the element expression). Either way nothing is filtered. + if not jump_skip and not fall_skip: + return BranchKind.TERNARY, both + + # Otherwise the branch could be a filter, or a conditional expression that + # merely happens to sit inside one. Only the latter leaves a value behind. + if _is_conditional_expression(state, instr): + return BranchKind.TERNARY, both + + # An edge that cannot reach an element contributes nothing to the filter, so + # there is no point walking it. Pruning those edges is also what keeps the + # executor out of the operand-cleanup block on a chained comparison's + # failing edge. + live = [ + edge + for edge, start in ( + (BranchEdge.TAKE_JUMP, instr.argval), + (BranchEdge.FALL_THROUGH, following.offset), + ) + if _reachable_outcomes(state, start)[1] + ] + return BranchKind.FILTER, live or [BranchEdge.TAKE_JUMP] + + +def _negate(condition: ast.expr) -> ast.expr: + """Logical negation, cancelling a `not` rather than stacking another one.""" + if isinstance(condition, ast.UnaryOp) and isinstance(condition.op, ast.Not): + return condition.operand + return ast.UnaryOp(op=ast.Not(), operand=condition) + + +def _conjoin(conditions: list[ast.expr]) -> ast.expr | None: + """Combine the entries of a ``comprehension.ifs`` list into one expression.""" + if not conditions: + return None + elif len(conditions) == 1: + return conditions[0] + else: + return ast.BoolOp(op=ast.And(), values=list(conditions)) + + +def _disjoin(left: ast.expr, right: ast.expr) -> ast.expr: + """Combine two conditions with ``or``, flattening nested disjunctions. + + Two rewrites are applied while combining. Duplicate disjuncts are dropped, + because paths through independent filters repeat them. And ``X or (not X and + Y)`` becomes ``X or Y``: enumerating paths records the negation of every + test a path declined, so a later disjunct restates the negation of an + earlier one. Dropping it is not merely tidier -- `or` short-circuits, so the + earlier disjunct has already been evaluated, and leaving the negation in + would evaluate it a second time, which is visibly wrong when the condition + contains an assignment expression. + + Conditions are keyed by ``ast.dump`` exactly once each: these lists get long + and the expressions large, so re-dumping per comparison dominates. + """ + values: list[ast.expr] = [] + for side in (left, right): + if isinstance(side, ast.BoolOp) and isinstance(side.op, ast.Or): + values.extend(side.values) + else: + values.append(side) + + unique: list[ast.expr] = [] + seen: set[str] = set() + for value in values: + key = ast.dump(value) + if key in seen: + continue + seen.add(key) + + # Absorb the negations of the disjuncts already accepted. + if isinstance(value, ast.BoolOp) and isinstance(value.op, ast.And): + kept = [ + conjunct + for conjunct in value.values + if not ( + isinstance(conjunct, ast.UnaryOp) + and isinstance(conjunct.op, ast.Not) + and ast.dump(conjunct.operand) in seen + ) + ] + if kept and len(kept) < len(value.values): + conjoined = _conjoin(kept) + assert conjoined is not None + value = conjoined + + unique.append(value) + + return unique[0] if len(unique) == 1 else ast.BoolOp(op=ast.Or(), values=unique) + + +def _merge_filters_into( + node: typing.Any, other: typing.Any, mutate: bool = True +) -> bool: + """Walk two results in parallel, OR-ing the filters where they disagree. + + Everything outside a ``comprehension.ifs`` has to match exactly; the ifs are + where the paths are allowed to differ, and are combined rather than + compared. Filters are not recursed into, so a nested comprehension inside a + filter is treated as part of that filter's condition. + + Returns False if the two results differ somewhere they may not. With + ``mutate=False`` nothing is written, which allows compatibility to be tested + before paying for a deep copy -- most candidate pairs do not merge, and the + copy dominates otherwise. + """ + if type(node) is not type(other): + return False + + if isinstance(node, ast.comprehension): + if ast.dump(node.target) != ast.dump(other.target): + return False + if not _merge_filters_into(node.iter, other.iter, mutate): + return False + if not mutate: + return True + + guard, other_guard = _conjoin(node.ifs), _conjoin(other.ifs) + if guard is None or other_guard is None: + # One path reached the element unconditionally, so the filter as a + # whole is unconditional at this generator. + node.ifs = [] + elif ast.dump(guard) != ast.dump(other_guard): + node.ifs = [_disjoin(guard, other_guard)] + return True + + if not isinstance(node, ast.AST): + return bool(node == other) + + for name in node._fields: + mine, theirs = getattr(node, name, None), getattr(other, name, None) + if isinstance(mine, list) or isinstance(theirs, list): + if not isinstance(mine, list) or not isinstance(theirs, list): + return False + if len(mine) != len(theirs): + return False + if not all(_merge_filters_into(a, b, mutate) for a, b in zip(mine, theirs)): + return False + elif isinstance(mine, ast.AST) or isinstance(theirs, ast.AST): + if not isinstance(mine, ast.AST) or not isinstance(theirs, ast.AST): + return False + if not _merge_filters_into(mine, theirs, mutate): + return False + elif mine != theirs: + return False + + return True + + +def _merge_filters(left: ast.expr, right: ast.expr) -> ast.expr | None: + """Combine two paths that differ only in which filter conditions they met. + + Each path through a filter records the conjunction that got it to the + element; the filter as a whole is the disjunction over all such paths. + Returns ``None`` when the results differ by more than their filters. + """ + # A marker anywhere means some conditional expression is still unresolved, + # and unresolved arms must be spliced before anything can be OR-ed. + if any(isinstance(n, Skipped) for n in ast.walk(left)): + return None + if any(isinstance(n, Skipped) for n in ast.walk(right)): + return None + + if not _merge_filters_into(left, right, mutate=False): + return None + + merged = copy.deepcopy(left) + return merged if _merge_filters_into(merged, right) else None + + +def _skipped_offset(key: str) -> int: + """Sort key for `.SKIPPED_` markers, so merging is deterministic.""" + return int(key.rsplit("_", 1)[-1]) + + +def _merge_at_ifexp(left: ast.expr, right: ast.expr) -> ast.expr: + """ + Merge two expression ASTs obtained from two branches of symbolic execution. + """ + if isinstance(left, ast.Constant) and left.value is None: + return copy.deepcopy(right) + elif isinstance(right, ast.Constant) and right.value is None: + return copy.deepcopy(left) + + assert type(left) == type(right) + + lb, rb = BranchIdentifier(), BranchIdentifier() + lb.visit(left) + rb.visit(right) + + # A conditional expression: each path filled in one arm and left a marker in + # the other, so splice the two together. Sorted for determinism -- set + # iteration order over the marker names varies with PYTHONHASHSEED. + common_keys = set(lb.branching) & set(rb.branching) + differing = [ + key + for key in sorted(common_keys, key=_skipped_offset) + if lb.branching[key].testval != rb.branching[key].testval + ] + + # Only copy once it is known there is something to splice; this runs for + # every candidate pair of paths, most of which have nothing in common. + merged: ast.expr = copy.deepcopy(left) if differing else left + for key in differing: + visited = ReplaceSkipped(key, rb.branching[key].value).visit(merged) + assert isinstance(visited, ast.expr) + merged = visited + spliced = bool(differing) + + # The paths may *also* have satisfied different filter conditions on the way + # to the element, so combine those too rather than picking one arbitrarily. + combined = _merge_filters(merged, right) + if combined is not None: + return combined + if spliced: + return merged + + if ast.dump(left) == ast.dump(right): + return copy.deepcopy(left) + + raise ValueError("No differing branches found to merge") + + +def _specialization_guard_edge( + state: ReconstructionState, instr: dis.Instruction +) -> BranchEdge | None: + """The edge past a 3.14 inlined-builtin guard, or None if this isn't one. + + 3.14 may inline `any()`/`all()` over a generator, guarding the inlined code + with `loaded_name is ` and keeping an ordinary call on the + other edge. Only that other edge still contains the call to reconstruct, so + the guard is treated as though the identity test failed. + """ + if not state.stack: + return None + condition = state.stack[-1] + if not isinstance(condition, ast.Compare): + return None + if not any(isinstance(c, CommonConstant) for c in condition.comparators): + return None + + # Follow the edge taken when the identity test is false. + if instr.opname == "POP_JUMP_IF_FALSE": + return BranchEdge.TAKE_JUMP + elif instr.opname == "POP_JUMP_IF_TRUE": + return BranchEdge.FALL_THROUGH + else: + return None + + +def _merge_all(results: list[ast.expr]) -> ast.expr: + """Combine every path's result into one expression. + + Merging is not associative: a path that still carries an unfilled + conditional-expression arm can only combine with the path that took the + other arm, which need not be its neighbour. So rather than folding in + order, repeatedly merge whichever pair actually combines. + """ + pending = list(results) + while len(pending) > 1: + for i, j in itertools.combinations(range(len(pending)), 2): + try: + merged = _merge_at_ifexp(pending[i], pending[j]) + except (ValueError, AssertionError): + continue + pending = [merged] + [p for k, p in enumerate(pending) if k not in (i, j)] + break + else: + raise ValueError("Could not merge the paths of symbolic execution") + + return pending[0] + + +def _decide_branch( + state: ReconstructionState, instr: dis.Instruction, edge: BranchEdge +) -> ReconstructionState: + """Record which edge of ``instr`` the path being explored takes.""" + return replace(state, branches={**state.branches, instr.offset: edge}) + + +def _symbolic_exec(code: types.CodeType, freevars: dict[str, ast.expr]) -> ast.expr: + """Execute bytecode symbolically, following control flow.""" + continuations: list[ReconstructionState] = [ + ReconstructionState( + code=code, + instruction=next(iter(dis.get_instructions(code))), + stack=[Placeholder(), Placeholder()] + if current_version() == PythonVersion.PY_312 + and code.co_flags & inspect.CO_GENERATOR + else [Placeholder()], + freevars=freevars, + ) + ] + + results: list[ast.expr] = [] + + while continuations: + state = continuations.pop() + while not state.finished: + instr = state.instruction + if instr.opname in BRANCH_OPS and instr.offset not in state.branches: + forced = _specialization_guard_edge(state, instr) + if forced is not None: + state = _decide_branch(state, instr, forced) + else: + _, live = _classify_branch(state, instr) + # Explore the first live edge now; queue the rest for later. + continuations.extend( + _decide_branch(state, instr, edge) for edge in live[1:] + ) + state = _decide_branch(state, instr, live[0]) + + state = OP_HANDLERS[state.instruction.opname](state, state.instruction) + results.append(state.result) + + assert results, "No results from symbolic execution" + result = _merge_all(results) + assert not any(isinstance(n, Skipped) for n in ast.walk(result)), ( + "Every conditional expression arm must have been filled in" + ) + return result + + +# ============================================================================ +# GENERATOR COMPREHENSION HANDLERS +# ============================================================================ + + +@register_handler("RETURN_GENERATOR", version=PythonVersion.PY_312) +def handle_return_generator_312( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # RETURN_GENERATOR is the first instruction in generator expressions in Python 3.13+ + assert len(state.stack) == 2 and all( + isinstance(x, Null | Placeholder) for x in state.stack + ), "RETURN_GENERATOR must be the first instruction" + new_result = ast.GeneratorExp(elt=Placeholder(), generators=[]) + return replace(state, stack=[new_result, Null()]) + + +@register_handler("RETURN_GENERATOR", version=PythonVersion.PY_313) +@register_handler("RETURN_GENERATOR", version=PythonVersion.PY_314) +def handle_return_generator( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # RETURN_GENERATOR is the first instruction in generator expressions in Python 3.13+ + assert len(state.stack) == 1 and isinstance(state.stack[0], Null | Placeholder), ( + "RETURN_GENERATOR must be the first instruction" + ) + return replace( + state, stack=[ast.GeneratorExp(elt=Placeholder(), generators=[]), Null()] + ) + + +@register_handler("YIELD_VALUE", version=PythonVersion.PY_312) +@register_handler("YIELD_VALUE", version=PythonVersion.PY_313) +@register_handler("YIELD_VALUE", version=PythonVersion.PY_314) +def handle_yield_value( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # YIELD_VALUE pops a value from the stack and yields it + # This is the expression part of the generator + assert isinstance(state.result, Placeholder) + new_result = copy.deepcopy(state.stack[0]) + assert isinstance(new_result, ast.GeneratorExp), ( + "YIELD_VALUE must be called after RETURN_GENERATOR" + ) + assert len(new_result.generators) > 0, "YIELD_VALUE should have generators" + assert any(isinstance(x, Placeholder) for x in ast.walk(new_result.elt)) + new_result.elt = ReplacePlaceholder(ensure_ast(state.stack[-1])).visit( + new_result.elt + ) + new_stack = [new_result] + state.stack[1:] + return replace(state, stack=new_stack, result=new_result) + + +# ============================================================================ +# LIST COMPREHENSION HANDLERS +# ============================================================================ + + +@register_handler("BUILD_LIST", version=PythonVersion.PY_312) +@register_handler("BUILD_LIST", version=PythonVersion.PY_313) +@register_handler("BUILD_LIST", version=PythonVersion.PY_314) +def handle_build_list( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert instr.arg is not None + size: int = instr.arg + + if size == 0: + # Check if this looks like the start of a list comprehension pattern + # In nested comprehensions, BUILD_LIST(0) starts a new list comprehe + new_ret = ast.ListComp(elt=Placeholder(), generators=[]) + new_stack = state.stack + [new_ret] + return replace(state, stack=new_stack) + else: + # BUILD_LIST with elements - create a regular list + elements = [ensure_ast(elem) for elem in state.stack[-size:]] + new_stack = state.stack[:-size] + elt_node = ast.List(elts=elements, ctx=ast.Load()) + new_stack = new_stack + [elt_node] + return replace(state, stack=new_stack) + + +@register_handler("LIST_APPEND", version=PythonVersion.PY_312) +@register_handler("LIST_APPEND", version=PythonVersion.PY_313) +@register_handler("LIST_APPEND", version=PythonVersion.PY_314) +def handle_list_append( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert isinstance(state.stack[-instr.argval - 1], ast.ListComp) + + # add the body to the comprehension + comp: ast.ListComp = copy.deepcopy(state.stack[-instr.argval - 1]) + assert any(isinstance(x, Placeholder) for x in ast.walk(comp.elt)) + comp.elt = ReplacePlaceholder(state.stack[-1]).visit(comp.elt) + + # swap the return value + new_stack = state.stack[:-1] + new_stack[-instr.argval] = comp + + return replace(state, stack=new_stack) + + +# ============================================================================ +# SET COMPREHENSION HANDLERS +# ============================================================================ + + +@register_handler("BUILD_SET", version=PythonVersion.PY_312) +@register_handler("BUILD_SET", version=PythonVersion.PY_313) +@register_handler("BUILD_SET", version=PythonVersion.PY_314) +def handle_build_set( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert instr.arg is not None + size: int = instr.arg + + if size == 0: + new_result = ast.SetComp(elt=Placeholder(), generators=[]) + new_stack = state.stack + [new_result] + return replace(state, stack=new_stack) + else: + elements = [ensure_ast(elem) for elem in state.stack[-size:]] + new_stack = state.stack[:-size] + elt_node = ast.Set(elts=elements) + new_stack = new_stack + [elt_node] + return replace(state, stack=new_stack) + + +@register_handler("SET_ADD", version=PythonVersion.PY_312) +@register_handler("SET_ADD", version=PythonVersion.PY_313) +@register_handler("SET_ADD", version=PythonVersion.PY_314) +def handle_set_add( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert isinstance(state.stack[-instr.argval - 1], ast.SetComp) + + # add the body to the comprehension + comp: ast.SetComp = copy.deepcopy(state.stack[-instr.argval - 1]) + assert any(isinstance(x, Placeholder) for x in ast.walk(comp.elt)) + comp.elt = ReplacePlaceholder(state.stack[-1]).visit(comp.elt) + + # swap the return value + new_stack = state.stack[:-1] + new_stack[-instr.argval] = comp + + return replace(state, stack=new_stack) + + +# ============================================================================ +# DICT COMPREHENSION HANDLERS +# ============================================================================ + + +@register_handler("BUILD_MAP", version=PythonVersion.PY_312) +@register_handler("BUILD_MAP", version=PythonVersion.PY_313) +@register_handler("BUILD_MAP", version=PythonVersion.PY_314) +def handle_build_map( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert instr.arg is not None + size: int = instr.arg + + if size == 0: + new_result = ast.DictComp(key=Placeholder(), value=Placeholder(), generators=[]) + new_stack = state.stack + [new_result] + return replace(state, stack=new_stack) + else: + # Pop key-value pairs for the dict. They sit on the stack in source + # order -- key first, then value -- so they must be read from the + # bottom of that slice up: a later duplicate key has to keep winning, + # and side effects have to happen in the order they were written. + pairs = state.stack[-2 * size :] + keys: list[ast.expr | None] = [ensure_ast(pairs[2 * i]) for i in range(size)] + values = [ensure_ast(pairs[2 * i + 1]) for i in range(size)] + new_stack = state.stack[: -2 * size] + + # Create dict AST + dict_node = ast.Dict(keys=keys, values=values) + new_stack = new_stack + [dict_node] + return replace(state, stack=new_stack) + + +@register_handler("MAP_ADD", version=PythonVersion.PY_312) +@register_handler("MAP_ADD", version=PythonVersion.PY_313) +@register_handler("MAP_ADD", version=PythonVersion.PY_314) +def handle_map_add( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert isinstance(state.stack[-instr.argval - 2], ast.DictComp) + + # add the body to the comprehension + comp: ast.DictComp = copy.deepcopy(state.stack[-instr.argval - 2]) + assert any(isinstance(x, Placeholder) for x in ast.walk(comp.key)) + assert any(isinstance(x, Placeholder) for x in ast.walk(comp.value)) + comp.key = ReplacePlaceholder(state.stack[-2]).visit(comp.key) + comp.value = ReplacePlaceholder(state.stack[-1]).visit(comp.value) + + # swap the return value + new_stack = state.stack[:-2] + new_stack[-instr.argval] = comp + + return replace(state, stack=new_stack) + + +# ============================================================================ +# LOOP CONTROL HANDLERS +# ============================================================================ + + +@register_handler("RETURN_VALUE", version=PythonVersion.PY_312) +@register_handler("RETURN_VALUE", version=PythonVersion.PY_313) +def handle_return_value( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert isinstance(state.result, Placeholder) + assert len(state.stack) == 2 + new_result = ReplacePlaceholder(ensure_ast(state.stack[-1])).visit(state.stack[-2]) + new_stack = state.stack[:-1] + return replace(state, stack=new_stack, result=new_result) + + +def _unyielded_comprehension(state: ReconstructionState) -> CompExp | None: + """The comprehension of a body the compiler proved unreachable, if any. + + An always-false filter lets the compiler drop the whole body: the loop is + still walked, but nothing is ever yielded or appended, so no element + expression survives. The partly built comprehension still carries its + generators, so it can be rebuilt with a filter that is never satisfied -- + which iterates exactly as the original did and produces nothing. + """ + for item in reversed(state.stack): + if not isinstance(item, CompExp) or not item.generators: + continue + + element = item.value if isinstance(item, ast.DictComp) else item.elt + if not isinstance(element, Placeholder): + continue + + unreachable = copy.deepcopy(item) + never = ast.Constant(value=None) + if isinstance(unreachable, ast.DictComp): + unreachable.key, unreachable.value = never, copy.deepcopy(never) + else: + unreachable.elt = never + unreachable.generators[-1].ifs = [ast.Constant(value=False)] + return unreachable + + return None + + +@register_handler("RETURN_VALUE", version=PythonVersion.PY_314) +def handle_return_value_314( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # Two things changed in 3.14. RETURN_CONST is gone, so a generator's + # trailing `return None` now arrives as LOAD_CONST + RETURN_VALUE; and + # RETURN_VALUE's stack effect is 0 rather than -1, the returned value being + # discarded along with the frame. The value therefore stays on the stack. + if not isinstance(state.result, Placeholder): + assert ( + isinstance(state.stack[-1], ast.Constant) and state.stack[-1].value is None + ), "A generator may only fall off the end returning None" + return state + + unreachable = _unyielded_comprehension(state) + if unreachable is not None: + return replace(state, result=unreachable) + + assert len(state.stack) == 2 + new_result = ReplacePlaceholder(ensure_ast(state.stack[-1])).visit(state.stack[-2]) + return replace(state, result=new_result) + + +@register_handler("RETURN_CONST", version=PythonVersion.PY_312) +@register_handler("RETURN_CONST", version=PythonVersion.PY_313) +def handle_return_const( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # RETURN_CONST returns a constant value (replaces some LOAD_CONST + RETURN_VALUE patterns) + # Similar to RETURN_VALUE but with a constant + if isinstance(state.result, Placeholder): + unreachable = _unyielded_comprehension(state) + if unreachable is not None: + return replace(state, result=unreachable) + return replace(state, result=ensure_ast(instr.argval)) + else: + assert instr.argval is None + return state + + +@register_handler("FOR_ITER", version=PythonVersion.PY_312) +@register_handler("FOR_ITER", version=PythonVersion.PY_313) +@register_handler("FOR_ITER", version=PythonVersion.PY_314) +def handle_for_iter( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # FOR_ITER pops an iterator from the stack and pushes the next item + # If the iterator is exhausted, it jumps to the target instruction + assert len(state.stack) > 0, "FOR_ITER must have an iterator on the stack" + + if state.loops.get(instr.offset, 0) > 0: + return replace(state, stack=state.stack + [Null()]) + + # The iterator should be on top of stack + iterator: ast.expr = state.stack[-1] + + for pos, item in zip(reversed(range(len(state.stack))), reversed(state.stack)): + if not isinstance(item, CompExp): + continue + + element = item.value if isinstance(item, ast.DictComp) else item.elt + new_result = copy.deepcopy(item) + + if isinstance(element, Placeholder): + loop_iter = ensure_ast(iterator) + elif isinstance(element, ast.IfExp) and any( + isinstance(x, Placeholder) for x in ast.walk(element) + ): + # A conditional expression was being built up in the element slot, + # but it turned out to be this loop's iterable, as in + # `for y in (a if c else b)`. Move it back out and plug the value + # this path produced into the arm still awaiting one. + if isinstance(new_result, ast.DictComp): + new_result.key, new_result.value = Placeholder(), Placeholder() + else: + new_result.elt = Placeholder() + + plugged = ReplacePlaceholder(ensure_ast(iterator)).visit( + copy.deepcopy(element) + ) + assert isinstance(plugged, ast.expr) + loop_iter = plugged + else: + continue + + # The loop target is not named until the STORE_* that follows. + loop_info = ast.comprehension( + target=TargetHole(), iter=loop_iter, ifs=[], is_async=0 + ) + new_result.generators.append(loop_info) + new_stack = ( + state.stack[:pos] + + [new_result] + + state.stack[pos + 1 :] + + [loop_info.target] + ) + return replace(state, stack=new_stack) + + raise TypeError("FOR_ITER did not find partial comprehension on stack") + + +@register_handler("GET_ITER", version=PythonVersion.PY_312) +@register_handler("GET_ITER", version=PythonVersion.PY_313) +@register_handler("GET_ITER", version=PythonVersion.PY_314) +def handle_get_iter( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # GET_ITER converts the top stack item to an iterator + # For AST reconstruction, we typically don't need to change anything + # since the iterator will be used directly in the comprehension + return state + + +@register_handler("END_FOR", version=PythonVersion.PY_312) +def handle_end_for_312( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # END_FOR marks the end of a for loop, followed by POP_TOP (in 3.12) + new_stack = state.stack[:-2] + return replace(state, stack=new_stack) + + +@register_handler("END_FOR", version=PythonVersion.PY_313) +@register_handler("END_FOR", version=PythonVersion.PY_314) +def handle_end_for( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # END_FOR marks the end of a for loop - no action needed for AST reconstruction + new_stack = state.stack[:-1] + return replace(state, stack=new_stack) + + +@register_handler("RERAISE", version=PythonVersion.PY_312) +@register_handler("RERAISE", version=PythonVersion.PY_313) +@register_handler("RERAISE", version=PythonVersion.PY_314) +def handle_reraise( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # RERAISE re-raises an exception - generally ignore for AST reconstruction + return state + + +# ============================================================================ +# VARIABLE OPERATIONS HANDLERS +# ============================================================================ + + +def _literal_elements(value: ast.expr, count: int | None = None) -> list[ast.expr]: + """The elements of a literal sequence, for destructuring a known value.""" + assert isinstance(value, ast.Tuple | ast.List), ( + f"Cannot unpack {type(value).__name__}; expected a literal sequence" + ) + assert count is None or len(value.elts) == count, ( + f"Expected {count} values to unpack, got {len(value.elts)}" + ) + return [ensure_ast(element) for element in value.elts] + + +def _bind_local( + state: ReconstructionState, var_name: str, value: ast.expr +) -> ReconstructionState: + """Record that a local now stands for ``value``, popping it off the stack. + + Reached when a store is not filling in a loop target: 3.14 unrolls a + single-iteration loop over a literal, assigning its targets outright. The + loop is gone from the bytecode, so the names it bound are not in scope in + the reconstruction and their uses are substituted instead. + """ + bindings = {**state.bindings, var_name: ensure_ast(value)} + return replace(state, stack=state.stack[:-1], bindings=bindings) + + +def _read_local(state: ReconstructionState, var_name: str) -> ast.expr: + """The expression a local name stands for at this point on this path.""" + if var_name == DummyIterName().id: + return DummyIterName() + elif var_name in state.bindings: + # Bound to a known expression rather than by a loop, so the name itself + # is not in scope in the reconstruction; use what it was bound to. + return copy.deepcopy(state.bindings[var_name]) + else: + return ast.Name(id=var_name, ctx=ast.Load()) + + +@register_handler("LOAD_FAST", version=PythonVersion.PY_312) +@register_handler("LOAD_FAST", version=PythonVersion.PY_313) +@register_handler("LOAD_FAST", version=PythonVersion.PY_314) +@register_handler("LOAD_FAST_CHECK", version=PythonVersion.PY_312) +@register_handler("LOAD_FAST_CHECK", version=PythonVersion.PY_313) +@register_handler("LOAD_FAST_CHECK", version=PythonVersion.PY_314) +def handle_load_fast( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_FAST_CHECK differs only in raising when the local is unbound, which + # says nothing about the expression being reconstructed. + return replace(state, stack=state.stack + [_read_local(state, instr.argval)]) + + +@register_handler("LOAD_DEREF", version=PythonVersion.PY_312) +@register_handler("LOAD_DEREF", version=PythonVersion.PY_313) +@register_handler("LOAD_DEREF", version=PythonVersion.PY_314) +def handle_load_deref( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_DEREF loads a value out of a cell. When the cell is one this code + # object *creates* (a comprehension target captured by a nested lambda, say) + # the name is bound inside the reconstructed tree too, so it can stand. A + # free variable is different: its cell belongs to an enclosing scope that + # the reconstructed tree does not reproduce, so a bare name would silently + # become a global lookup. Write the captured value in instead. + var_name = instr.argval + if var_name in state.code.co_freevars and var_name in state.freevars: + loaded: ast.expr = copy.deepcopy(state.freevars[var_name]) + else: + loaded = ast.Name(id=var_name, ctx=ast.Load()) + return replace(state, stack=state.stack + [loaded]) + + +@register_handler("LOAD_CLOSURE", version=PythonVersion.PY_312) +@register_handler("LOAD_CLOSURE", version=PythonVersion.PY_313) +@register_handler("LOAD_CLOSURE", version=PythonVersion.PY_314) +def handle_load_closure( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_CLOSURE loads a closure variable + var_name = instr.argval + new_stack = state.stack + [ast.Name(id=var_name, ctx=ast.Load())] + return replace(state, stack=new_stack) + + +@register_handler("LOAD_CONST", version=PythonVersion.PY_312) +@register_handler("LOAD_CONST", version=PythonVersion.PY_313) +@register_handler("LOAD_CONST", version=PythonVersion.PY_314) +def handle_load_const( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + const_value = instr.argval + # A nested lambda or comprehension arrives as a code object, and its own + # free variables reach back through this scope to the same cells, so the + # captured values have to travel with it. + loaded = ( + _reconstruct_code(const_value, state.freevars) + if isinstance(const_value, types.CodeType) + else ensure_ast(const_value) + ) + return replace(state, stack=state.stack + [loaded]) + + +@register_handler("LOAD_GLOBAL", version=PythonVersion.PY_312) +@register_handler("LOAD_GLOBAL", version=PythonVersion.PY_313) +@register_handler("LOAD_GLOBAL", version=PythonVersion.PY_314) +def handle_load_global( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + global_name = instr.argval + + if instr.argrepr.endswith(" + NULL"): + new_stack = state.stack + [ast.Name(id=global_name, ctx=ast.Load()), Null()] + elif instr.argrepr.startswith("NULL + "): + new_stack = state.stack + [Null(), ast.Name(id=global_name, ctx=ast.Load())] + else: + new_stack = state.stack + [ast.Name(id=global_name, ctx=ast.Load())] + return replace(state, stack=new_stack) + + +@register_handler("LOAD_NAME", version=PythonVersion.PY_312) +@register_handler("LOAD_NAME", version=PythonVersion.PY_313) +@register_handler("LOAD_NAME", version=PythonVersion.PY_314) +def handle_load_name( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_NAME is similar to LOAD_GLOBAL but for names in the global namespace + name = instr.argval + new_stack = state.stack + [ast.Name(id=name, ctx=ast.Load())] + return replace(state, stack=new_stack) + + +def _is_assignment_expression(state: ReconstructionState) -> bool: + """Is the top of the stack a COPY of the value beneath it? + + That duplication is how an assignment expression keeps its value after + binding it: `COPY 1` then a `STORE_*`. `handle_copy` pushes the very same + node, so identity is what distinguishes it from two equal-looking values. + """ + return len(state.stack) >= 2 and state.stack[-1] is state.stack[-2] + + +def _handle_assignment_expression( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + """Rebuild `(name := value)` from the COPY/STORE pair that implements it.""" + target = ast.Name(id=instr.argval, ctx=ast.Store()) + named = ast.NamedExpr(target=target, value=ensure_ast(state.stack[-1])) + return replace(state, stack=state.stack[:-2] + [named]) + + +@register_handler("STORE_GLOBAL", version=PythonVersion.PY_312) +@register_handler("STORE_GLOBAL", version=PythonVersion.PY_313) +@register_handler("STORE_GLOBAL", version=PythonVersion.PY_314) +def handle_store_global( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # A comprehension has no globals of its own, so the only way it stores one + # is an assignment expression, which binds in the enclosing scope. + assert _is_assignment_expression(state), ( + "STORE_GLOBAL outside an assignment expression" + ) + return _handle_assignment_expression(state, instr) + + +@register_handler("STORE_DEREF", version=PythonVersion.PY_312) +@register_handler("STORE_DEREF", version=PythonVersion.PY_313) +@register_handler("STORE_DEREF", version=PythonVersion.PY_314) +def handle_store_deref( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # STORE_DEREF stores into a closure variable: either an assignment + # expression binding in an enclosing function, or a loop target that an + # inner comprehension captures. + if _is_assignment_expression(state): + return _handle_assignment_expression(state, instr) + return handle_store_fast(state, instr) + + +@register_handler("STORE_FAST", version=PythonVersion.PY_312) +@register_handler("STORE_FAST", version=PythonVersion.PY_313) +@register_handler("STORE_FAST", version=PythonVersion.PY_314) +def handle_store_fast( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + if _is_assignment_expression(state): + # An assignment expression whose target is local to this code object, + # as in a comprehension inlined into the lambda that binds the name. + return _handle_assignment_expression(state, instr) + + if isinstance(state.stack[-1], ast.Name) and state.stack[-1].id == instr.argval: + # If the variable is already on the stack, we can skip adding it again + # This is common in nested comprehensions where the same variable is reused + return replace(state, stack=state.stack[:-1]) + + if not isinstance(state.stack[-1], TargetHole): + return _bind_local(state, instr.argval, state.stack[-1]) + + new_stack = _bind_target_hole( + state.stack, state.stack[-1], ast.Name(id=instr.argval, ctx=ast.Store()) + ) + return replace(state, stack=new_stack[:-1]) + + +@register_handler("STORE_FAST_LOAD_FAST", version=PythonVersion.PY_313) +@register_handler("STORE_FAST_LOAD_FAST", version=PythonVersion.PY_314) +def handle_store_fast_load_fast( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # STORE_FAST_LOAD_FAST stores and then loads the same variable (optimization) + # The instruction has two names: store_name and load_name + # In Python 3.13, this is often used for loop variables + + # In Python 3.13, the instruction argument contains both names + # argval should be a tuple (store_name, load_name) + assert isinstance(instr.argval, tuple) + store_name, load_name = instr.argval + + if _is_assignment_expression(state): + # `(name := value)` whose result is read straight back, as in + # `(z := w) + z`. The duplicate becomes the assignment expression and + # the reload becomes a plain reference to the name just bound. + named = ast.NamedExpr( + target=ast.Name(id=store_name, ctx=ast.Store()), + value=ensure_ast(state.stack[-1]), + ) + reload = ast.Name(id=load_name, ctx=ast.Load()) + return replace(state, stack=state.stack[:-2] + [named, reload]) + + if not isinstance(state.stack[-1], TargetHole): + # A plain assignment followed by a load, as 3.14 emits when it unrolls + # a single-iteration loop over a literal. + bound = _bind_local(state, store_name, state.stack[-1]) + return replace(bound, stack=bound.stack + [_read_local(bound, load_name)]) + + new_stack = _bind_target_hole( + state.stack, state.stack[-1], ast.Name(id=store_name, ctx=ast.Store()) + ) + new_var = ast.Name(id=load_name, ctx=ast.Load()) + return replace(state, stack=new_stack[:-1] + [new_var]) + + +@register_handler("STORE_FAST_STORE_FAST", version=PythonVersion.PY_313) +@register_handler("STORE_FAST_STORE_FAST", version=PythonVersion.PY_314) +def handle_store_fast_store_fast( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # STORE_FAST_STORE_FAST stores STACK[-1] into the first named variable and + # STACK[-2] into the second. It is emitted for unpacking targets, so both + # values are loop-target holes belonging to the same comprehension. + assert isinstance(instr.argval, tuple) + first_name, second_name = instr.argval + + if not isinstance(state.stack[-1], TargetHole): + # Not loop targets: a pair of plain assignments, as 3.14 emits when it + # unrolls a single-iteration loop over a literal. + bound = _bind_local(state, first_name, state.stack[-1]) + return _bind_local(bound, second_name, bound.stack[-1]) + + new_stack = _bind_target_hole( + state.stack, state.stack[-1], ast.Name(id=first_name, ctx=ast.Store()) + ) + new_stack = _bind_target_hole( + new_stack, new_stack[-2], ast.Name(id=second_name, ctx=ast.Store()) + ) + return replace(state, stack=new_stack[:-2]) + + +@register_handler("LOAD_FAST_AND_CLEAR", version=PythonVersion.PY_312) +@register_handler("LOAD_FAST_AND_CLEAR", version=PythonVersion.PY_313) +@register_handler("LOAD_FAST_AND_CLEAR", version=PythonVersion.PY_314) +def handle_load_fast_and_clear( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_FAST_AND_CLEAR pushes a local variable onto the stack and clears it + # For AST reconstruction, we treat this the same as LOAD_FAST + return replace(state, stack=state.stack + [_read_local(state, instr.argval)]) + + +@register_handler("LOAD_FAST_LOAD_FAST", version=PythonVersion.PY_313) +@register_handler("LOAD_FAST_LOAD_FAST", version=PythonVersion.PY_314) +def handle_load_fast_load_fast( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_FAST_LOAD_FAST loads two variables (optimization in Python 3.13) + # The instruction argument contains both variable names + if isinstance(instr.argval, tuple): + var1, var2 = instr.argval + else: + # Fallback: assume both names are the same + var1 = var2 = instr.argval + + new_stack = state.stack + [_read_local(state, var1), _read_local(state, var2)] + + return replace(state, stack=new_stack) + + +@register_handler("MAKE_CELL", version=PythonVersion.PY_312) +@register_handler("MAKE_CELL", version=PythonVersion.PY_313) +@register_handler("MAKE_CELL", version=PythonVersion.PY_314) +def handle_make_cell( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # MAKE_CELL creates a new cell in slot i for closure variables + # This is used when variables from outer scopes are captured by inner scopes + # For AST reconstruction purposes, this is just a variable scoping mechanism + # that we can ignore since the AST doesn't track low-level closure details + return state + + +@register_handler("COPY_FREE_VARS", version=PythonVersion.PY_312) +@register_handler("COPY_FREE_VARS", version=PythonVersion.PY_313) +@register_handler("COPY_FREE_VARS", version=PythonVersion.PY_314) +def handle_copy_free_vars( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # COPY_FREE_VARS copies n free (closure) variables from the closure into the frame + # This removes the need for special code on the caller's side when calling closures + # For AST reconstruction purposes, this is just a variable scoping mechanism + # that we can ignore since the AST doesn't track runtime variable management + return state + + +# ============================================================================ +# STACK MANAGEMENT HANDLERS +# ============================================================================ + + +@register_handler("POP_TOP", version=PythonVersion.PY_312) +@register_handler("POP_TOP", version=PythonVersion.PY_313) +@register_handler("POP_TOP", version=PythonVersion.PY_314) +def handle_pop_top( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # POP_TOP removes the top item from the stack + # In generators, often used after YIELD_VALUE + # Also used to clean up the duplicated middle value in failed chained comparisons + new_stack = state.stack[:-1] + return replace(state, stack=new_stack) + + +# Python 3.13 replacement for stack manipulation +@register_handler("SWAP", version=PythonVersion.PY_312) +@register_handler("SWAP", version=PythonVersion.PY_313) +@register_handler("SWAP", version=PythonVersion.PY_314) +def handle_swap( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # SWAP exchanges the top two stack items (replaces ROT_TWO in many cases) + assert instr.arg is not None + depth = instr.arg + stack_size = len(state.stack) + + if depth > stack_size: + # Not enough items on stack - this might be a pattern where some items were optimized away + # For AST reconstruction, we can often ignore certain stack manipulations + return state + + # For other depths, swap TOS with the item at specified depth + assert depth <= stack_size, f"SWAP depth {depth} exceeds stack size {stack_size}" + idx = stack_size - depth + new_stack = state.stack.copy() + new_stack[-1], new_stack[idx] = new_stack[idx], new_stack[-1] + return replace(state, stack=new_stack) + + +@register_handler("COPY", version=PythonVersion.PY_312) +@register_handler("COPY", version=PythonVersion.PY_313) +@register_handler("COPY", version=PythonVersion.PY_314) +def handle_copy( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # COPY duplicates the item at the specified depth + assert instr.arg is not None + depth = instr.arg + stack_size = len(state.stack) + if depth > stack_size: + raise ValueError(f"COPY depth {depth} exceeds stack size {stack_size}") + idx = stack_size - depth + copied_item = state.stack[idx] + new_stack = state.stack + [copied_item] + return replace(state, stack=new_stack) + + +@register_handler("PUSH_NULL", version=PythonVersion.PY_312) +@register_handler("PUSH_NULL", version=PythonVersion.PY_313) +@register_handler("PUSH_NULL", version=PythonVersion.PY_314) +def handle_push_null( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + return replace(state, stack=state.stack + [Null()]) + + +# ============================================================================ +# BINARY ARITHMETIC/LOGIC OPERATION HANDLERS +# ============================================================================ + + +def handle_binop( + op: ast.operator, state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + right = ensure_ast(state.stack[-1]) + left = ensure_ast(state.stack[-2]) + new_stack = state.stack[:-2] + [ast.BinOp(left=left, op=op, right=right)] + return replace(state, stack=new_stack) + + +# Python 3.12+ BINARY_OP handler +@register_handler("BINARY_OP", version=PythonVersion.PY_312) +@register_handler("BINARY_OP", version=PythonVersion.PY_313) +def handle_binary_op( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # BINARY_OP in Python 3.12+ consolidates all binary operations + # The operation type is determined by the instruction argument + assert instr.arg is not None + + # Map argument values to AST operators based on Python 3.12+ implementation + op_map: collections.abc.Mapping[int, ast.operator] = { + 0: ast.Add(), # + + 1: ast.BitAnd(), # & + 2: ast.FloorDiv(), # // + 3: ast.LShift(), # << + 4: ast.MatMult(), # @ + 5: ast.Mult(), # * + 6: ast.Mod(), # % + 7: ast.BitOr(), # | + 8: ast.Pow(), # ** + 9: ast.RShift(), # >> + 10: ast.Sub(), # - + 11: ast.Div(), # / + 12: ast.BitXor(), # ^ + } + + op = op_map.get(instr.arg) + if op is None: + raise TypeError(f"Unknown binary operation: {instr.arg}") + + return handle_binop(op, state, instr) + + +# 3.14 folded subscripting into BINARY_OP; `dis._nb_ops` names the oparg +# NB_SUBSCR. Looked up rather than hard-coded, since it sits past the in-place +# operators and so moves whenever one is added. +_NB_OPS: list[tuple[str, str]] = getattr(dis, "_nb_ops", []) +NB_SUBSCR: int | None = next( + (i for i, (name, _) in enumerate(_NB_OPS) if name == "NB_SUBSCR"), None +) + + +@register_handler("BINARY_OP", version=PythonVersion.PY_314) +def handle_binary_op_314( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # As in 3.13, except that BINARY_OP now also implements `a[b]`, which used + # to be its own BINARY_SUBSCR opcode. + if instr.arg is not None and instr.arg == NB_SUBSCR: + return handle_binary_subscr(state, instr) + return handle_binary_op(state, instr) + + +@register_handler("LOAD_SMALL_INT", version=PythonVersion.PY_314) +def handle_load_small_int( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_SMALL_INT pushes an int in range(256) held in the oparg itself, + # rather than going through co_consts. + assert isinstance(instr.argval, int) + return replace(state, stack=state.stack + [ensure_ast(instr.argval)]) + + +@register_handler("LOAD_FAST_BORROW", version=PythonVersion.PY_314) +def handle_load_fast_borrow( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # A borrowed reference differs only in ownership, which the AST does not + # model, so this is LOAD_FAST as far as reconstruction is concerned. + return handle_load_fast(state, instr) + + +@register_handler("LOAD_FAST_BORROW_LOAD_FAST_BORROW", version=PythonVersion.PY_314) +def handle_load_fast_borrow_load_fast_borrow( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + return handle_load_fast_load_fast(state, instr) + + +@register_handler("LOAD_COMMON_CONSTANT", version=PythonVersion.PY_314) +def handle_load_common_constant( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # Pushes one of a small hardcoded set of constants. In a comprehension this + # only shows up in the guard of an inlined builtin; see CommonConstant. + name = getattr(instr.argval, "__name__", str(instr.argval)) + return replace(state, stack=state.stack + [CommonConstant(id=name)]) + + +@register_handler("NOT_TAKEN", version=PythonVersion.PY_314) +def handle_not_taken( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # A no-op marking the not-taken edge of a branch for sys.monitoring. + return state + + +@register_handler("POP_ITER", version=PythonVersion.PY_314) +def handle_pop_iter( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # POP_ITER discards the exhausted iterator that FOR_ITER left behind. In + # 3.13 the same cleanup was spelled END_FOR followed by POP_TOP. + return replace(state, stack=state.stack[:-1]) + + +# ============================================================================ +# UNARY OPERATION HANDLERS +# ============================================================================ + + +def handle_unary_op( + op: ast.unaryop, state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + operand = ensure_ast(state.stack[-1]) + new_stack = state.stack[:-1] + [ast.UnaryOp(op=op, operand=operand)] + return replace(state, stack=new_stack) + + +UNARY_OPS: dict[str, ast.unaryop] = { + "UNARY_NEGATIVE": ast.USub(), + "UNARY_INVERT": ast.Invert(), + "UNARY_NOT": ast.Not(), +} + +# These three behave identically in every dialect this module supports; 3.13's +# "requires an exact bool operand" note on UNARY_NOT constrains the operand, not +# the reconstruction. +for _opname, _op in UNARY_OPS.items(): + for _version in ( + PythonVersion.PY_312, + PythonVersion.PY_313, + PythonVersion.PY_314, + ): + register_handler( + _opname, functools.partial(handle_unary_op, _op), version=_version + ) + + +@register_handler("CONVERT_VALUE", version=PythonVersion.PY_313) +@register_handler("CONVERT_VALUE", version=PythonVersion.PY_314) +def handle_convert_value( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # CONVERT_VALUE applies a conversion to the value on top of stack + # Used for f-string conversions like !r, !s, !a + # The conversion type is stored in instr.arg: + # 0 = None, 1 = str (!s), 2 = repr (!r), 3 = ascii (!a) + assert len(state.stack) > 0, "CONVERT_VALUE requires a value on stack" + assert instr.arg is not None, "CONVERT_VALUE requires conversion type" + + # Wrap the value with conversion information + value = state.stack[-1] + converted = ConvertedValue(value, instr.arg) + new_stack = state.stack[:-1] + [converted] + + return replace(state, stack=new_stack) + + +@register_handler("CALL_INTRINSIC_1", version=PythonVersion.PY_312) +@register_handler("CALL_INTRINSIC_1", version=PythonVersion.PY_313) +@register_handler("CALL_INTRINSIC_1", version=PythonVersion.PY_314) +def handle_call_intrinsic_1( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # CALL_INTRINSIC_1 calls an intrinsic function with one argument + if instr.argrepr == "INTRINSIC_LIST_TO_TUPLE": + assert isinstance(state.stack[-1], ast.List), ( + "Expected a list for LIST_TO_TUPLE" + ) + tuple_node = ast.Tuple(elts=state.stack[-1].elts, ctx=ast.Load()) + return replace(state, stack=state.stack[:-1] + [tuple_node]) + elif instr.argrepr == "INTRINSIC_UNARY_POSITIVE": + assert len(state.stack) > 0 + new_val = ast.UnaryOp(op=ast.UAdd(), operand=state.stack[-1]) + return replace(state, stack=state.stack[:-1] + [new_val]) + elif instr.argrepr == "INTRINSIC_STOPITERATION_ERROR": + return state + else: + raise TypeError(f"Unsupported generator intrinsic operation: {instr.argrepr}") + + +@register_handler("TO_BOOL", version=PythonVersion.PY_313) +@register_handler("TO_BOOL", version=PythonVersion.PY_314) +def handle_to_bool( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # TO_BOOL converts the top stack item to a boolean + # For AST reconstruction, we typically don't need an explicit bool() call + # since the boolean context is usually handled by the conditional jump that follows + # However, for some cases we might need to preserve the explicit conversion + + # For now, leave the value as-is since the jump instruction will handle the boolean logic + return state + + +# ============================================================================ +# COMPARISON OPERATION HANDLERS +# ============================================================================ + +CMP_OPMAP: dict[str, ast.cmpop] = { + "<": ast.Lt(), + "<=": ast.LtE(), + ">": ast.Gt(), + ">=": ast.GtE(), + "==": ast.Eq(), + "!=": ast.NotEq(), +} + + +@register_handler("COMPARE_OP", version=PythonVersion.PY_312) +@register_handler("COMPARE_OP", version=PythonVersion.PY_313) +@register_handler("COMPARE_OP", version=PythonVersion.PY_314) +def handle_compare_op( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert instr.arg is not None and instr.argval in dis.cmp_op, ( + f"Unsupported comparison operation: {instr.argval}" + ) + + right = ensure_ast(state.stack[-1]) + left = ensure_ast(state.stack[-2]) + + # Map comparison operation codes to AST operators + op_name = instr.argval + compare_node = ast.Compare(left=left, ops=[CMP_OPMAP[op_name]], comparators=[right]) + new_stack = state.stack[:-2] + [compare_node] + return replace(state, stack=new_stack) + + +@register_handler("CONTAINS_OP", version=PythonVersion.PY_312) +@register_handler("CONTAINS_OP", version=PythonVersion.PY_313) +@register_handler("CONTAINS_OP", version=PythonVersion.PY_314) +def handle_contains_op( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + right = ensure_ast(state.stack[-1]) # Container + left = ensure_ast(state.stack[-2]) # Item to check + + # instr.arg determines if it's 'in' (0) or 'not in' (1) + op = ast.NotIn() if instr.arg else ast.In() + + compare_node = ast.Compare(left=left, ops=[op], comparators=[right]) + new_stack = state.stack[:-2] + [compare_node] + return replace(state, stack=new_stack) + + +@register_handler("IS_OP", version=PythonVersion.PY_312) +@register_handler("IS_OP", version=PythonVersion.PY_313) +@register_handler("IS_OP", version=PythonVersion.PY_314) +def handle_is_op( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + right = ensure_ast(state.stack[-1]) + left = ensure_ast(state.stack[-2]) + + # instr.arg determines if it's 'is' (0) or 'is not' (1) + op = ast.IsNot() if instr.arg else ast.Is() + + compare_node = ast.Compare(left=left, ops=[op], comparators=[right]) + new_stack = state.stack[:-2] + [compare_node] + return replace(state, stack=new_stack) + + +# ============================================================================ +# FUNCTION CALL HANDLERS +# ============================================================================ + + +@register_handler("KW_NAMES", version=PythonVersion.PY_312) +def handle_kw_names( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # KW_NAMES names the trailing arguments of the CALL that follows it. + # Python 3.13 replaced this pair with a single CALL_KW instruction. + assert isinstance(instr.argval, tuple), "KW_NAMES requires a tuple of names" + assert all(isinstance(name, str) for name in instr.argval) + assert state.kw_names is None, "KW_NAMES must be consumed by the following CALL" + return replace(state, kw_names=instr.argval) + + +@register_handler("CALL", version=PythonVersion.PY_312) +def handle_call_312( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # CALL in Python 3.12 handles both function and method calls + # Stack layout: [..., callable or self, callable or NULL] + assert instr.arg is not None + arg_count: int = instr.arg + + # Check if this is a method call (no NULL on top) + if isinstance(state.stack[-arg_count - 2], Null): + # Regular function call: [..., NULL, callable, *args] + func = ensure_ast(state.stack[-arg_count - 1]) + args = ( + [ensure_ast(arg) for arg in state.stack[-arg_count:]] + if arg_count > 0 + else [] + ) + new_stack = state.stack[: -arg_count - 2] + else: + # Method call: [..., callable, self, *args] + func = ensure_ast(state.stack[-arg_count - 2]) + self_arg = ensure_ast(state.stack[-arg_count - 1]) + remaining_args = ( + [ensure_ast(arg) for arg in state.stack[-arg_count:]] + if arg_count > 0 + else [] + ) + args = [self_arg] + remaining_args + new_stack = state.stack[: -arg_count - 2] + + # A preceding KW_NAMES names the trailing `len(kw_names)` positional slots. + keywords: list[ast.keyword] = [] + if state.kw_names is not None: + assert 0 < len(state.kw_names) <= arg_count + keywords = [ + ast.keyword(arg=name, value=value) + for name, value in zip(state.kw_names, args[-len(state.kw_names) :]) + ] + args = args[: -len(state.kw_names)] + + if isinstance(func, CompLambda): + assert len(args) == 1 and not keywords + return replace(state, stack=new_stack + [func.inline(args[0])], kw_names=None) + else: + # Create function call AST + call_node = ast.Call(func=func, args=args, keywords=keywords) + new_stack = new_stack + [call_node] + return replace(state, stack=new_stack, kw_names=None) + + +@register_handler("CALL", version=PythonVersion.PY_313) +@register_handler("CALL", version=PythonVersion.PY_314) +def handle_call( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # CALL pops function and arguments from stack (replaces CALL_FUNCTION in Python 3.13) + assert instr.arg is not None + arg_count: int = instr.arg + + func = ensure_ast(state.stack[-arg_count - 2]) + + # Pop arguments and function + args = ( + [ensure_ast(arg) for arg in state.stack[-arg_count:]] if arg_count > 0 else [] + ) + if not isinstance(state.stack[-arg_count - 1], Null): + args = [ensure_ast(state.stack[-arg_count - 1])] + args + + new_stack = state.stack[: -arg_count - 2] + if isinstance(func, CompLambda): + assert len(args) == 1 + return replace(state, stack=new_stack + [func.inline(args[0])]) + else: + # Create function call AST + call_node = ast.Call(func=func, args=args, keywords=[]) + new_stack = new_stack + [call_node] + return replace(state, stack=new_stack) + + +@register_handler("CALL_KW", version=PythonVersion.PY_313) +@register_handler("CALL_KW", version=PythonVersion.PY_314) +def handle_call_kw( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # CALL_KW pops function, arguments, and keyword names from stack + assert instr.arg is not None + arg_count: int = instr.arg + assert arg_count > 0, "CALL_KW requires at least one argument" + + func = ensure_ast(state.stack[-arg_count - 3]) + assert not isinstance(func, CompLambda | Null) + + kw_names = state.stack[-1] + assert isinstance(kw_names, ast.Tuple), "Expected a tuple of keyword names" + assert len(kw_names.elts) > 0, "Expected at least one keyword name" + + # Pop arguments, function, and keyword names + keywords = [] + for i, kw in enumerate(reversed(kw_names.elts)): + assert isinstance(kw, ast.Constant) and isinstance(kw.value, str) + keywords += [ast.keyword(arg=kw.value, value=ensure_ast(state.stack[-2 - i]))] + keywords.reverse() + + args = [ensure_ast(a) for a in state.stack[-arg_count - 1 : -len(keywords) - 1]] + if not isinstance(state.stack[-arg_count - 2], Null): + args = [ensure_ast(state.stack[-arg_count - 2])] + args + + # Create function call AST + call_node = ast.Call(func=func, args=args, keywords=keywords) + new_stack = state.stack[: -arg_count - 3] + [call_node] + return replace(state, stack=new_stack) + + +# Flags shared by MAKE_FUNCTION (3.12) and SET_FUNCTION_ATTRIBUTE (3.13) +MAKE_FUNCTION_DEFAULTS = 0x01 +MAKE_FUNCTION_KWDEFAULTS = 0x02 +MAKE_FUNCTION_ANNOTATIONS = 0x04 +MAKE_FUNCTION_CLOSURE = 0x08 +MAKE_FUNCTION_ANNOTATE = 0x10 # added in 3.14 +MAKE_FUNCTION_FLAGS = ( + MAKE_FUNCTION_DEFAULTS, + MAKE_FUNCTION_KWDEFAULTS, + MAKE_FUNCTION_ANNOTATIONS, + MAKE_FUNCTION_CLOSURE, +) + + +def _apply_function_attribute( + func: ast.Lambda | CompLambda, flag: int, value: ast.expr +) -> ast.Lambda | CompLambda: + """Attach one function attribute to a reconstructed lambda.""" + if flag == MAKE_FUNCTION_CLOSURE: + # The body has already resolved each free variable: to a captured value + # if the cell came from outside the comprehension, and otherwise to the + # name the reconstructed tree binds it under. See `handle_load_deref`. + return func + if flag == MAKE_FUNCTION_ANNOTATE: + # A lambda has no annotations, and the AST does not carry the lazy + # annotate function 3.14 attaches to annotated functions. + return func + + assert isinstance(func, ast.Lambda) and not isinstance(func, CompLambda), ( + "Only lambdas carry defaults; comprehensions take exactly one argument" + ) + + if flag == MAKE_FUNCTION_DEFAULTS: + # A tuple of defaults for the *trailing* positional parameters. + assert isinstance(value, ast.Tuple), "Expected a tuple of default values" + func.args.defaults = list(value.elts) + elif flag == MAKE_FUNCTION_KWDEFAULTS: + # A dict mapping keyword-only parameter names to their defaults. + assert isinstance(value, ast.Dict), "Expected a dict of keyword defaults" + by_name = { + key.value: val + for key, val in zip(value.keys, value.values) + if isinstance(key, ast.Constant) + } + func.args.kw_defaults = [by_name.get(a.arg) for a in func.args.kwonlyargs] + else: + raise NotImplementedError("Function annotations are not supported") + + return func + + +def _split_callable( + first: ast.expr, second: ast.expr +) -> tuple[ast.expr, ast.expr | None]: + """Separate the callable from the NULL-or-self slot beside it. + + Which of the two comes first varies: LOAD_GLOBAL and LOAD_ATTR report the + order in their argrepr, and it differs between 3.13 and 3.14. + """ + if isinstance(first, Null): + return second, None + elif isinstance(second, Null): + return first, None + else: + return first, second + + +def _build_variadic_call( + func: ast.expr, + self_arg: ast.expr | None, + positional: ast.expr, + keyword_mapping: ast.expr | None, +) -> ast.Call: + """Assemble `func(*positional, **keyword_mapping)`. + + CALL_FUNCTION_EX receives its arguments already collected into a sequence + and a mapping, with the original mix of plain and starred arguments no + longer distinguishable. Spelling every argument as unpacked reproduces the + call exactly, even where the source did not use `*` for all of them. + """ + args: list[ast.expr] = [] if self_arg is None else [ensure_ast(self_arg)] + args.append(ast.Starred(value=ensure_ast(positional), ctx=ast.Load())) + keywords = ( + [] + if keyword_mapping is None + else [ast.keyword(arg=None, value=ensure_ast(keyword_mapping))] + ) + return ast.Call(func=ensure_ast(func), args=args, keywords=keywords) + + +@register_handler("CALL_FUNCTION_EX", version=PythonVersion.PY_312) +@register_handler("CALL_FUNCTION_EX", version=PythonVersion.PY_313) +def handle_call_function_ex( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # Stack: callable and NULL-or-self, the positional sequence, and -- only + # when the low bit of the oparg is set -- the keyword mapping. + size = 4 if instr.arg else 3 + keyword_mapping = state.stack[-1] if instr.arg else None + positional = state.stack[-2] if instr.arg else state.stack[-1] + func, self_arg = _split_callable(state.stack[-size], state.stack[-size + 1]) + + call = _build_variadic_call(func, self_arg, positional, keyword_mapping) + return replace(state, stack=state.stack[:-size] + [call]) + + +@register_handler("CALL_FUNCTION_EX", version=PythonVersion.PY_314) +def handle_call_function_ex_314( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # 3.14 always reserves the keyword-mapping slot, pushing NULL into it when + # the call has no `**` argument, so the layout is a fixed four slots. + keyword_mapping = None if isinstance(state.stack[-1], Null) else state.stack[-1] + func, self_arg = _split_callable(state.stack[-4], state.stack[-3]) + + call = _build_variadic_call(func, self_arg, state.stack[-2], keyword_mapping) + return replace(state, stack=state.stack[:-4] + [call]) + + +@register_handler("MAKE_FUNCTION", version=PythonVersion.PY_312) +def handle_make_function_312( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # MAKE_FUNCTION in Python 3.12 uses flags to determine stack consumption. + # Unlike 3.10 there is no qualified name on the stack, and unlike 3.13 the + # extra attributes travel with this instruction rather than with a following + # SET_FUNCTION_ATTRIBUTE. They are pushed in ascending flag order, below the + # code object. + assert instr.arg is not None + assert isinstance(state.stack[-1], ast.Lambda | CompLambda), ( + "Expected a function object (Lambda or CompLambda) on the stack." + ) + + set_flags = [flag for flag in MAKE_FUNCTION_FLAGS if instr.arg & flag] + attributes = state.stack[-1 - len(set_flags) : -1] + + func = copy.deepcopy(state.stack[-1]) + for flag, value in zip(set_flags, attributes): + func = _apply_function_attribute(func, flag, value) + + new_stack = state.stack[: -1 - len(set_flags)] + [func] + return replace(state, stack=new_stack) + + +# Python 3.13 version +@register_handler("MAKE_FUNCTION", version=PythonVersion.PY_313) +@register_handler("MAKE_FUNCTION", version=PythonVersion.PY_314) +def handle_make_function( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # MAKE_FUNCTION in Python 3.13 is simplified: it only takes a code object from the stack + # and creates a function from it. No flags, no extra attributes on the stack. + # All extra attributes are handled by separate SET_FUNCTION_ATTRIBUTE instructions. + + # Pop the function object from the stack (it's the only thing expected) + # Conversion from CodeType to ast.Lambda should have happened already + assert isinstance(state.stack[-1], ast.Lambda | CompLambda), ( + "Expected a function object (Lambda or CompLambda) on the stack." + ) + return state + + +@register_handler("SET_FUNCTION_ATTRIBUTE", version=PythonVersion.PY_313) +@register_handler("SET_FUNCTION_ATTRIBUTE", version=PythonVersion.PY_314) +def handle_set_function_attribute( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # SET_FUNCTION_ATTRIBUTE sets one attribute on a function object. Python + # 3.13 uses it in place of the MAKE_FUNCTION flags; the stack holds the + # attribute value below the function, and only the function is left behind. + assert instr.arg is not None + assert isinstance(state.stack[-1], ast.Lambda | CompLambda), ( + "Expected a function object (Lambda or CompLambda) on the stack." + ) + + func = _apply_function_attribute( + copy.deepcopy(state.stack[-1]), instr.arg, state.stack[-2] + ) + return replace(state, stack=state.stack[:-2] + [func]) + + +# ============================================================================ +# OBJECT ACCESS HANDLERS +# ============================================================================ + + +@register_handler("LOAD_ATTR", version=PythonVersion.PY_312) +@register_handler("LOAD_ATTR", version=PythonVersion.PY_313) +@register_handler("LOAD_ATTR", version=PythonVersion.PY_314) +def handle_load_attr( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_ATTR loads an attribute from the object on top of stack + obj = ensure_ast(state.stack[-1]) + attr_name = instr.argval + + # Create attribute access AST + attr_node = ast.Attribute(value=obj, attr=attr_name, ctx=ast.Load()) + if instr.argrepr.endswith(" + NULL|self"): + new_stack = state.stack[:-1] + [attr_node, Null()] + elif instr.argrepr.startswith("NULL|self + "): + new_stack = state.stack[:-1] + [Null(), attr_node] + else: + new_stack = state.stack[:-1] + [attr_node] + return replace(state, stack=new_stack) + + +@register_handler("BINARY_SUBSCR", version=PythonVersion.PY_312) +@register_handler("BINARY_SUBSCR", version=PythonVersion.PY_313) +def handle_binary_subscr( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # BINARY_SUBSCR implements obj[index] - pops index and obj from stack + index = ensure_ast(state.stack[-1]) # Index is on top + obj = ensure_ast(state.stack[-2]) # Object is below index + new_stack = state.stack[:-2] + + # Create subscript access AST + subscr_node = ast.Subscript(value=obj, slice=index, ctx=ast.Load()) + new_stack = new_stack + [subscr_node] + return replace(state, stack=new_stack) + + +@register_handler("BINARY_SLICE", version=PythonVersion.PY_312) +@register_handler("BINARY_SLICE", version=PythonVersion.PY_313) +@register_handler("BINARY_SLICE", version=PythonVersion.PY_314) +def handle_binary_slice( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # BINARY_SLICE implements obj[start:end] - pops start, end, and obj from stack + end = ensure_ast(state.stack[-1]) + start = ensure_ast(state.stack[-2]) + container = ensure_ast(state.stack[-3]) # Object is below start and end + sliced = ast.Subscript( + value=container, + slice=ast.Slice(lower=start, upper=end, step=None), + ctx=ast.Load(), + ) + new_stack = state.stack[:-3] + [sliced] + return replace(state, stack=new_stack) + + +# ============================================================================ +# OTHER CONTAINER BUILDING HANDLERS +# ============================================================================ + + +@register_handler("UNPACK_SEQUENCE", version=PythonVersion.PY_312) +@register_handler("UNPACK_SEQUENCE", version=PythonVersion.PY_313) +@register_handler("UNPACK_SEQUENCE", version=PythonVersion.PY_314) +def handle_unpack_sequence( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # UNPACK_SEQUENCE splits a comprehension loop target into `arg` sub-targets, + # as in ((k, v) for k, v in items). The names are not known yet, so the + # single target hole is refined into a tuple of fresh holes, which the + # following STORE_* instructions bind one at a time. + # + # CPython pushes the unpacked values right-to-left, so element 0 ends up on + # top of the stack and is consumed by the first STORE_*. + assert instr.arg is not None + unpack_count: int = instr.arg + + if not isinstance(state.stack[-1], TargetHole): + # Destructuring a known value rather than a loop target, as 3.14 emits + # when it unrolls a single-iteration loop over a literal. + elements = _literal_elements(state.stack[-1], unpack_count) + return replace(state, stack=state.stack[:-1] + list(reversed(elements))) + + holes = [TargetHole() for _ in range(unpack_count)] + new_stack = _bind_target_hole( + state.stack, state.stack[-1], ast.Tuple(elts=list(holes), ctx=ast.Store()) + ) + return replace(state, stack=new_stack[:-1] + list(reversed(holes))) + + +@register_handler("UNPACK_EX", version=PythonVersion.PY_312) +@register_handler("UNPACK_EX", version=PythonVersion.PY_313) +@register_handler("UNPACK_EX", version=PythonVersion.PY_314) +def handle_unpack_ex( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # UNPACK_EX handles a starred target, as in ((a, b) for a, *b in pairs). + # The low byte of the argument counts the targets before the starred one and + # the high byte counts those after it; the starred target itself collects + # whatever is left over. As with UNPACK_SEQUENCE the values are pushed + # right-to-left, so the first target ends up on top of the stack. + assert instr.arg is not None + before, after = instr.arg & 0xFF, instr.arg >> 8 + + if not isinstance(state.stack[-1], TargetHole): + # Destructuring a known value; the starred target collects the middle. + elements = _literal_elements(state.stack[-1]) + assert len(elements) >= before + after, "Too few values to unpack" + middle = elements[before : len(elements) - after] + unpacked: list[ast.expr] = [ + *elements[:before], + ast.List(elts=list(middle), ctx=ast.Load()), + *elements[len(elements) - after :], + ] + return replace(state, stack=state.stack[:-1] + list(reversed(unpacked))) + + holes = [TargetHole() for _ in range(before + 1 + after)] + elts: list[ast.expr] = list(holes) + elts[before] = ast.Starred(value=holes[before], ctx=ast.Store()) + + new_stack = _bind_target_hole( + state.stack, state.stack[-1], ast.Tuple(elts=elts, ctx=ast.Store()) + ) + return replace(state, stack=new_stack[:-1] + list(reversed(holes))) + + +@register_handler("BUILD_TUPLE", version=PythonVersion.PY_312) +@register_handler("BUILD_TUPLE", version=PythonVersion.PY_313) +@register_handler("BUILD_TUPLE", version=PythonVersion.PY_314) +def handle_build_tuple( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert instr.arg is not None + tuple_size: int = instr.arg + # Pop elements for the tuple + elements = ( + [ensure_ast(elem) for elem in state.stack[-tuple_size:]] + if tuple_size > 0 + else [] + ) + new_stack = state.stack[:-tuple_size] if tuple_size > 0 else state.stack + + # Create tuple AST + tuple_node = ast.Tuple(elts=elements, ctx=ast.Load()) + new_stack = new_stack + [tuple_node] + return replace(state, stack=new_stack) + + +@register_handler("BUILD_SLICE", version=PythonVersion.PY_312) +@register_handler("BUILD_SLICE", version=PythonVersion.PY_313) +@register_handler("BUILD_SLICE", version=PythonVersion.PY_314) +def handle_build_slice( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # BUILD_SLICE creates a slice object from the top of the stack + # The number of elements to pop is determined by the instruction argument + assert instr.arg is not None + slice_size: int = instr.arg + + if slice_size == 2: + # Slice with start and end: [start, end] + end = ensure_ast(state.stack[-1]) + start = ensure_ast(state.stack[-2]) + new_stack = state.stack[:-2] + slice_node = ast.Slice(lower=start, upper=end, step=None) + elif slice_size == 3: + # Slice with start, end, and step: [start, end, step] + step = ensure_ast(state.stack[-1]) + end = ensure_ast(state.stack[-2]) + start = ensure_ast(state.stack[-3]) + new_stack = state.stack[:-3] + slice_node = ast.Slice(lower=start, upper=end, step=step) + else: + raise ValueError(f"Unsupported slice size: {slice_size}") + + # Create slice AST + new_stack = new_stack + [slice_node] + return replace(state, stack=new_stack) + + +@register_handler("BUILD_CONST_KEY_MAP", version=PythonVersion.PY_312) +@register_handler("BUILD_CONST_KEY_MAP", version=PythonVersion.PY_313) +def handle_build_const_key_map( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # BUILD_CONST_KEY_MAP builds a dictionary with constant keys + # The keys are in a tuple on TOS, values are on the stack below + assert instr.arg is not None + assert isinstance(state.stack[-1], ast.Tuple), "Expected a tuple of keys" + map_size: int = instr.arg + # Pop the keys tuple and values + keys_tuple: ast.Tuple = state.stack[-1] + keys: list[ast.expr | None] = [ensure_ast(key) for key in keys_tuple.elts] + values = [ensure_ast(val) for val in state.stack[-map_size - 1 : -1]] + new_stack = state.stack[: -map_size - 1] + + # Create dictionary AST + dict_node = ast.Dict(keys=keys, values=values) + new_stack = new_stack + [dict_node] + return replace(state, stack=new_stack) + + +@register_handler("LIST_EXTEND", version=PythonVersion.PY_312) +@register_handler("LIST_EXTEND", version=PythonVersion.PY_313) +@register_handler("LIST_EXTEND", version=PythonVersion.PY_314) +def handle_list_extend( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LIST_EXTEND appends the contents of the iterable at TOS to the list + # further down the stack. That list is either the empty ListComp that + # BUILD_LIST(0) optimistically created -- a list display, not a + # comprehension after all -- or a partly built argument list for a call + # with a starred argument. + update = state.stack[-1] + target = state.stack[-instr.argval - 1] + + # A literal iterable contributes its elements directly; anything else has to + # stay unpacked, as in `[*whatever]`. + elements: list[ast.expr] + if isinstance(update, ast.Tuple | ast.List): + elements = [ensure_ast(e) for e in update.elts] + else: + elements = [ast.Starred(value=ensure_ast(update), ctx=ast.Load())] + + if isinstance(target, ast.ListComp) and not target.generators: + merged = ast.List(elts=elements, ctx=ast.Load()) + else: + assert isinstance(target, ast.List), "LIST_EXTEND expects a list to extend" + merged = ast.List(elts=list(target.elts) + elements, ctx=ast.Load()) + + new_stack = state.stack[:-1] + new_stack[-instr.argval] = merged + return replace(state, stack=new_stack) + + +@register_handler("DICT_MERGE", version=PythonVersion.PY_312) +@register_handler("DICT_MERGE", version=PythonVersion.PY_313) +@register_handler("DICT_MERGE", version=PythonVersion.PY_314) +def handle_dict_merge( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # DICT_MERGE folds the mapping at TOS into the one below it, rejecting + # duplicate keys. It assembles the keyword arguments of a call using `**`. + update = state.stack[-1] + target = state.stack[-instr.argval - 1] + + # An `ast.Dict` entry with a key of None is `**value`, which is how a + # mapping that is not a literal has to be spliced in. + def entries(node: ast.expr) -> tuple[list[ast.expr | None], list[ast.expr]]: + if isinstance(node, ast.Dict): + return list(node.keys), list(node.values) + return [None], [ensure_ast(node)] + + if isinstance(target, ast.DictComp) and not target.generators: + # BUILD_MAP(0) guessed at a dict comprehension; it was a `**` argument. + keys, values = entries(update) + else: + assert isinstance(target, ast.Dict), "DICT_MERGE expects a dict to merge into" + target_keys, target_values = entries(target) + update_keys, update_values = entries(update) + keys, values = target_keys + update_keys, target_values + update_values + + new_stack = state.stack[:-1] + new_stack[-instr.argval] = ast.Dict(keys=keys, values=values) + return replace(state, stack=new_stack) + + +@register_handler("SET_UPDATE", version=PythonVersion.PY_312) +@register_handler("SET_UPDATE", version=PythonVersion.PY_313) +@register_handler("SET_UPDATE", version=PythonVersion.PY_314) +def handle_set_update( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # The set being extended is actually in state.result instead of the stack + # because it was initially recognized as a list comprehension in BUILD_SET, + # while the actual result expression is in the stack where the set "should be" + # and needs to be put back into the state result slot + assert isinstance(state.stack[-instr.argval - 1], ast.SetComp) + assert isinstance(state.stack[-1], ast.Tuple | ast.List | ast.Set) + + new_val = ast.Set(elts=[ensure_ast(e) for e in state.stack[-1].elts]) + new_stack = state.stack[:-2] + [new_val] + + return replace(state, stack=new_stack) + + +@register_handler("DICT_UPDATE", version=PythonVersion.PY_312) +@register_handler("DICT_UPDATE", version=PythonVersion.PY_313) +@register_handler("DICT_UPDATE", version=PythonVersion.PY_314) +def handle_dict_update( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # The dict being extended is actually in state.result instead of the stack + # because it was initially recognized as a list comprehension in BUILD_MAP, + # while the actual result expression is in the stack where the dict "should be" + # and needs to be put back into the state result slot + assert isinstance(state.stack[-instr.argval - 1], ast.DictComp) + assert isinstance(state.stack[-1], ast.Dict) + + new_val = ast.Dict( + keys=[ensure_ast(e) for e in state.stack[-1].keys], + values=[ensure_ast(e) for e in state.stack[-1].values], + ) + new_stack = state.stack[:-2] + [new_val] + + return replace(state, stack=new_stack) + + +@register_handler("BUILD_STRING", version=PythonVersion.PY_312) +@register_handler("BUILD_STRING", version=PythonVersion.PY_313) +@register_handler("BUILD_STRING", version=PythonVersion.PY_314) +def handle_build_string( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # BUILD_STRING concatenates strings from the stack + # For f-strings, it combines FormattedValue and Constant nodes + assert instr.arg is not None + string_size: int = instr.arg + + if string_size == 0: + # Empty string case + new_stack = state.stack + [ast.Constant(value="")] + return replace(state, stack=new_stack) + + # Pop elements for the string + elements = [ensure_ast(elem) for elem in state.stack[-string_size:]] + new_stack = state.stack[:-string_size] + + # Check if this is an f-string build (has FormattedValue nodes) + # or a regular string concatenation + if any(isinstance(elem, ast.JoinedStr) for elem in elements): + # This is an f-string - create JoinedStr + values = [] + for elem in elements: + if isinstance(elem, ast.JoinedStr): + values.extend(elem.values) + else: + values.append(elem) + return replace(state, stack=new_stack + [ast.JoinedStr(values=values)]) + elif all(isinstance(elem, ast.Constant) for elem in elements): + # This is regular string concatenation or format spec building + # If all elements are constants, we might be building a format spec + # Concatenate the constant strings + assert all( + isinstance(elem, ast.Constant) and isinstance(elem.value, str) + for elem in elements + ) + concat_str = "".join( + elem.value + for elem in elements + if isinstance(elem, ast.Constant) and isinstance(elem.value, str) + ) + return replace(state, stack=new_stack + [ast.Constant(value=concat_str)]) + else: + raise TypeError("Should not be here?") + + +@register_handler("FORMAT_VALUE", version=PythonVersion.PY_312) +def handle_format_value( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # FORMAT_VALUE formats a string with a value in Python 3.12 + # Flag bits: (flags & 0x03) = conversion, (flags & 0x04) = has format spec + assert instr.arg is not None, "FORMAT_VALUE requires flags argument" + assert len(state.stack) >= 1, "Not enough items on stack for FORMAT_VALUE" + + flags = instr.arg + + # Check if there's a format specification + has_format_spec = bool(flags & 0x04) + + if has_format_spec: + # Pop format spec and value + assert len(state.stack) >= 2, ( + "FORMAT_VALUE with format spec needs 2 stack items" + ) + format_spec = ensure_ast(state.stack[-1]) + value = ensure_ast(state.stack[-2]) + new_stack = state.stack[:-2] + + # Wrap format spec in JoinedStr if it's a constant + if isinstance(format_spec, ast.Constant): + format_spec_node = ast.JoinedStr(values=[format_spec]) + else: + assert isinstance(format_spec, ast.JoinedStr) + format_spec_node = format_spec + else: + # Just pop the value + value = ensure_ast(state.stack[-1]) + new_stack = state.stack[:-1] + format_spec_node = None + + # Determine conversion type from flags + conversion_flags = flags & 0x03 + conversion_map = { + 0: -1, # No conversion + 1: 115, # str (!s) + 2: 114, # repr (!r) + 3: 97, # ascii (!a) + } + conversion = conversion_map[conversion_flags] + + # Create formatted value AST + formatted_node = ast.FormattedValue( + value=value, conversion=conversion, format_spec=format_spec_node + ) + new_stack = new_stack + [ast.JoinedStr(values=[formatted_node])] + return replace(state, stack=new_stack) + + +@register_handler("FORMAT_SIMPLE", version=PythonVersion.PY_313) +@register_handler("FORMAT_SIMPLE", version=PythonVersion.PY_314) +def handle_format_simple( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # FORMAT_SIMPLE formats a string with a single value + # Pops the value and the format string from the stack + assert len(state.stack) >= 1, "Not enough items on stack for FORMAT_SIMPLE" + value = state.stack[-1] + + # Check if the value was converted + if isinstance(value, ConvertedValue): + conversion = value.ast_conversion + value = value.value + else: + conversion = -1 + value = ensure_ast(value) + + # Create formatted string AST + formatted_node = ast.FormattedValue( + value=value, conversion=conversion, format_spec=None + ) + new_stack = state.stack[:-1] + [ast.JoinedStr(values=[formatted_node])] + return replace(state, stack=new_stack) + + +@register_handler("FORMAT_WITH_SPEC", version=PythonVersion.PY_313) +@register_handler("FORMAT_WITH_SPEC", version=PythonVersion.PY_314) +def handle_format_with_spec( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # FORMAT_WITH_SPEC formats a value with a format specifier + # Stack order in Python 3.13: format_spec on top, value below + assert len(state.stack) >= 2, "Not enough items on stack for FORMAT_WITH_SPEC" + format_spec = ensure_ast(state.stack[-1]) # Format spec is on top + value = state.stack[-2] # Value is below + + # Check if the value was converted + if isinstance(value, ConvertedValue): + conversion = value.ast_conversion + value = value.value + else: + conversion = -1 + value = ensure_ast(value) + + # Create formatted string AST with specifier + # The format_spec should be wrapped in a JoinedStr if it's a simple constant + if isinstance(format_spec, ast.Constant): + format_spec_node = ast.JoinedStr(values=[format_spec]) + else: + # Already a JoinedStr from nested formatting + assert isinstance(format_spec, ast.JoinedStr) + format_spec_node = format_spec + + formatted_node = ast.FormattedValue( + value=value, conversion=conversion, format_spec=format_spec_node + ) + new_stack = state.stack[:-2] + [ast.JoinedStr(values=[formatted_node])] + return replace(state, stack=new_stack) + + +# ============================================================================ +# CONDITIONAL JUMP HANDLERS +# ============================================================================ + + +def _handle_pop_jump_if( + f_condition: Callable[[ast.expr], ast.expr], + state: ReconstructionState, + instr: dis.Instruction, +) -> ReconstructionState: + # Generic handler for POP_JUMP_IF_* instructions. Pops a value from the + # stack; `condition` is true exactly when the jump is taken. + condition: ast.expr = f_condition(ensure_ast(state.stack[-1])) + + # An inlined-builtin guard is an implementation detail of the interpreter, + # not part of the comprehension: drop it rather than record it as a filter. + if _specialization_guard_edge(state, instr) is not None: + return replace(state, stack=state.stack[:-1]) + + kind, _ = _classify_branch(state, instr) + edge = state.branches.get(instr.offset, BranchEdge.TAKE_JUMP) + + if kind is BranchKind.TERNARY: + return _handle_conditional_expression(state, instr, condition, edge) + + # A filter. The guard is the condition under which *this* path carries on + # toward the element, so it is negated when the path falls through. + guard = condition if edge is BranchEdge.TAKE_JUMP else _negate(condition) + return _attach_filter(state, guard) + + +def _attach_filter( + state: ReconstructionState, guard: ast.expr | None +) -> ReconstructionState: + """Conjoin ``guard`` to the filters of the innermost unfinished comprehension.""" + for pos, item in zip(reversed(range(len(state.stack))), reversed(state.stack)): + if not isinstance(item, CompExp): + continue + + elt: ast.expr = item.value if isinstance(item, ast.DictComp) else item.elt + new_result: CompExp = copy.deepcopy(item) + + if isinstance(elt, Placeholder): + resolved = guard + elif isinstance(elt, ast.IfExp) and any( + isinstance(x, Placeholder) for x in ast.walk(elt) + ): + # A conditional expression was being built up in the element slot, + # but it turned out to be part of this filter's condition. Move it + # back out, plugging the guard into the arm still awaiting a value. + if isinstance(new_result, ast.DictComp): + new_result.key, new_result.value = Placeholder(), Placeholder() + else: + new_result.elt = Placeholder() + + if guard is None: + resolved = None + else: + plugged = ReplacePlaceholder(guard).visit(copy.deepcopy(elt)) + assert isinstance(plugged, ast.expr) + resolved = plugged + else: + continue + + if resolved is not None: + ifs = new_result.generators[-1].ifs + combined = _conjoin(ifs + [resolved]) + assert combined is not None + new_result.generators[-1].ifs = [combined] + + new_stack = state.stack[:pos] + [new_result] + state.stack[pos + 1 : -1] + return replace(state, stack=new_stack) + + raise TypeError("No comprehension context found for filter condition") + + +def _handle_conditional_expression( + state: ReconstructionState, + instr: dis.Instruction, + condition: ast.expr, + edge: BranchEdge, +) -> ReconstructionState: + """Start an ``ast.IfExp``, marking the arm this path did not take.""" + for pos, item in zip(reversed(range(len(state.stack))), reversed(state.stack)): + if any(isinstance(x, Placeholder) for x in ast.walk(item)): + body: Skipped | Placeholder + orelse: Skipped | Placeholder + skipped = Skipped(id=f".SKIPPED_{instr.offset}") + if edge is BranchEdge.FALL_THROUGH: + body, orelse = skipped, Placeholder() + else: + body, orelse = Placeholder(), skipped + + new_ifexp = ast.IfExp(test=condition, body=body, orelse=orelse) + new_result = ReplacePlaceholder(new_ifexp).visit(copy.deepcopy(item)) + new_stack = state.stack[:pos] + [new_result] + state.stack[pos + 1 : -1] + return replace(state, stack=new_stack) + + raise TypeError("No placeholder found for conditional expression") + + +@register_handler("POP_JUMP_IF_TRUE", version=PythonVersion.PY_312) +@register_handler("POP_JUMP_IF_TRUE", version=PythonVersion.PY_313) +@register_handler("POP_JUMP_IF_TRUE", version=PythonVersion.PY_314) +def handle_pop_jump_if_true( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # POP_JUMP_IF_TRUE pops a value from the stack and jumps if it's true + # In Python 3.13, this is used for filter conditions where True means continue + return _handle_pop_jump_if(lambda c: c, state, instr) + + +@register_handler("POP_JUMP_IF_FALSE", version=PythonVersion.PY_312) +@register_handler("POP_JUMP_IF_FALSE", version=PythonVersion.PY_313) +@register_handler("POP_JUMP_IF_FALSE", version=PythonVersion.PY_314) +def handle_pop_jump_if_false( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # POP_JUMP_IF_FALSE pops a value from the stack and jumps if it's false + # In comprehensions, this is used for filter conditions + return _handle_pop_jump_if( + lambda c: ast.UnaryOp(op=ast.Not(), operand=c), state, instr + ) + + +@register_handler("POP_JUMP_IF_NONE", version=PythonVersion.PY_312) +@register_handler("POP_JUMP_IF_NONE", version=PythonVersion.PY_313) +@register_handler("POP_JUMP_IF_NONE", version=PythonVersion.PY_314) +def handle_pop_jump_if_none( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # POP_JUMP_IF_NONE pops a value and jumps if it's None + return _handle_pop_jump_if( + lambda c: ast.Compare( + left=c, ops=[ast.Is()], comparators=[ast.Constant(value=None)] + ), + state, + instr, + ) + + +@register_handler("POP_JUMP_IF_NOT_NONE", version=PythonVersion.PY_312) +@register_handler("POP_JUMP_IF_NOT_NONE", version=PythonVersion.PY_313) +@register_handler("POP_JUMP_IF_NOT_NONE", version=PythonVersion.PY_314) +def handle_pop_jump_if_not_none( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # POP_JUMP_IF_NOT_NONE pops a value and jumps if it's not None + return _handle_pop_jump_if( + lambda c: ast.Compare( + left=c, ops=[ast.IsNot()], comparators=[ast.Constant(value=None)] + ), + state, + instr, + ) + + +@register_handler("JUMP_FORWARD", version=PythonVersion.PY_312) +@register_handler("JUMP_FORWARD", version=PythonVersion.PY_313) +@register_handler("JUMP_FORWARD", version=PythonVersion.PY_314) +def handle_jump_forward( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # JUMP_FORWARD is used to jump forward in the code + # In generator expressions, this is often used to skip code in conditional logic + return state + + +@register_handler("JUMP_BACKWARD", version=PythonVersion.PY_312) +@register_handler("JUMP_BACKWARD", version=PythonVersion.PY_313) +@register_handler("JUMP_BACKWARD", version=PythonVersion.PY_314) +def handle_jump_backward( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # JUMP_BACKWARD is used to jump back to the beginning of a loop (replaces JUMP_ABSOLUTE in 3.13) + # In generator expressions, this typically indicates the end of the loop body + return state + + +@register_handler("JUMP_BACKWARD_NO_INTERRUPT", version=PythonVersion.PY_312) +@register_handler("JUMP_BACKWARD_NO_INTERRUPT", version=PythonVersion.PY_313) +@register_handler("JUMP_BACKWARD_NO_INTERRUPT", version=PythonVersion.PY_314) +def handle_jump_backward_no_interrupt( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + raise TypeError( + "JUMP_BACKWARD_NO_INTERRUPT instruction should not appear in generator comprehensions" + ) + + +@register_handler("JUMP_NO_INTERRUPT", version=PythonVersion.PY_312) +@register_handler("JUMP_NO_INTERRUPT", version=PythonVersion.PY_313) +@register_handler("JUMP_NO_INTERRUPT", version=PythonVersion.PY_314) +def handle_jump_no_interrupt( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + raise TypeError( + "JUMP_NO_INTERRUPT instruction should not appear in generator comprehensions" + ) + + +@register_handler("JUMP", version=PythonVersion.PY_312) +@register_handler("JUMP", version=PythonVersion.PY_313) +@register_handler("JUMP", version=PythonVersion.PY_314) +def handle_jump( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + raise TypeError("JUMP instruction should not appear in generator comprehensions") + + +@register_handler("EXTENDED_ARG", version=PythonVersion.PY_312) +@register_handler("EXTENDED_ARG", version=PythonVersion.PY_313) +@register_handler("EXTENDED_ARG", version=PythonVersion.PY_314) +def handle_extended_arg( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # EXTENDED_ARG prefixes an instruction whose argument does not fit in a + # byte. `dis` has already folded it into the following instruction's `arg`, + # so there is nothing left to do here. + return state + + +@register_handler("RESUME", version=PythonVersion.PY_312) +@register_handler("RESUME", version=PythonVersion.PY_313) +@register_handler("RESUME", version=PythonVersion.PY_314) +def handle_resume( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # RESUME is used for resuming execution after yield/await - mostly no-op for AST reconstruction + return state + + +@register_handler("SEND", version=PythonVersion.PY_312) +@register_handler("SEND", version=PythonVersion.PY_313) +@register_handler("SEND", version=PythonVersion.PY_314) +def handle_send( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + raise TypeError("SEND instruction should not appear in generator comprehensions") + + +# ============================================================================ +# UTILITY FUNCTIONS +# ============================================================================ + + +@functools.singledispatch +def ensure_ast(value) -> ast.expr: + """Ensure value is an AST node""" + raise TypeError(f"Cannot convert {type(value)} to AST node") + + +@ensure_ast.register +def _ensure_ast_ast(value: ast.expr) -> ast.expr: + """If already an AST node, return it as is""" + return value + + +@ensure_ast.register(int) +@ensure_ast.register(float) +@ensure_ast.register(str) +@ensure_ast.register(bytes) +@ensure_ast.register(bool) +@ensure_ast.register(complex) +@ensure_ast.register(type(None)) +def _ensure_ast_constant(value) -> ast.Constant: + return ast.Constant(value=value) + + +@ensure_ast.register +def _ensure_ast_tuple(value: tuple) -> ast.Tuple: + """Convert tuple to AST""" + return ast.Tuple(elts=[ensure_ast(v) for v in value], ctx=ast.Load()) + + +def _unconsumed(value: Iterator) -> typing.Any: + """Return the items an iterator has not yet yielded, as a concrete sequence. + + Built-in sequence iterators pickle as ``(iter, (underlying,), index)``, where + ``index`` is how far the iterator has advanced (absent or ``None`` when it + does not apply). ``reversed`` objects pickle with ``reversed`` as the + callable and count *down* from the end of the underlying sequence. + """ + reduced = value.__reduce__() + assert isinstance(reduced, tuple) and len(reduced) >= 2, ( + f"Cannot recover the contents of {type(value)}" + ) + if not reduced[1]: # an exhausted iterator pickles with no arguments + return () + + underlying = reduced[1][0] + index = reduced[2] if len(reduced) > 2 and reduced[2] is not None else 0 + return underlying[index::-1] if reduced[0] is reversed else underlying[index:] + + +@ensure_ast.register(type(iter((1,)))) +def _ensure_ast_tuple_iterator(value: Iterator) -> ast.Tuple: + return ensure_ast(tuple(_unconsumed(value))) # type: ignore + + +@ensure_ast.register +def _ensure_ast_list(value: list) -> ast.List: + return ast.List(elts=[ensure_ast(v) for v in value], ctx=ast.Load()) + + +@ensure_ast.register(type(iter([1]))) +@ensure_ast.register(type(iter({1: 2}.values()))) +@ensure_ast.register(type(iter({1: 2}.items()))) +@ensure_ast.register(type(iter(reversed([1])))) +@ensure_ast.register(reversed) +def _ensure_ast_list_iterator(value: Iterator) -> ast.List: + return ensure_ast(list(_unconsumed(value))) # type: ignore + + +@ensure_ast.register(type(iter("ab"))) # str_ascii_iterator +@ensure_ast.register(type(iter("\xe9b"))) # str_iterator +@ensure_ast.register(type(iter(b"ab"))) +@ensure_ast.register(type(iter(bytearray(b"ab")))) +def _ensure_ast_str_iterator(value: Iterator) -> ast.Constant: + remainder = _unconsumed(value) + # bytearray iteration yields ints, exactly as bytes iteration does + return ensure_ast( # type: ignore + bytes(remainder) if isinstance(remainder, bytearray) else remainder + ) + + +@ensure_ast.register(set) +@ensure_ast.register(frozenset) +def _ensure_ast_set(value: set | frozenset) -> ast.Set: + return ast.Set(elts=[ensure_ast(v) for v in value]) + + +@ensure_ast.register(type(iter({1}))) +def _ensure_ast_set_iterator(value: Iterator) -> ast.Set: + return ensure_ast(set(_unconsumed(value))) # type: ignore + + +@ensure_ast.register +def _ensure_ast_dict(value: dict) -> ast.Dict: + return ast.Dict( + keys=[ensure_ast(k) for k in value.keys()], + values=[ensure_ast(v) for v in value.values()], + ) + + +@ensure_ast.register(type(iter({1: 2}))) +def _ensure_ast_dict_iterator(value: Iterator) -> ast.expr: + return ensure_ast(_unconsumed(value)) + + +@ensure_ast.register(types.BuiltinFunctionType) +@ensure_ast.register(type) +def _ensure_ast_builtin(value: typing.Callable) -> ast.Name: + """A built-in callable is referred to by name, which resolves via builtins. + + Covers both built-in functions (``abs``) and built-in types used as + callables (``bool``), which appear as the predicate of a ``filter`` or the + function of a ``map``. + """ + name = getattr(value, "__name__", None) + assert name and getattr(builtins, name, None) is value, ( + f"Cannot reference non-builtin callable {value!r}" + ) + return ast.Name(id=name, ctx=ast.Load()) + + +@ensure_ast.register(zip) +@ensure_ast.register(enumerate) +@ensure_ast.register(map) +@ensure_ast.register(filter) +def _ensure_ast_iterator_adaptor(value: Iterator) -> ast.Call: + """Rebuild zip/enumerate/map/filter from the arguments they pickle with. + + These wrap other iterators rather than a concrete sequence, so unlike a list + or range iterator they cannot be materialised -- but ``__reduce__`` hands + back their constituent parts, each of which ``ensure_ast`` can handle in + turn. Any already-consumed prefix is reflected in the inner iterators. + + A ``zip`` also pickles its strictness as reduction state, which has to be + carried over as a keyword argument: a strict ``zip`` raises on ragged input + where a lax one stops at the shortest iterable. + """ + reduced = value.__reduce__() + if isinstance(reduced, str): + raise TypeError(f"Cannot convert {type(value)} to AST node") + func, args = reduced[:2] + keywords = [] + if isinstance(value, zip) and len(reduced) > 2 and reduced[2]: + keywords.append(ast.keyword(arg="strict", value=ast.Constant(value=True))) + return ast.Call( + func=ast.Name(id=func.__name__, ctx=ast.Load()), + args=[ensure_ast(arg) for arg in args], + keywords=keywords, + ) + + +@ensure_ast.register +def _ensure_ast_slice(value: slice) -> ast.Slice: + """A constant slice, as 3.14 emits for `s[1:3]` alongside BINARY_OP/NB_SUBSCR.""" + return ast.Slice( + lower=None if value.start is None else ensure_ast(value.start), + upper=None if value.stop is None else ensure_ast(value.stop), + step=None if value.step is None else ensure_ast(value.step), + ) + + +@ensure_ast.register +def _ensure_ast_range(value: range) -> ast.Call: + return ast.Call( + func=ast.Name(id="range", ctx=ast.Load()), + args=[ensure_ast(value.start), ensure_ast(value.stop), ensure_ast(value.step)], + keywords=[], + ) + + +@ensure_ast.register(type(iter(range(1)))) +def _ensure_ast_range_iterator(value: Iterator) -> ast.Call: + return ensure_ast(_unconsumed(value)) # type: ignore + + +def _cell_values( + code: types.CodeType, closure: tuple[types.CellType, ...] | None +) -> dict[str, typing.Any]: + """Read the cells a function closed over, by free-variable name. + + A cell that is still empty -- a recursive definition not yet bound, say -- + is left out. + """ + values: dict[str, typing.Any] = {} + for name, cell in zip(code.co_freevars, closure or ()): + try: + values[name] = cell.cell_contents + except ValueError: + continue + return values + + +def _freevar_bindings( + code: types.CodeType, captured: collections.abc.Mapping[str, typing.Any] +) -> dict[str, ast.expr]: + """AST nodes for the values ``code``'s free variables were captured from. + + A free variable with no captured value to be found is left out, and the + reconstructed tree keeps the bare name, which is as close as it can get. + + An iterator is refused rather than written in. `ensure_ast` spells one as + the elements it has left, which is what the outermost iterable wants -- it + is about to be consumed anyway -- but a captured iterator is a value the + body can do anything with, and a list of its remaining elements is not the + same object. + """ + bindings: dict[str, ast.expr] = {} + for name in code.co_freevars: + if name not in captured: + continue + value = captured[name] + try: + if isinstance(value, Iterator): + raise TypeError("an iterator has no AST spelling in value position") + bindings[name] = ensure_ast(value) + except (TypeError, AssertionError) as exc: + raise TypeError( + f"Cannot represent {value!r}, " + f"the value captured in free variable {name!r}: {exc}" + ) from exc + return bindings + + +def _reconstruct_code( + value: types.CodeType, freevars: dict[str, ast.expr] +) -> ast.Lambda | CompLambda: + """Reconstruct a lambda or comprehension body from its code object. + + ``freevars`` maps the names this code captures from enclosing scopes to the + values found in their cells; see `handle_load_deref`. + """ + assert inspect.iscode(value), "Input must be a code object" + + name: str = value.co_name.split(".")[-1] + + # Check preconditions + if name in {"", "", "", ""}: + assert name == "" or sys.version_info < (3, 13) + assert name != "" or value.co_flags & inspect.CO_GENERATOR + assert value.co_flags & inspect.CO_NEWLOCALS + assert value.co_argcount == 1 + assert value.co_kwonlyargcount == value.co_posonlyargcount == 0 + assert DummyIterName().id in value.co_varnames + elif name == "": + assert not value.co_flags & inspect.CO_GENERATOR + assert value.co_flags & inspect.CO_NEWLOCALS + assert DummyIterName().id not in value.co_varnames + else: + raise TypeError(f"Unsupported code object type: {value.co_name}") + + # Symbolic execution to reconstruct the AST + result: ast.expr = _symbolic_exec(value, freevars) + + # Check postconditions + assert not any(isinstance(x, ast.stmt) for x in ast.walk(result)), ( + "Final return value must not contain statement nodes" + ) + assert not any( + isinstance( + x, + Placeholder + | Skipped + | TargetHole + | CommonConstant + | Null + | CompLambda + | ConvertedValue, + ) + for x in ast.walk(result) + ), "Final return value must not contain temporary nodes" + assert not any(x.arg == ".0" for x in ast.walk(result) if isinstance(x, ast.arg)), ( + "Final return value must not contain .0 argument" + ) + assert not any( + isinstance(x, ast.Name) and x.id == ".0" + for x in ast.walk(result) + if not isinstance(x, DummyIterName) + ), "Final return value must not contain .0 names" + assert sum(1 for x in ast.walk(result) if isinstance(x, DummyIterName)) <= 1, ( + "Final return value must contain at most 1 dummy iterator names" + ) + assert all(x.generators for x in ast.walk(result) if isinstance(x, CompExp)), ( + "Return value must have generators if not a lambda" + ) + + if name == "" and isinstance(result, ast.expr): + # co_varnames lists parameters first: positional, keyword-only, then + # *args and **kwargs if present. Default values are not part of the code + # object -- they are pushed by the caller and attached by MAKE_FUNCTION + # (3.12) or SET_FUNCTION_ATTRIBUTE (3.13). + names = value.co_varnames + n_args, n_kwonly = value.co_argcount, value.co_kwonlyargcount + n_params = n_args + n_kwonly + + vararg = kwarg = None + if value.co_flags & inspect.CO_VARARGS: + vararg = ast.arg(arg=names[n_params]) + n_params += 1 + if value.co_flags & inspect.CO_VARKEYWORDS: + kwarg = ast.arg(arg=names[n_params]) + + args = ast.arguments( + posonlyargs=[ast.arg(arg=arg) for arg in names[: value.co_posonlyargcount]], + args=[ast.arg(arg=arg) for arg in names[value.co_posonlyargcount : n_args]], + vararg=vararg, + kwonlyargs=[ast.arg(arg=arg) for arg in names[n_args : n_args + n_kwonly]], + kw_defaults=[None] * n_kwonly, + kwarg=kwarg, + defaults=[], + ) + return ast.Lambda(args=args, body=result) + elif name == "" and isinstance(result, ast.GeneratorExp): + return CompLambda(body=result) + elif name == "" and isinstance(result, ast.DictComp): + return CompLambda(body=result) + elif name == "" and isinstance(result, ast.ListComp): + return CompLambda(body=result) + elif name == "" and isinstance(result, ast.SetComp): + return CompLambda(body=result) + else: + raise TypeError(f"Invalid result for type {name}: {result}") + + +@ensure_ast.register +def _ensure_ast_codeobj(value: types.CodeType) -> ast.Lambda | CompLambda: + """A bare code object has no cells attached, so free variables stay names.""" + return _reconstruct_code(value, {}) + + +@ensure_ast.register +def _ensure_ast_lambda(value: types.LambdaType) -> ast.Lambda: + assert inspect.isfunction(value) and value.__name__.endswith(""), ( + "Input must be a lambda function" + ) + code: types.CodeType = value.__code__ + result = _reconstruct_code( + code, _freevar_bindings(code, _cell_values(code, value.__closure__)) + ) + assert isinstance(result, ast.Lambda), "Lambda body must be an AST Lambda node" + assert not isinstance(result, CompLambda), "Lambda must not be a CompLambda" + + # Default values are not in the code object: they were evaluated where the + # lambda was written and attached to the function. A lambda built inside the + # comprehension gets them from the stack instead -- see + # `_apply_function_attribute` -- but one arriving as a live object carries + # them here, and dropping them would leave parameters with no way to be + # filled. They cover the *trailing* positional parameters. + if value.__defaults__: + result.args.defaults = [ensure_ast(d) for d in value.__defaults__] + if value.__kwdefaults__: + by_name = value.__kwdefaults__ + result.args.kw_defaults = [ + ensure_ast(by_name[arg.arg]) if arg.arg in by_name else None + for arg in result.args.kwonlyargs + ] + return result + + +@ensure_ast.register +def _ensure_ast_genexpr(genexpr: types.GeneratorType) -> ast.GeneratorExp: + assert inspect.isgenerator(genexpr), "Input must be a generator expression" + assert inspect.getgeneratorstate(genexpr) == inspect.GEN_CREATED, ( + "Generator must be in created state" + ) + assert genexpr.gi_frame is not None, "Generator must not be exhausted" + # A generator holds its cells in its frame rather than in a __closure__, and + # an unstarted frame already has them all copied in. + frame_locals = genexpr.gi_frame.f_locals + genexpr_ast = _reconstruct_code( + genexpr.gi_code, _freevar_bindings(genexpr.gi_code, frame_locals) + ) + assert isinstance(genexpr_ast, CompLambda) + geniter_ast = ensure_ast(frame_locals[".0"]) + result = genexpr_ast.inline(geniter_ast) + assert isinstance(result, ast.GeneratorExp) + assert inspect.getgeneratorstate(genexpr) == inspect.GEN_CREATED, ( + "Generator must stay in created state" + ) + return result + + +# ============================================================================ +# MAIN RECONSTRUCTION FUNCTION +# ============================================================================ + + +def disassemble( + genexpr: Generator[typing.Any, typing.Any, typing.Any], +) -> ast.Expression: + """ + Reconstruct an AST from a generator expression's bytecode. + + This function analyzes the bytecode of a generator object and reconstructs + an abstract syntax tree (AST) that represents the original comprehension + expression. The reconstruction process simulates the Python VM's execution + of the bytecode, building AST nodes instead of executing operations. + + The reconstruction handles complex comprehension features including: + - Multiple nested loops + - Filter conditions (if clauses) + - Complex expressions in the yield/result part + - Tuple unpacking in loop variables + - Various operators and function calls + + Args: + genexpr (Generator[object, None, None]): The generator object to analyze. + Must be a freshly created generator that has not been iterated yet + (in 'GEN_CREATED' state). + + Returns: + ast.Expression: An AST node representing the reconstructed comprehension. + + Raises: + ValueError: If the input is not a generator or if the generator + has already been started (not in 'GEN_CREATED' state). + TypeError: If some part of the comprehension has no AST spelling -- + an outermost iterable or a captured free variable whose value + cannot be written into the tree. + + Example: + >>> # Generator expression + >>> g = (x * 2 for x in range(10) if x % 2 == 0) + >>> ast_node = disassemble(g) + >>> isinstance(ast_node, ast.Expression) + True + + >>> # The reconstructed AST can be compiled and evaluated + >>> import ast + >>> code = compile(ast_node, '', 'eval') + >>> result = eval(code) + >>> list(result) + [0, 4, 8, 12, 16] + + Note: + The reconstruction is based on bytecode analysis and may not perfectly + preserve the original source code formatting or variable names in all + cases. However, the semantic behavior of the reconstructed AST should + match the original comprehension, subject to the limits set out under + "What is recovered, and what is not" in this module's docstring: + evaluating the result re-runs every expression in it, so a stateful + filter or element expression can answer differently the second time, + and the outermost iterable is a snapshot rather than an expression. + """ + if not inspect.isgenerator(genexpr): + raise ValueError( + f"Input must be a generator expression, got {type(genexpr).__name__}" + ) + if inspect.getgeneratorstate(genexpr) != inspect.GEN_CREATED: + raise ValueError( + "Input must be a generator expression that has not been started, " + f"got one in state {inspect.getgeneratorstate(genexpr)}" + ) + return ast.fix_missing_locations(ast.Expression(ensure_ast(genexpr))) diff --git a/effectful/internals/product_n.py b/effectful/internals/product_n.py index ca7c2f0b3..e96a2f33a 100644 --- a/effectful/internals/product_n.py +++ b/effectful/internals/product_n.py @@ -79,8 +79,8 @@ def recurse(x): else: result = type(expr)(recurse(tuple(expr.items()))) elif isinstance(expr, collections.abc.Sequence): - if isinstance(expr, str | bytes): - result = expr + if isinstance(expr, str | bytes | range): + return expr elif ( isinstance(expr, tuple) and hasattr(expr, "_fields") diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py new file mode 100644 index 000000000..b0b9256ae --- /dev/null +++ b/effectful/ops/monoid.py @@ -0,0 +1,2098 @@ +import collections.abc +import functools +import itertools +import operator +import typing +from collections import UserDict, defaultdict +from collections.abc import Callable, Generator, Iterable, Mapping, Sequence, Sized +from dataclasses import dataclass +from graphlib import TopologicalSorter +from typing import Annotated, Any + +import effectful.ops.syntax +from effectful.internals.runtime import interpreter +from effectful.ops.semantics import coproduct, evaluate, fvsof, fwd, handler, typeof +from effectful.ops.syntax import ( + ObjectInterpretation, + Scoped, + _MappingTerm, + _NumberTerm, + as_dict, + defdata, + deffn, + defop, + implements, + ite, + range_, + syntactic_eq, + syntactic_hash, +) +from effectful.ops.types import Expr, Interpretation, NotHandled, Operation, Term + +type Stream[T] = Iterable[T] + +type Streams = Mapping[Operation[[], Any], Stream[Any]] + +type Body[T] = ( + Iterable[T] + | Callable[..., Body[T]] + | T + | Mapping[Any, Body[T]] + | Interpretation[T, Body[T]] +) + + +def outer_stream(streams: Streams) -> Iterable[tuple[Operation, Stream, Streams]]: + """Returns the streams that can be ordered outermost in the loop nest as + well as the remaining streams in the nest. + + """ + stream_vars = set(streams.keys()) + pred = {k: fvsof(v) & stream_vars for k, v in streams.items()} + topo = TopologicalSorter(pred) + topo.prepare() + return ( + (op, streams[op], {k: v for (k, v) in streams.items() if k != op}) + for op in topo.get_ready() + ) + + +def inner_stream( + streams: dict[Operation, Expr], +) -> Iterable[tuple[dict[Operation, Expr], Operation, Expr]]: + """Returns the streams that can be ordered innermost in the loop nest as + well as the remaining streams in the nest. + + """ + stream_vars = set(streams.keys()) + + no_dependents = set() + succ = defaultdict(set) + for k, v in streams.items(): + preds = fvsof(v) & stream_vars + if preds: + for pred in preds: + succ[pred].add(k) + else: + no_dependents.add(k) + + topo = TopologicalSorter(succ) + topo.prepare() + return ( + ({k: v for (k, v) in streams.items() if k != op}, op, streams[op]) + for op in set(topo.get_ready()) | no_dependents + ) + + +def inner_streams_first(streams: dict[Operation, Expr]) -> Iterable[Operation]: + """Iterable over streams where dependent streams precede their dependencies.""" + stream_vars = set(streams.keys()) + + no_dependents = set() + succ = defaultdict(set) + for k, v in streams.items(): + preds = fvsof(v) & stream_vars + if preds: + for pred in preds: + succ[pred].add(k) + else: + no_dependents.add(k) + + topo = TopologicalSorter(succ) + return topo.static_order() + + +class Monoid[W]: + """A monoid with ``plus`` and ``reduce`` :class:`Operation` s.""" + + __name__: str + identity: W + + def __init__(self, identity: W, name: str): + self.__name__ = name + self.identity = identity + + def __repr__(self): + return f"Monoid({self.__name__!r}, {self.identity!r})" + + def __eq__(self, other): + return id(self) == id(other) + + def __hash__(self): + return hash(id(self)) + + @Operation.define + def plus(self, *args: W) -> W: + """Monoid addition. Handlers supply per-monoid and broadcasting + behavior; the default rule only handles identity and zero cases (for + monoids that have a zero). + + """ + if hasattr(self, "zero") and any(a is self.zero for a in args): + return self.zero + + nonident_args = [a for a in args if a is not self.identity] + if len(nonident_args) != len(args): + return self.plus(*nonident_args) + + return defdata(self.plus, *nonident_args) # type: ignore[return-value] + + @Operation.define + def reduce[A, B, U: Body]( + self, + body: Annotated[U, Scoped[A | B]], + streams: Annotated[Streams, Scoped[A]], + ) -> Annotated[U, Scoped[B]]: + """Reduce ``body`` over ``streams``. Handlers supply per-monoid and + broadcasting behavior. + + """ + raise NotHandled + + @Operation.define + def weighted[T]( + self, stream: Stream[T], weight: Callable[[T], W] | Operation[[T], W] + ) -> Stream[T]: + """A stream paired with a per-element weight. ``var`` is an + :class:`Operation` standing for "an element of ``stream``"; ``weight`` + is an expression that uses ``var`` and evaluates to the weight of that + element. + + """ + raise NotHandled + + @Operation.define + def mask(self, value: W, cond: Any = True) -> W: + raise NotHandled + + @Operation.define + def delta[K](self, index: K, weight: W) -> Mapping[K, W]: + raise NotHandled + + +class MonoidWithZero[T](Monoid[T]): + zero: T + + def __init__(self, name: str, identity: T, zero: T): + super().__init__(name=name, identity=identity) + self.zero = zero + + +Min = Monoid(name="Min", identity=float("inf")) +Max = Monoid(name="Max", identity=-float("inf")) +ArgMin = Monoid(name="ArgMin", identity=(Min.identity, None)) +ArgMax = Monoid(name="ArgMax", identity=(Max.identity, None)) +Sum = Monoid(name="Sum", identity=0) +Product = MonoidWithZero(name="Product", identity=1, zero=0) +CartesianProduct: MonoidWithZero[Sequence[Mapping]] = MonoidWithZero( + name="CartesianProduct", identity=[{}], zero=[] +) +Union: Monoid[Sequence[Mapping]] = Monoid(name="Union", identity=[]) +And = MonoidWithZero(name="And", identity=True, zero=False) +Or = Monoid(name="Or", identity=False) + + +def _conjuncts(mask) -> Sequence[Term]: + """Return the conjuncts of an ``And`` mask as a flat tuple.""" + match mask: + case Term(And.plus, elems, {}): + return elems + case _: + return (mask,) + + +def _is_simple_range(obj) -> bool: + return isinstance(obj, range) and obj.start == 0 and obj.step == 1 + + +@dataclass +class _ExtensiblePredicate[T]: + elems: set[T] + + def register(self, t: T) -> None: + self.elems.add(t) + + def __call__(self, t: T) -> bool: + return t in self.elems + + +is_commutative = _ExtensiblePredicate({Max, Min, Sum, Product, And, Or}) +is_idempotent = _ExtensiblePredicate({Max, Min, And, Or}) + + +@dataclass +class _ExtensibleBinaryRelation[S, T]: + tuples: set[tuple[S, T]] + + def __init__(self, *args): + self.tuples = set() + for s, t in args: + self.register(s, t) + + def register(self, s: S, t: T) -> None: + self.tuples.add((s, t)) + + def __call__(self, s: S, t: T) -> bool: + return (s, t) in self.tuples + + +class _ExtensiblePartialInvolution[S](_ExtensibleBinaryRelation[S, S]): + def register(self, s: S, t: S) -> None: + for existing, image in self.tuples: + if existing == s and image != t: + raise ValueError(f"{s!r} already matched to {image!r}") + if existing == t and image != s: + raise ValueError(f"{t!r} already matched to {image!r}") + super().register(s, t) + super().register(t, s) + + def of(self, t: S) -> S | None: + for a, b in self.tuples: + if a == t: + return b + return None + + +distributes_over: _ExtensibleBinaryRelation[Monoid, Monoid] = _ExtensibleBinaryRelation( + (Max, Min), + (Min, Max), + (Sum, Min), + (Sum, Max), + (Product, Sum), + (CartesianProduct, Union), + (And, Or), +) + + +def _is_monoid_plus(op: Operation) -> bool: + """True if ``op`` is the ``plus`` operation of some :class:`Monoid`.""" + owner = getattr(op, "__self__", None) + return isinstance(owner, Monoid) and op is owner.plus + + +def _is_monoid_reduce(op: Operation) -> bool: + """True if ``op`` is the ``reduce`` operation of some :class:`Monoid`.""" + owner = getattr(op, "__self__", None) + return isinstance(owner, Monoid) and op is owner.reduce + + +def _is_monoid_weighted(op: Operation) -> bool: + """True if ``op`` is the ``weighted`` operation of some :class:`Monoid`.""" + owner = getattr(op, "__self__", None) + return isinstance(owner, Monoid) and op is owner.weighted + + +def _is_monoid_mask(op: Operation) -> bool: + """True if ``op`` is the ``mask`` operation of some :class:`Monoid`.""" + owner = getattr(op, "__self__", None) + return isinstance(owner, Monoid) and op is owner.mask + + +def _is_monoid_delta(op: Operation) -> bool: + """True if ``op`` is the ``delta`` operation of some :class:`Monoid`.""" + owner = getattr(op, "__self__", None) + return isinstance(owner, Monoid) and op is owner.delta + + +class GetitemDelta(ObjectInterpretation): + """M.delta(i, v)[q] ≡ M.mask(v, i == q)""" + + @implements(_MappingTerm.__getitem__) + def _(self, value, index): + if isinstance(value, Term) and _is_monoid_delta(value.op): + return value.op.__self__.mask(value.args[1], value.args[0] == index) + return fwd() + + +class MaskBool(ObjectInterpretation): + @implements(Monoid.mask) + def _(self, monoid, value, cond): + if isinstance(cond, bool): + return value if cond else monoid.identity + return fwd() + + +class MaskFusion(ObjectInterpretation): + """M.mask(M.mask(value, i1), i2) ≡ M.mask(value, And.plus(i1, i2))""" + + @implements(Monoid.mask) + def _(self, monoid, value, cond): + if ( + isinstance(value, Term) + and _is_monoid_mask(value.op) + and value.op.__self__ == monoid + ): + return monoid.mask(value.args[0], And.plus(value.args[1], cond)) + return fwd() + + +is_equality = _ExtensiblePredicate({_NumberTerm.__eq__}) + + +class ReduceEqualityMaskRange(ObjectInterpretation): + """M.reduce(M.mask(v, And.plus(i = x, *m)), {i: range(N)} ∪ S) ≡ + M.mask(M.reduce(M.mask(v, *m), {i: [x]} ∪ S), And.plus(0 <= x, x < N)) + + The equality constraint ``i = x`` on a range-stream reduce is discharged by + a gather (the stream becomes the singleton ``[x]``) guarded by a bounds + check. + + When the reduce body is a ``plus`` of the same monoid, the rule distributes + the reduce over the plus -- but only when doing so exposes an eliminable + equality mask in some summand. This is a *targeted* split (it leaves + ``ReduceSplit`` conservative): summands whose reduced index appears in an + equality become gathers, while the rest stay as ordinary masked reduces. + """ + + @staticmethod + def _match_eq(cond, streams): + """If ``cond`` is ``stream_op == key`` (either order) where ``stream_op`` + is a ``range(0, N)`` stream and ``key`` is stream-independent, return + ``(stream_op, key)``; otherwise ``None``.""" + + def test(op, stream_op, mask_key): + return ( + is_equality(op) + and stream_op in streams + and _is_simple_range(streams[stream_op]) + and not (fvsof(mask_key) & set(streams)) + ) + + match cond: + case Term(op, (Term(stream_op, (), {}), mask_key), {}) if test( + op, stream_op, mask_key + ): + return (stream_op, mask_key) + case Term(op, (mask_key, Term(stream_op, (), {})), {}) if test( + op, stream_op, mask_key + ): + return (stream_op, mask_key) + case _: + return None + + def _eliminate(self, monoid, value, mask, streams): + """Discharge one eliminable equality constraint via a gather, or return + ``None`` if no constraint is eliminable.""" + conds = _conjuncts(mask) + for i, cond in enumerate(conds): + matched = self._match_eq(cond, streams) + if matched is None: + continue + stream_op, mask_key = matched + stream = streams[stream_op] + return monoid.reduce( + monoid.mask( + monoid.mask( + value, + And.plus(stream.start <= mask_key, mask_key < stream.stop), + ), + And.plus(*(c for (j, c) in enumerate(conds) if i != j)), + ), + {stream_op: (mask_key,)} + | {k: v for (k, v) in streams.items() if k != stream_op}, + ) + return None + + def _summand_eliminable(self, monoid, summand, streams): + return ( + isinstance(summand, Term) + and _is_monoid_mask(summand.op) + and summand.op.__self__ == monoid + and self._eliminate(monoid, summand.args[0], summand.args[1], streams) + is not None + ) + + @implements(Monoid.reduce) + def _(self, monoid, body, streams): + if not isinstance(body, Term): + return fwd() + + # single mask body: discharge an equality constraint directly + if _is_monoid_mask(body.op) and body.op.__self__ == monoid: + result = self._eliminate(monoid, body.args[0], body.args[1], streams) + return result if result is not None else fwd() + + # plus body: distribute the reduce only when it exposes an eliminable + # equality mask in some summand + if _is_monoid_plus(body.op) and body.op.__self__ == monoid: + if any(self._summand_eliminable(monoid, s, streams) for s in body.args): + return monoid.plus(*(monoid.reduce(s, streams) for s in body.args)) + + return fwd() + + +class ReduceMaskHoist(ObjectInterpretation): + """M.reduce(M.mask(v, c), S) ≡ M.mask(M.reduce(v, S), c) when ``c`` does not + depend on any stream in ``S``. + + A reduce-stream-independent condition gates the whole reduction uniformly, + so it can be lifted out: when ``c`` holds both sides are ``reduce(v, S)``, + and when it fails both are the identity (``reduce(identity, S) = identity``). + This holds for any monoid -- no commutativity or distributivity required. + """ + + @implements(Monoid.reduce) + def _(self, monoid, body, streams): + if not ( + streams + and isinstance(body, Term) + and _is_monoid_mask(body.op) + and body.op.__self__ == monoid + ): + return fwd() + value, cond = body.args + if fvsof(cond) & set(streams): + return fwd() + return monoid.mask(monoid.reduce(value, streams), cond) + + +class PlusEmpty(ObjectInterpretation): + """plus() = 0""" + + @implements(Monoid.plus) + def plus(self, monoid, *args): + if not args: + return monoid.identity + return fwd() + + +class PlusSingle(ObjectInterpretation): + """plus(x) = x""" + + @implements(Monoid.plus) + def plus(self, _, *args): + if len(args) == 1: + return args[0] + return fwd() + + +class PlusAssoc(ObjectInterpretation): + """x + (y + z) = (x + y) + z = x + y + z""" + + @implements(Monoid.plus) + def plus(self, monoid, *args): + def is_nested_plus(x): + return isinstance(x, Term) and x.op is monoid.plus + + if any(is_nested_plus(x) for x in args): + flat_args = itertools.chain.from_iterable( + t.args if is_nested_plus(t) else (t,) for t in args + ) + assert len(args) > 0 + return monoid.plus(*flat_args) + return fwd() + + +class PlusDistr(ObjectInterpretation): + """x * (y + z) = x * y + x * z""" + + @implements(Monoid.plus) + def plus(self, monoid: Monoid, *args): + if len(args) < 2: + return fwd() + + for i, a in enumerate(args): + if ( + isinstance(a, Term) + and _is_monoid_plus(a.op) + and distributes_over(monoid, (inner_monoid := a.op.__self__)) + and not distributes_over(inner_monoid, monoid) + ): + if i > 0: + return monoid.plus( + *args[: i - 1], + inner_monoid.plus( + *(monoid.plus(args[i - 1], x) for x in a.args) + ), + *args[i + 1 :], + ) + else: + return monoid.plus( + inner_monoid.plus( + *(monoid.plus(x, args[i + 1]) for x in a.args) + ), + *args[i + 2 :], + ) + return fwd() + + +class PlusOrder(ObjectInterpretation): + """Normalize plus ordering for commutative monoids. + + x ⊕ y ⊕ x = x ⊕ x ⊕ y + + """ + + @staticmethod + def _term_sort_key(t: Term) -> tuple[int, int]: + return (syntactic_hash(t), id(t)) + + @implements(Monoid.plus) + def plus(self, monoid, *args): + if not is_commutative(monoid): + return fwd() + + sorted_args = tuple( + sorted(range(len(args)), key=lambda i: self._term_sort_key(args[i])) + ) + if sorted_args == tuple(range(len(args))): + return fwd() + return monoid.plus(*(args[i] for i in sorted_args)) + + +class PlusConsecutiveDups(ObjectInterpretation): + """Normalize duplicate arguments for idempotent monoids. + + x ⊕ x ⊕ y = x ⊕ y + + """ + + @implements(Monoid.plus) + def plus(self, monoid, *args): + if not is_idempotent(monoid): + return fwd() + + dedup_args = tuple( + i + for i in range(len(args)) + if i == 0 or not syntactic_eq(args[i - 1], args[i]) + ) + if dedup_args == tuple(range(len(args))): + return fwd() + return monoid.plus(*(args[i] for i in dedup_args)) + + +class ReduceEmpty(ObjectInterpretation): + @implements(Monoid.reduce) + def _(self, monoid, body, streams): + if streams: + return fwd() + + return body + + +class ReducePartial(ObjectInterpretation): + unrolled: collections.abc.Set[Operation] | None + + def __init__(self, unrolled: collections.abc.Set[Operation] | None = None): + self.unrolled = unrolled + + @implements(Monoid.reduce) + def _(self, monoid, body, streams): + if not streams: + return fwd() + + for stream_key, stream_body, streams_tail in outer_stream(streams): + if self.unrolled is not None and stream_key not in self.unrolled: + continue + if isinstance(stream_body, Term): + continue + stream_values_iter = iter(stream_body) + + # if we iterate and get a term instead of a real iterator, skip + if isinstance(stream_values_iter, Term): + continue + + new_reduces = [] + for stream_val in stream_values_iter: + with handler({stream_key: deffn(stream_val)}): + eval_args = evaluate((body, streams_tail)) + assert isinstance(eval_args, tuple) + new_reduces.append( + monoid.reduce(*eval_args) if streams_tail else eval_args[0] + ) + return monoid.plus(*new_reduces) + return fwd() + + +class ReduceFusion(ObjectInterpretation): + """Implements the identity + reduce(R, S1, reduce(R, S2, body)) = reduce(R, S1 ∪ S2, body) + """ + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + if isinstance(body, Term) and body.op is monoid.reduce: + return monoid.reduce(body.args[0], streams | body.args[1]) + return fwd() + + +class ReduceSplit(ObjectInterpretation): + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + if not is_commutative(monoid): + return fwd() + if not (isinstance(body, Term) and body.op is monoid.plus): + return fwd() + + return monoid.plus(*(monoid.reduce(a, streams) for a in body.args)) + + +@Operation.define +def choose_contraction(factors: Sequence[Any], streams: Streams) -> Operation: + """Used by `ReduceFactorization` to choose a contraction when there is + ambiguity. Takes the factors and streams that are eligible for contraction + (innermost and non-universal). + + The default behavior is to return the first support-minimal stream in the + streams dictionary. + + """ + assert len(streams) > 0 + + factors = [(a, fvsof(a)) for a in factors] + support: dict = { + k: frozenset(i for i, (_, fvs) in enumerate(factors) if k in fvs) + for k in streams + } + for v, f_v in support.items(): + if any(u_sup < f_v for u, u_sup in support.items() if u is not v): + continue + return v + assert False, "expected at least one subset-minimal stream" + + +class Factor(ObjectInterpretation): + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + """reduce(⊗(F_v ∪ F_rest), {v} ∪ S) = reduce(⊗F_rest ⊗ reduce(⊗F_v, {v}), S) + + where F_v = factors mentioning v, F_rest = the others. Fires only when + v has no dependents among the remaining streams (so it can be innermost) + and F_rest is nonempty (universal variables stay in the outer core). + + A reduce-monoid mask wrapping the plus is handled too: + + reduce(M.mask(⊗(F_v ∪ F_rest), c), {v} ∪ S) + = reduce(M.mask(⊗F_rest ⊗ reduce(⊗F_v, {v}), c), S) + + Because the mask gates the whole product, every factor effectively depends + on the condition's variables; folding ``fvsof(c)`` into each factor's free + variables keeps those streams in the outer core (a stream the mask depends + on is treated as universal and never pulled into the inner reduce). The + mask therefore stays outside the inner reduce unchanged. Soundness relies + on ``M.identity`` annihilating the inner monoid's plus (the semiring zero), + so masking distributes over the inner product. + """ + if not (is_commutative(monoid) and isinstance(body, Term)): + return fwd() + + # Optionally peel an outer mask of the reduce monoid. + cond = None + plus_term = body + if _is_monoid_mask(body.op) and body.op.__self__ is monoid: + plus_term, cond = body.args + + if not ( + isinstance(plus_term, Term) + and _is_monoid_plus(plus_term.op) + and distributes_over(plus_term.op.__self__, monoid) + ): + return fwd() + + inner = plus_term.op.__self__ + stream_keys = set(streams) + cond_fvs = fvsof(cond) if cond is not None else set() + factors = [(a, fvsof(a) | cond_fvs) for a in plus_term.args] + + # candidates: innermost-eligible (no remaining stream depends on v), + # non-universal (some factor doesn't mention v) + eligible = {} + for k, v in streams.items(): + if any(k in fvsof(vv) for kk, vv in streams.items() if k is not kk): + continue + if len({i for i, (_, fvs) in enumerate(factors) if k in fvs}) == len( + factors + ): + continue # v is universal: leave it in the outer core + eligible[k] = v + + if not eligible: + return fwd() + if len(eligible) == 1: + inner_stream = next(iter(eligible)) + else: + inner_stream = choose_contraction(plus_term.args, eligible) + + inner_factor_ids = frozenset( + i for i, (_, fvs) in enumerate(factors) if inner_stream in fvs + ) + + inner_factors = [factors[i][0] for i in sorted(inner_factor_ids)] + inner_stream_keys = {inner_stream} + inner_deps = set().union( + *(factors[i][1] for i in inner_factor_ids), + fvsof(streams[inner_stream]) & stream_keys, + ) + + outer_factors = [ + a for i, (a, _) in enumerate(factors) if i not in inner_factor_ids + ] + outer_stream_keys = stream_keys - inner_stream_keys + outer_factor_deps = set().union( + *(vars for i, (_, vars) in enumerate(factors) if i not in inner_factor_ids) + ) + + # find all streams that are used in the inner factors/streams and are + # not used by the outer factors/streams + # this has to be done iteratively, because moving a stream inward + # reduces the outer dependency set + # ensures that no future factorization application creates a reduce that + # fuses with with the inner reduce + for s in inner_streams_first(streams): + outer_stream_deps = ( + set().union(*(fvsof(streams[k]) for k in outer_stream_keys)) + & stream_keys + ) + outer_deps = outer_factor_deps | outer_stream_deps + if s in inner_deps and s not in outer_deps: + inner_stream_keys |= {s} + inner_deps |= stream_keys & fvsof(streams[s]) + outer_stream_keys -= {s} + + inner_streams = {k: v for (k, v) in streams.items() if k in inner_stream_keys} + inner_red = monoid.reduce(inner.plus(*inner_factors), inner_streams) + + rest_streams = {k: s for k, s in streams.items() if k in outer_stream_keys} + new_body = inner.plus(*outer_factors, inner_red) + if cond is not None: + new_body = monoid.mask(new_body, cond) + return monoid.reduce(new_body, rest_streams) if rest_streams else new_body + + +class ReduceUnfactor(ObjectInterpretation): + """Undo one :class:`Factor` layer beneath a compatible product. + + This rule is deliberately kept out of ``NormalizeIntp``: together with + ``Factor`` it would form a rewrite cycle. It is used only while preparing + a candidate for cartesian-product inversion. + """ + + @implements(Monoid.reduce) + def reduce(self, monoid: Monoid, body, streams): + if not is_commutative(monoid): + return fwd() + if not (isinstance(body, Term) and _is_monoid_plus(body.op)): + return fwd() + + product_monoid = body.op.__self__ + if not distributes_over(product_monoid, monoid): + return fwd() + + for i, factor in enumerate(body.args): + if not (isinstance(factor, Term) and factor.op is monoid.reduce): + continue + + inner_body, inner_streams = factor.args + assert isinstance(inner_streams, Mapping) + + merged_streams = dict(streams) | dict(inner_streams) + + inner_factors = ( + inner_body.args + if isinstance(inner_body, Term) and inner_body.op is product_monoid.plus + else (inner_body,) + ) + return monoid.reduce( + product_monoid.plus( + *body.args[:i], *inner_factors, *body.args[i + 1 :] + ), + merged_streams, + ) + return fwd() + + +class ReduceDistributeCartesianProduct(ObjectInterpretation): + """Eliminates a reduce over a cartesian product. + ∑_x₁ ∑_x₂ ... ∑_xₙ ∏_i f(xᵢ) = ∏_i ∑_xᵢ f(xᵢ) + This transform is also called inversion in the lifting + literature (e.g. [1]). + + More specifically, this transform implements the identity + reduce(⨁, reduce(⨂, body2, {vv: v()}), {v: reduce(×, body1, S1)} ∪ S2) + = reduce(⨁, reduce(⨂, reduce(⨁, body2, {vv: body1}), S1), S2) + where × is the cartesian product and ⨂ distributes over ⨁. + + The body may also be a ``⨂``-plus of such reductions. Each reduction is + row-substituted independently, and the row positions determine which plate + variables are unified. Ordinary row-independent factors remain outside the + peeled plate reduction. + + Note: This could be generalized to grouped inversion [2]. + + [1] Braz, Rd, Eyal Amir, and Dan Roth. "Lifted first-order + probabilistic inference." IJCAI. 2005. + [2] Taghipour, Nima, et al. "Completeness results for lifted + variable elimination." AISTATS. 2013. + """ + + @implements(Monoid.reduce) + def reduce(self, monoid: Monoid, body, streams): + if not isinstance(body, Term): + return fwd() + + if not any( + isinstance(v, Term) and v.op == CartesianProduct.reduce + for v in streams.values() + ): + return fwd() + + # ``Factor`` may have moved product factors into nested additive + # reductions. Normalize the whole candidate so ReduceUnfactor can see + # and merge both stream bundles. This isolated interpretation cannot + # cycle with Factor. Build the candidate outside the active handler; + # otherwise invoking ``monoid.reduce`` here would recursively redispatch + # this rule. + with interpreter(CartesianProductNormalizeIntp): + candidate = monoid.reduce(body, streams) + body, streams = candidate.args + + inner_reduces: tuple[Term, ...] + if isinstance(body, Term) and _is_monoid_reduce(body.op): + inner_reduces = (body,) + outer_factors: tuple = () + inner_monoid = body.op.__self__ + elif isinstance(body, Term) and _is_monoid_plus(body.op): + inner_monoid = body.op.__self__ + inner_reduces = tuple( + arg + for arg in body.args + if isinstance(arg, Term) and arg.op is inner_monoid.reduce + ) + if not inner_reduces: + return fwd() + outer_factors = tuple( + arg + for arg in body.args + if not (isinstance(arg, Term) and arg.op is inner_monoid.reduce) + ) + else: + return fwd() + + if not distributes_over(inner_monoid, monoid): + return fwd() + + class InvalidIndexError(Exception): ... + + def drop_elem(ls, index): + return tuple(x for (i, x) in enumerate(ls) if i != index) + + for stream_key, stream_body in streams.items(): + # stream is cartesian + if not ( + isinstance(stream_body, Term) + and stream_body.op is CartesianProduct.reduce + ): + continue + + # Product reductions that do not use this row are ordinary outer + # factors. Every factor that does use it must be plate-reduced. + row_inner_reduces = tuple( + reduce for reduce in inner_reduces if stream_key in fvsof(reduce) + ) + row_outer_factors = outer_factors + tuple( + reduce for reduce in inner_reduces if stream_key not in fvsof(reduce) + ) + if not row_inner_reduces or stream_key in fvsof(row_outer_factors): + continue + + (cprod_body, cprod_streams) = stream_body.args + assert isinstance(cprod_streams, dict) + + # plates are rectangular + if not all( + isinstance(plate_stream, range) + for plate_stream in cprod_streams.values() + ): + continue + + # stream body is a sequence of mappings from plate index to domain value + match cprod_body: + case Term( + Union.reduce, + ( + [Term(effectful.ops.syntax.as_dict, ((idx, union_body),), {})], + union_streams, + ), + {}, + ) if isinstance(idx, Sequence) and set( + i.op for i in idx if isinstance(i, Term) + ) >= set(cprod_streams): + pass + case _: + continue + + assert len(idx) > 0 + + # inner product folds over all plates + plate_index, plate_op = next( + (j, i.op) for (j, i) in enumerate(idx) if i.op in cprod_streams + ) + plate_range = cprod_streams[plate_op] + + def row_substitute(inner_body, inner_streams): + """Peel ``plate_index`` off every ``stream_key[...]`` in one + summand. Asserts the summand's bundle contains a stream that + folds over the full cartesian product row and returns that plate + variable alongside the substituted body.""" + if stream_key in fvsof(inner_streams): + raise InvalidIndexError() + + inner_plate_op = None + + # substitute all instances of row[i, *rest] -> row[*rest] + def _getitem(mapping, idx1): + nonlocal inner_plate_op + + idx1 = idx1 if isinstance(idx1, Sequence) else (idx1,) + if isinstance(mapping, Term) and mapping.op == stream_key: + if not ( + isinstance(idx1, Sequence) + and len(idx1) > plate_index + and isinstance(idx1[plate_index], Term) + and idx1[plate_index].op in inner_streams + and syntactic_eq( + inner_streams[idx1[plate_index].op], plate_range + ) + ): + raise InvalidIndexError() + + if inner_plate_op is None: + inner_plate_op = idx1[plate_index].op + elif inner_plate_op != idx1[plate_index].op: + raise InvalidIndexError() + + return fwd(mapping, drop_elem(idx1, plate_index)) + return fwd() + + subst = handler({_MappingTerm.__getitem__: _getitem})(evaluate)( + inner_body + ) + if inner_plate_op is None and inner_streams: + # A nontrivial reduction that does not fold over the row + # cannot be unified with the peeled plate. Streamless + # reductions are ordinary product factors introduced by + # sum-of-products expansion and are retained unchanged. + raise InvalidIndexError() + return subst, inner_plate_op + + try: + substituted = [ + row_substitute(*reduce.args) for reduce in row_inner_reduces + ] + except InvalidIndexError: + continue + + # Unify reductions by the variable used at this row position, not + # merely by equal ranges. Equal-sized axes are otherwise ambiguous. + shared_plate_op = plate_op + combined_factors = [] + for (subst_body, inner_plate_op), inner_reduce in zip( + substituted, row_inner_reduces + ): + _, inner_streams = inner_reduce.args + assert isinstance(inner_streams, Mapping) + assert inner_plate_op is not None + if inner_plate_op is not shared_plate_op: + subst_body = handler({inner_plate_op: shared_plate_op})(evaluate)( + subst_body + ) + + inner_tail_streams = { + k: v for (k, v) in inner_streams.items() if k != inner_plate_op + } + subst_body = inner_monoid.reduce(subst_body, inner_tail_streams) + combined_factors.append(subst_body) + + combined = ( + combined_factors[0] + if len(combined_factors) == 1 + else inner_monoid.plus(*combined_factors) + ) + + peeled_idx = drop_elem(idx, plate_index) + peeled_cprod_streams = { + k: v for (k, v) in cprod_streams.items() if k != plate_op + } + if not peeled_cprod_streams and not peeled_idx: + + def _to_body(mapping, key): + if isinstance(mapping, Term) and mapping.op == stream_key: + return union_body + return fwd() + + subst_combined = handler({_MappingTerm.__getitem__: _to_body})( + evaluate + )(combined) + inner_reduce_body = monoid.reduce(subst_combined, union_streams) + else: + peeled_body = Union.reduce( + [as_dict((peeled_idx, union_body))], union_streams + ) + peeled_cprod = CartesianProduct.reduce( + peeled_body, peeled_cprod_streams + ) + inner_reduce_body = monoid.reduce(combined, {stream_key: peeled_cprod}) + + peeled_reduce = inner_monoid.reduce( + inner_reduce_body, {shared_plate_op: plate_range} + ) + result_body = inner_monoid.plus(peeled_reduce, *row_outer_factors) + + # Include any extra sum streams outermost. In particular, the + # non-reduce product factors above remain outside the plate fold. + tail_streams = {k: v for (k, v) in streams.items() if k != stream_key} + result = monoid.reduce(result_body, tail_streams) + return result + + return fwd() + + +class _GroundRow(collections.abc.Mapping): + """A cartesian-product row that mints a variable per index it is asked for. + + Standing this in for a row variable passes a concrete key straight + to :meth:`__getitem__` here, which answers with the variable for that plate + assignment -- minting it, and the stream it ranges over, on first sight. + Asking twice for the same assignment gives the same variable, which is what + ties the factors of a chain together. + + A key that is still symbolic never reaches here at all: the subscript stays + a term until the fold supplying the index is expanded, after which the same + traversal passes over it again with the key concrete. + """ + + variables: dict[tuple, Operation[[], typing.Any]] + streams: dict[Operation[[], typing.Any], collections.abc.Iterable] + sealed: bool + + def __init__( + self, + cprod_body: collections.abc.Iterable, + cprod_streams: dict[Operation, collections.abc.Iterable], + ): + """cprod_body is the body and cprod_streams is ths streams of a CartesianProduct.reduce expression""" + shape = self._row_shape(cprod_body) + assert shape is not None + (self._index, self._value), self._value_streams = shape + self._plates = cprod_streams + self.variables = {} + self.streams = {} + # A subscript whose index is not concrete yet holds the row until the + # fold supplying that index is expanded, so while grounding is under + # way the row does legitimately sit inside a term. + self.sealed = False + + def __getitem__(self, key: tuple) -> Term: + key = key if isinstance(key, tuple) else (key,) + if key not in self.variables: + fresh = defop(self._value.op) + self.streams[fresh] = self._value_streams[self._value.op] + self.variables[key] = fresh + return self.variables[key]() + + @staticmethod + def _row_shape(cprod_body) -> tuple[tuple[tuple, Term], Mapping] | None: + """The index and value of a cartesian product's rows, with the value's + domains, or ``None`` if the body is not a row stream. + + A row stream is a union of one-entry rows: for each value its domain takes, + a row mapping ``idx`` to that value. The value must be a bare call of its + domain variable, since a key's variable is that variable renamed. + """ + match cprod_body: + case Term( + Union.reduce, + ( + [Term(effectful.ops.syntax.as_dict, ((idx, value),), {})], + Mapping() as value_streams, + ), + {}, + ) if ( + isinstance(idx, Sequence) + and all(isinstance(i, Term) and not i.args for i in idx) + and isinstance(value, Term) + and not value.args + and value.op in value_streams + ): + return (tuple(idx), value), value_streams + case _: + return None + + def residual(self) -> Mapping | None: + """What is left of the row once the keys the body used are taken out. + + Grounding consumes the keys the body subscripts; the rest are still + assigned a value by every row, so they still multiply the number of + rows. Rather than accounting for them, they stay in the expression -- + the same cartesian product over the plate values no key consumed, which + :class:`ReduceDistributeCartesianProduct` also leaves behind when it + peels one plate off a product over several. + + ``None`` if what is left cannot be written as a cartesian product: the + complement of a set of keys is only a product of plate ranges when + there is one plate. + """ + consumed = { + plate: {key[position] for key in self.variables} + for position, plate in enumerate(i.op for i in self._index) + } + remaining = { + plate: [v for v in values if v not in consumed.get(plate, ())] + for plate, values in self._plates.items() + } + if not any(remaining.values()): + return {} + if len(self._plates) > 1: + return None + return {plate: values for (plate, values) in remaining.items() if values} + + def __iter__(self): + return iter(()) + + def __len__(self) -> int: + return 0 + + +@evaluate.register +def _(expr: _GroundRow): + if expr.sealed: + raise ValueError(f"{expr} should not have lived long enough to get here") + return expr + + +class ReduceGroundCartesianProduct(ObjectInterpretation): + """Replace a cartesian-product stream with one stream per plate assignment. + + reduce(M, body, {X: CartesianProduct.reduce(Union.reduce([as_dict((idx, w))], D), P)} ∪ S) + ═══════════════════════════════════════════════════════════════════════════ + reduce(M, body[X[a] := w_a], {w_a: D[idx := a]} ∪ S) + + where ``a`` ranges over the plate assignments the body actually subscripts. + + Where the body folds over the plates uniformly, + :class:`ReduceDistributeCartesianProduct` inverts the reduction instead, + which is cheaper and leaves the plate fold intact. This rule is the fallback + for bodies that admit no per-plate factorization -- a chain ``X[t], X[t+1]`` + couples adjacent plate indices, so inversion does not apply. Grounding it + costs one variable per assignment (``|P|`` variables over ``D``) instead of + enumerating the ``|D|^|P|`` rows, after which the resulting reduction over + ordinary range streams is an ordinary variable-elimination problem that + :class:`Factor` solves. + """ + + @staticmethod + def _plate_folds(body, monoid) -> collections.abc.Set[Operation]: + """The stream variables of the product folds nested in this reduction""" + folds: tuple[Term, ...] + if isinstance(body, Term) and _is_monoid_reduce(body.op): + inner_monoid, folds = body.op.__self__, (body,) + elif isinstance(body, Term) and _is_monoid_plus(body.op): + inner_monoid = body.op.__self__ + folds = tuple( + arg + for arg in body.args + if isinstance(arg, Term) and arg.op is inner_monoid.reduce + ) + else: + return frozenset() + + if not folds or not distributes_over(inner_monoid, monoid): + return frozenset() + return {var for fold in folds for var in typing.cast(Mapping, fold.args[1])} + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + plate_vars = self._plate_folds(body, monoid) + if not plate_vars: + return fwd() + + for var, stream in streams.items(): + if ( + isinstance(stream, Term) + and stream.op is CartesianProduct.reduce + and _GroundRow._row_shape(stream.args[0]) is not None + ): + row = _GroundRow(stream.args[0], stream.args[1]) + else: + continue + + rest = {k: v for (k, v) in streams.items() if k is not var} + with handler(ReducePartial(plate_vars)), handler({var: lambda: row}): + grounded, tail = evaluate((body, rest)) + + # Past this point a row in the result is a subscript whose index + # never became concrete, which is a half-ground term rather than a + # step of the rewrite. + row.sealed = True + try: + fvsof((grounded, tail)) + except ValueError: + continue + + # Whatever the body did not consume stays a cartesian product over + # the plate values no key took, still bound to the row variable, so + # its multiplicity is carried by the expression rather than counted + # here. Nothing is left when every key was consumed, which is the + # ordinary case. + plates = row.residual() + if plates is None or not row.variables: + # Nothing was consumed, so rebinding the row would reproduce + # the term this rule was given. + continue + residual = ( + {var: CartesianProduct.reduce(stream.args[0], plates)} if plates else {} + ) + return monoid.reduce(grounded, {**row.streams, **residual, **tail}) + return fwd() + + +class ReduceUnion(ObjectInterpretation): + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + for k, v in streams.items(): + if isinstance(v, Term) and v.op == Union.reduce: + union_body, union_streams = v.args + return monoid.reduce(body, streams | {k: union_body} | union_streams) + return fwd() + + +class ReduceWeightedStream(ObjectInterpretation): + """reduce(M, body, {x: WM.weighted(s, v, w), ...}) = reduce(M, WM.plus(w[v:=x()], body), {x: s, ...}) + + requires distributes_over(WM, M). + + The substitution ``v -> x`` is done by beta-reducing ``deffn(w, v)`` on + ``x()`` — symbolic, no Python dispatch on the weight expression. + """ + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + for k, v in streams.items(): + if isinstance(v, Term) and _is_monoid_weighted(v.op): + v_stream, v_weight = v.args + v_monoid = v.op.__self__ + if not distributes_over(v_monoid, monoid): + continue + w_at_k = v_weight(k()) + weighted_body = v_monoid.plus(w_at_k, body) + new_streams = {**streams, k: v_stream} + return monoid.reduce(weighted_body, new_streams) + return fwd() + + +class MonoidOverCallable(ObjectInterpretation): + """``monoid.reduce(f, streams) = lambda *a: monoid.reduce(f(*a), streams)``.""" + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + if isinstance(body, Term) or not isinstance(body, Callable): + return fwd() + return lambda *a, **k: monoid.reduce(body(*a, **k), streams) + + @implements(Monoid.plus) + def plus(self, monoid, *args): + if not args or any( + isinstance(arg, Term) or not isinstance(arg, Callable) for arg in args + ): + return fwd() + return lambda *a, **k: monoid.plus(*(arg(*a, **k) for arg in args)) + + +class MonoidOverMapping(ObjectInterpretation): + """``monoid.reduce({k: v_k}, streams) = {k: monoid.reduce(v_k, streams)}``.""" + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + if isinstance(body, Term) or not isinstance(body, Mapping): + return fwd() + return {k: monoid.reduce(v, streams) for (k, v) in body.items()} + + @implements(Monoid.plus) + def plus(self, monoid, *args): + if not args or not isinstance(args[0], Mapping): + return fwd() + + if isinstance(args[0], Interpretation): + keys = args[0].keys() + for b in args[1:]: + if not isinstance(b, Interpretation): + raise TypeError(f"Expected interpretation but got {b}") + if not keys == b.keys(): + raise ValueError( + f"Expected interpretation of {keys} but got {b.keys()}" + ) + return {k: monoid.plus(*(handler(b)(b[k]) for b in args)) for k in keys} + + for b in args[1:]: + if not isinstance(b, Mapping): + raise TypeError(f"Expected mapping but got {b}") + all_values = collections.defaultdict(list) + for d in args: + for k, v in d.items(): + all_values[k].append(v) + return {k: monoid.plus(*vs) for (k, vs) in all_values.items()} + + +def _scalar_args(args): + """True iff ``args`` is non-empty and every arg is a concrete int/float.""" + return ( + bool(args) + and not any(isinstance(x, Term) for x in args) + and all(isinstance(x, int | float) for x in args) + ) + + +class DeltaConcrete(ObjectInterpretation): + @implements(Monoid.delta) + def _(self, _, k, v): + if not fvsof((k, v)): + return {k: v} + return fwd() + + +class SumPlus(ObjectInterpretation): + """Scalar implementation of :data:`Sum`.""" + + @implements(Sum.plus) + def plus(self, *args): + if not _scalar_args(args): + return fwd() + return sum(args) + + +class MinPlus(ObjectInterpretation): + """Scalar implementation of :data:`Min`.""" + + @implements(Min.plus) + def plus(self, *args): + if not _scalar_args(args): + return fwd() + return min(args) + + +class MaxPlus(ObjectInterpretation): + """Scalar implementation of :data:`Max`.""" + + @implements(Max.plus) + def plus(self, *args): + if not _scalar_args(args): + return fwd() + return max(args) + + +class ProductPlus(ObjectInterpretation): + """Scalar implementation of :data:`Product`.""" + + @implements(Product.plus) + def plus(self, *args): + if not _scalar_args(args): + return fwd() + return functools.reduce(operator.mul, args) + + +class ArgMinPlus(ObjectInterpretation): + """Scalar score implementation of :data:`ArgMin`.""" + + @implements(ArgMin.plus) + def plus(self, *args): + if not args or not all(isinstance(a, tuple) for a in args): + return fwd() + if any(isinstance(a[0], Term) for a in args): + return fwd() + if not all(isinstance(a[0], int | float) for a in args): + return fwd() + return min(args, key=lambda a: a[0]) + + +class ArgMaxPlus(ObjectInterpretation): + """Scalar score implementation of :data:`ArgMax`.""" + + @implements(ArgMax.plus) + def plus(self, *args): + if not args or not all(isinstance(a, tuple) for a in args): + return fwd() + if any(isinstance(a[0], Term) for a in args): + return fwd() + if not all(isinstance(a[0], int | float) for a in args): + return fwd() + return max(args, key=lambda a: a[0]) + + +def _disjoint_merge[K, V](*dicts: Mapping[K, V]) -> Mapping[K, V]: + merged = {} + for d in dicts: + if isinstance(d, Term): + return fwd() + for key, value in d.items(): + if key in merged: + raise ValueError(f"Duplicate key found: '{key}'") + merged[key] = value + return merged + + +class CartesianProductPlus(ObjectInterpretation): + """Pure-Python implementation of :data:`CartesianProduct`.""" + + @implements(CartesianProduct.plus) + def plus(self, *args): + if not args: + return fwd() + if any(isinstance(x, Term) for x in args): + return fwd() + if not all(isinstance(x, Iterable) for x in args): + return fwd() + return [_disjoint_merge(*vals) for vals in itertools.product(*args)] + + +class UnionPlus(ObjectInterpretation): + @implements(Union.plus) + def plus(self, *args): + if not args: + return fwd() + if any(isinstance(x, Term) for x in args): + return fwd() + if not all(isinstance(x, Iterable) for x in args): + return fwd() + return list(itertools.chain(*args)) + + +is_scalar = _ExtensiblePredicate({Min, Max, Sum, Product, And, Or}) + + +class MonoidOverSequence(ObjectInterpretation): + @implements(Monoid.plus) + def plus(self, monoid, *args): + if ( + not is_scalar(monoid) + or not args + or not isinstance(args[0], tuple | list | Generator) + ): + return fwd() + zipped = zip(*args, strict=True) + result = (monoid.plus(*vs) for vs in zipped) + if isinstance(args[0], tuple | list): + return type(args[0])(result) + return result + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + if not is_scalar(monoid) or not isinstance(body, tuple | list | Generator): + return fwd() + result = (monoid.reduce(x, streams) for x in body) + if isinstance(body, tuple | list): + return type(body)(result) + return result + + +@Operation.define +def as_float(x: int) -> float: + if isinstance(x, Term): + raise NotHandled + return float(x) + + +class PlusCastFloat(ObjectInterpretation): + @implements(Monoid.plus) + def plus(self, monoid, *args): + typs = [typeof(a) for a in args] + if any(issubclass(t, float) for t in typs) and any( + issubclass(t, int) for t in typs + ): + args = [ + as_float(a) if issubclass(t, int) else a + for (a, t) in zip(args, typs, strict=True) + ] + return monoid.plus(*args) + return fwd() + + +class EliminateSingletonStreams(ObjectInterpretation): + """Eliminate a length-1 stream by substituting its sole element. + + reduce(M, body, {k: (v,)} ∪ S) = reduce(M, body[k := v], S[k := v]) + + Fires only when the sole element ``v`` is a :class:`Term`, i.e. a *symbolic* + singleton. This is exactly the form ``ReduceArrayGather`` produces (a gather + ``(a[i()],)``) and, more generally, every dependent singleton that + :class:`ReducePartial` cannot peel -- a non-outermost stream whose element + references another stream var. Concrete enumerated streams (``[0]``, + ``range(1)``) and monoid sentinels (``CartesianProduct.identity == [()]``) + have non-``Term`` elements and are left to ``ReducePartial`` / the + per-monoid rules. + + Unlike ``ReducePartial``, this peels the stream wherever it sits in the loop + nest and substitutes symbolically rather than unrolling, leaving a + vectorized index range (e.g. the gather's range) intact instead of + materializing it. + """ + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + # Eliminate *all* symbolic length-1 streams in one pass via a + # simultaneous substitution. Doing them together (rather than one per + # invocation) keeps an interleaving reduction rule -- e.g. + # ``ReduceArray`` consuming a now-live index range -- from firing + # between eliminations, so sibling index ranges stay together and fuse + # into a single reduction. + singletons = { + k: vs[0] + for k, vs in streams.items() + if not isinstance(vs, Term) + and isinstance(vs, collections.abc.Sequence) + and len(vs) == 1 + } + if not singletons: + return fwd() + + subs = {k: deffn(v) for k, v in singletons.items()} + new_body = handler(subs)(evaluate)(body) + new_streams = { + kk: handler(subs)(evaluate)(vv) + for kk, vv in streams.items() + if kk not in singletons + } + # reduce over no streams is a single (empty) assignment, i.e. the body + # itself -- not the monoid identity. + return monoid.reduce(new_body, new_streams) + + +class WhereHoist(ObjectInterpretation): + """Hoist :func:`ite` out of monoid ``reduce`` and ``plus``. + + A stream-independent selection commutes with reduction, while monoid + addition distributes pointwise over either selected branch. + """ + + @implements(Monoid.plus) + def plus(self, monoid, *args): + if len(args) < 2: + return fwd() + + for i, arg in enumerate(args): + if not ( + isinstance(arg, Term) + and arg.op is ite + and len(arg.args) == 3 + and not arg.kwargs + ): + continue + + cond, when_true, when_false = arg.args + return ite( + cond, + monoid.plus(*args[:i], when_true, *args[i + 1 :]), + monoid.plus(*args[:i], when_false, *args[i + 1 :]), + ) + + return fwd() + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + if not ( + streams + and isinstance(body, Term) + and body.op is ite + and len(body.args) == 3 + and not body.kwargs + ): + return fwd() + + cond, when_true, when_false = body.args + if fvsof(cond) & set(streams): + return fwd() + + return ite( + cond, + monoid.reduce(when_true, streams), + monoid.reduce(when_false, streams), + ) + + +class ReduceWhereEqualityPeel(ObjectInterpretation): + """Peel stream-independent conjuncts off a ``where`` equality guard. + + Given a condition ``outer & inner`` where ``outer`` is independent of the + reduced streams and ``inner`` contains an equality selecting one of them:: + + M.reduce(ite(outer & inner, x, y), S) + == ite(outer, + M.reduce(ite(inner, x, y), S), + M.reduce(y, S)) + + This exposes the stream equality to gather/elimination rules while moving + the remaining output guard above the reduction. + """ + + @staticmethod + def _stream_equality(cond, streams): + def matches(op, stream_term, other): + return ( + is_equality(op) + and isinstance(stream_term, Term) + and not stream_term.args + and not stream_term.kwargs + and stream_term.op in streams + and not (fvsof(other) & set(streams)) + ) + + return ( + isinstance(cond, Term) + and len(cond.args) == 2 + and ( + matches(cond.op, cond.args[0], cond.args[1]) + or matches(cond.op, cond.args[1], cond.args[0]) + ) + ) + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + if not ( + streams + and isinstance(body, Term) + and body.op is ite + and len(body.args) == 3 + and not body.kwargs + ): + return fwd() + + cond, when_true, when_false = body.args + if not (isinstance(cond, Term) and cond.op is And.plus): + return fwd() + + stream_ops = set(streams) + dependent = [c for c in cond.args if fvsof(c) & stream_ops] + independent = [c for c in cond.args if not (fvsof(c) & stream_ops)] + if not independent or not any( + self._stream_equality(c, streams) for c in dependent + ): + return fwd() + + return ite( + And.plus(*independent), + monoid.reduce(ite(And.plus(*dependent), when_true, when_false), streams), + monoid.reduce(when_false, streams), + ) + + +complement: _ExtensiblePartialInvolution[Operation] = _ExtensiblePartialInvolution( + (_NumberTerm.__ne__, _NumberTerm.__eq__), +) + + +class ReduceWhereToMasks(ObjectInterpretation): + """Split an equality-guarded ``where`` reduction into masked reductions. + + For a conjunction of stream-dependent equalities ``eqs``:: + + M.reduce(where(eqs, a, b), S) + == M.plus(M.reduce(M.mask(a, eqs), S), + M.reduce(M.mask(b, not(eqs)), S)) + + De Morgan's law represents ``not(eqs)`` as the disjunction of the + corresponding disequalities. The two masks are complementary and hence + partition the stream assignments without double counting. + """ + + @staticmethod + def _combine(op, terms): + return terms[0] if len(terms) == 1 else op(*terms) + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + if not (isinstance(body, Term) and body.op == ite): + return fwd() + + cond, when_true, when_false = body.args + conds = _conjuncts(cond) + all_have_compl = all( + isinstance(t, Term) and complement.of(t.op) is not None for t in conds + ) + stream_ops = set(streams) + all_stream_dep = all(bool(fvsof(t) & stream_ops) for t in conds) + if not (conds and all_have_compl and all_stream_dep): + return fwd() + + disequalities = tuple(complement.of(t.op)(*t.args, **t.kwargs) for t in conds) + return monoid.plus( + monoid.reduce(monoid.mask(when_true, cond), streams), + monoid.reduce( + monoid.mask(when_false, self._combine(Or.plus, disequalities)), + streams, + ), + ) + + +class ReduceDisjunctiveDisequalityMask(ObjectInterpretation): + """Partition a disjunctive disequality mask into disjoint reductions. + + For example, the overlapping regions ``i != x`` and ``j != y`` are + rewritten as ``i != x`` and ``i == x and j != y``:: + + reduce(mask(v, (i != x) or (j != y)), streams) + == plus( + reduce(mask(v, i != x), streams), + reduce(mask(v, (i == x) and (j != y)), streams), + ) + + More generally, disjunct ``k`` is conjoined with the complements of all + preceding disjuncts. The resulting masks are pairwise disjoint, so their + reductions can be combined with the reduction monoid. In particular, + each result has a conjunctive mask that :class:`ReduceArrayScan` can + eliminate. + """ + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams: Streams): + match body: + case Term(mask_op, (value, Term(Or.plus, disjuncts, {})), {}) if ( + _is_monoid_mask(mask_op) and mask_op.__self__ == monoid + ): + pass + case _: + return fwd() + + if len(disjuncts) < 2 or not all( + isinstance(disjunct, Term) and is_equality(complement.of(disjunct.op)) + for disjunct in disjuncts + ): + return fwd() + + preceding_complements: list = [] + reductions = [] + for disjunct in disjuncts: + assert isinstance(disjunct, Term) + reductions.append( + monoid.reduce( + monoid.mask(value, And.plus(*preceding_complements, disjunct)), + streams, + ) + ) + comp = complement.of(disjunct.op) + assert comp is not None + preceding_complements.append(comp(*disjunct.args, **disjunct.kwargs)) + + return monoid.plus(*reductions) + + +class ReduceDependentRangeMask(ObjectInterpretation): + """Eliminate a dependent range by masking. + + reduce(M, streams ∪ {u: range(N), v: range(u())}, body) + ═══════════════════════════════════════════════════════════════════════════ + reduce(M, streams ∪ {u: range(N), v: range(N)}, where(v() < u(), body, M.identity)) + + Currently recognises only the lower-triangular form ``v: range(u())``: + constant start of 0, dependent stop equal to a bare call of another + stream var. + + Not yet supported: + + - **Upper-triangular** (``v: range(u(), N)`` — constant stop, dependent + start): bbox becomes ``range(0, N)`` (or ``range(0, bbox_N)``), guard + becomes ``v() >= u()``. Same shape of rewrite as lower-tri; differs + only in which side of the range carries the stream-var reference and + in the predicate direction. + - **Banded** (``v: range(u() - k, u() + k + 1)`` — two-sided dependent + bounds with constant width): bbox is ``range(0, N + k)`` (or similar + bounded by both endpoints' extents), guard is + ``(v() >= u() - k) & (v() < u() + k + 1)``. Needs both-sides + affine-bound recognition. + - **Strided dependent** (``v: range(0, u(), k)`` for ``k != 1``): bbox + stays ``range(0, N)`` and guard becomes + ``(v() < u()) & (v() % k == 0)`` (or equivalent), or alternatively + embed in a smaller bbox ``range(0, ceil(N/k))`` and remap the index. + - **Affine bounds** (``v: range(a*u() + b, c*u() + d)`` for affine + coefficients): bbox computed from ``ub(c*u() + d)`` over ``u``'s + range; guard is the conjunction of the two affine constraints. This + subsumes the upper/banded/strided cases under one affine recogniser. + - **Multi-stream-var dependent** (``v: range(u() + w())`` referencing + more than one outer stream var): bbox is the affine combination over + both referents' ranges; guard threads through all dependencies. + - **Reverse-order dependent ranges**: e.g. ``v: range(u(), 0, -1)``; + needs to handle negative step and the corresponding reverse + enumeration. + """ + + @implements(Monoid.reduce) + def _(self, monoid: Monoid, body, streams: Streams): + for u, u_stream in streams.items(): + # streams of the form k: range(X) + if not _is_simple_range(u_stream): + continue + + for v, v_stream in streams.items(): + if not (isinstance(v_stream, Term) and v_stream.op == range_): + continue + + start, stop, step = v_stream.start, v_stream.stop, v_stream.step # type: ignore[attr-defined] + if not ( + isinstance(start, int) + and start == 0 + and isinstance(step, int) + and step == 1 + and isinstance(stop, Term) + and stop.op == u + ): + continue + + fresh_streams = { + a: (u_stream if a == v else b) for (a, b) in streams.items() + } + fresh_body = monoid.mask(body, v() < u()) + return monoid.reduce(fresh_body, fresh_streams) + + return fwd() + + +class ReduceDisequalityMask(ObjectInterpretation): + """M.reduce(M.mask(v, And.plus(a != b, c)), S) + ≡ M.reduce(M.plus(M.mask(v, And.plus(a < b, c)), M.mask(v, And.plus(a > b, c)))) + """ + + @implements(Monoid.reduce) + def _(self, monoid, body, streams: Streams): + match body: + case Term(mask_op, (value, mask), {}) if ( + _is_monoid_mask(mask_op) and mask_op.__self__ == monoid + ): + pass + + case _: + return fwd() + + mask_elems = _conjuncts(mask) + + def _neq_to_plus(args, tail_mask_elems): + match args: + case (Term(stream_op, (), {}), index) if _is_simple_range( + streams.get(stream_op, None) + ): + pass + case _: + return None + return monoid.reduce( + monoid.plus( + monoid.mask(value, And.plus(stream_op() < index, *tail_mask_elems)), + monoid.mask(value, And.plus(stream_op() > index, *tail_mask_elems)), + ), + streams, + ) + + for i, elem in enumerate(mask_elems): + if not (isinstance(elem, Term) and is_equality(complement.of(elem.op))): + continue + + tail_mask_elems = [e for (j, e) in enumerate(mask_elems) if i != j] + ret = _neq_to_plus(elem.args, tail_mask_elems) or _neq_to_plus( + tuple(reversed(elem.args)), tail_mask_elems + ) + if ret is not None: + return ret + + return fwd() + + +class ContractLongestStream(ObjectInterpretation): + @implements(choose_contraction) + def _(self, factors, streams): + lengths = { + k: len(v) if not isinstance(v, Term) and isinstance(v, Sized) else 0 + for (k, v) in streams.items() + } + longest = max(lengths.values()) + longest_streams = {k: v for (k, v) in streams.items() if lengths[k] == longest} + if len(longest_streams) == len(streams): + return fwd() + return choose_contraction(factors, longest_streams) + + +class _ExtensibleInterpretation(UserDict, Interpretation): + def extend(self, *intps: Interpretation) -> typing.Self: + for intp in intps: + self.data = coproduct(self.data, intp) # type: ignore[assignment] + return self + + +EvaluateIntp = _ExtensibleInterpretation().extend( + ReducePartial(), + DeltaConcrete(), + SumPlus(), + MinPlus(), + MaxPlus(), + ProductPlus(), + ArgMinPlus(), + ArgMaxPlus(), + CartesianProductPlus(), + UnionPlus(), + ReduceEqualityMaskRange(), + ReduceWhereToMasks(), +) + +CartesianProductNormalizeIntp = functools.reduce( + coproduct, + typing.cast( + tuple[Interpretation, ...], + ( + PlusEmpty(), + PlusSingle(), + PlusAssoc(), + ReduceUnfactor(), + ), + ), +) +"""Structural preprocessing used exclusively for cartesian-product inversion. + +It intentionally excludes ``Factor`` and sum-of-products rewrites, which +would either cycle with ``ReduceUnfactor`` or move factors across a product +fold. +""" + + +NormalizeIntp = _ExtensibleInterpretation().extend( + GetitemDelta(), + MonoidOverSequence(), + MonoidOverMapping(), + MonoidOverCallable(), + ReduceEmpty(), + ReduceFusion(), + ReduceUnion(), + ReduceSplit(), + Factor(), + ReduceGroundCartesianProduct(), + ReduceDistributeCartesianProduct(), + ReduceWeightedStream(), + ReduceMaskHoist(), + EliminateSingletonStreams(), + PlusEmpty(), + PlusSingle(), + PlusAssoc(), + PlusDistr(), + PlusConsecutiveDups(), + PlusOrder(), + PlusCastFloat(), + MaskFusion(), + MaskBool(), + WhereHoist(), + ReduceWhereEqualityPeel(), + ReduceDisjunctiveDisequalityMask(), + ReduceDependentRangeMask(), + ReduceDisequalityMask(), + ContractLongestStream(), +) +"""``NormalizeIntp`` applies pure-Term rewrites (associativity, distributivity, +identity elimination, fusion, factorization, etc.) that drive a reduce +expression toward a *normal form*. + +Normal form +=========== + +The rules collectively push expressions toward an einsum / variable-elimination +shape: a factored sum-of-products over independent index ranges, with weights +carried symbolically by ``mask`` and ``delta`` rather than lowered to concrete +array ops. Throughout, "sum" means a ``reduce``/``plus`` of an additive monoid +``R`` and "product" a ``plus`` of a multiplicative monoid ``M`` with +``distributes_over(M, R)``. + +A fully normalized leaf expression matches the grammar (for additive ``R`` / +multiplicative ``M``):: + + nf ::= container[nf] # lambdas/dicts/sequences outermost + | M.mask(prod, orphan_cond) # only conjuncts no factor mentions + | prod + prod ::= M.plus(factor, ...) # flattened, n-ary + | factor + factor ::= M.mask(core, cond) # conjuncts the factor mentions + | core + core ::= atom # array, getitem, delta(idx, w), ... + | R.reduce(prod, ranges) # every factor in prod uses a range var + cond ::= And.plus(cmp, ...) # comparison atoms, stream-var first + ranges ::= { var: range(0, N, 1), ... } + +The invariants, grouped by the rules that establish them: + +A. Sums are factored over products (the einsum shape). + - :class:`Factor` pulls stream-invariant product factors into outer product + factors. :class:`ReduceSplit` distributes a commutative reduction over a + same-monoid ``plus``; each resulting reduction initially retains the full + stream mapping, after which factorization and stream-elimination rules may + simplify streams unused by an individual summand. + - Explicit (materialized) products-of-sums are expanded to sums-of-products + by :class:`PlusDistr`. (This is the opposite direction from factorization, + but it acts on explicit ``plus`` terms, never on ``reduce`` nodes, so the + two never fight.) + +B. Streams are independent ranges ``range(0, N, 1)``. + Every non-range binding has a rule that removes it: eager arrays + (:class:`~effectful.handlers.jax.monoid.ReduceArrayGather` + + :class:`EliminateSingletonStreams`), unions (:class:`ReduceUnion`), weighted + streams (:class:`ReduceWeightedStream`), cartesian products (eliminated by + inversion in :class:`ReduceDistributeCartesianProduct`), dependent ranges + (:class:`~effectful.handlers.jax.monoid.ReduceDependentRangeMask`), and + symbolic length-1 streams (:class:`EliminateSingletonStreams`). In particular + ``CartesianProduct.reduce`` is fully eliminated. + +C. ``plus``, ``mask``, and ``delta`` stay symbolic (not lowered). + - ``plus`` is only normalized structurally: flattened/associative + (:class:`PlusAssoc`), with nullary/unary collapsed + (:class:`PlusEmpty`/:class:`PlusSingle`) and int/float operands unified + (:class:`PlusCastFloat`). The scalar/array implementations live in + ``EvaluateIntp``, not here. + - ``mask`` floats to a canonical position: fused to a single layer with a + conjunctive (``And.plus``) condition (:class:`MaskFusion`), constant-bool + conditions discharged (:class:`MaskBool`), pushed *down* onto the factors + of a ``plus`` (:class:`MaskPushPlus`) -- each conjunct landing on the + factors that mention its variables, so a mask sits adjacent to a single + product factor where it can fuse with a gather (a conjunct no factor + mentions stays as a residual outer mask) -- and out of a ``reduce`` when + the condition is stream-independent (:class:`ReduceMaskHoist`). A mask + remaining inside a reduce body therefore depends on a reduced stream. + Comparison atoms are oriented stream-variable-first + (:class:`MaskOrderStreamOps`). + - ``delta`` is flattened: nested deltas merge (:class:`DeltaFusion`), + empty-index deltas collapse to their weight (:class:`DeltaEmpty`), and + subscripting a delta becomes a mask (:class:`GetitemDelta`). A normal-form + delta has a non-empty index, a single layer, and is never subscripted. + +D. Structural invariants. + - No same-monoid ``reduce`` is nested directly in a ``reduce`` body: they are + fused into one reduce over the combined stream set (:class:`ReduceFusion`), + so each reduce is maximal in its stream set per monoid. + - Containers are outermost: :class:`MonoidOverCallable`, + :class:`MonoidOverMapping`, and :class:`MonoidOverSequence` push + ``reduce``/``plus`` through lambdas, mappings, and sequences down to the + scalar/array leaves, so a monoid op never wraps a container. + +Termination and confluence +=========================== + +The rules run as a deterministic, priority-ordered, innermost normalization +(``evaluate`` rewrites args before parents, and each constructed replacement +re-enters the interpretation to a fixpoint). For a fixed *syntactic* input the +result is therefore unique. + +The normal form is **not fully canonical**, however: semantically equal but +differently-presented inputs can reach distinct (still semantically equal) +forms. Commutative ``plus`` arguments are sorted by :class:`PlusOrder`, and +consecutive syntactically equal arguments are removed for idempotent monoids by +:class:`PlusConsecutiveDups`. This provides limited AC normalization, but the +sort key uses object identity to break syntactic-hash ties, and duplicates are +preserved for non-idempotent monoids. Moreover, variable-elimination order +depends on ``streams`` insertion order via +:func:`choose_contraction`/:func:`outer_stream`, while argument order remains +significant for noncommutative monoids. Downstream code therefore relies only +on the *semantic* normal form (the result is ultimately evaluated to a concrete +array), not on full syntactic canonicity. + +Termination rests on per-family progress (each "lowering" rule strictly +consumes a resource -- a cartesian/weighted/union/array/singleton stream, a +nested reduce/mask/delta, or a liftable factor) plus redex-shape disjointness +that keeps the families from cycling (e.g. :class:`ReduceUnfactor` is kept out +of ``NormalizeIntp`` because it would form a rewrite cycle with +:class:`Factor`). There +is no single global measure and no confluence theorem: the system is +normalizing by construction, so a new rule that overlaps an existing redex shape +can introduce a loop or shift which normal form is reached. Note also that +:class:`PlusDistr` is exponential in the number of distributed sums. + +""" diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index 4fd5c481c..8937a3ae5 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -173,6 +173,7 @@ def evaluate[T]( @evaluate.register(object) @evaluate.register(str) @evaluate.register(bytes) +@evaluate.register(range) def _evaluate_object[T](expr: T, **kwargs) -> T: if dataclasses.is_dataclass(expr) and not isinstance(expr, type): return _evaluate_dataclass(expr, **kwargs) @@ -300,6 +301,13 @@ def _evaluate_list_view(expr, **kwargs): def _simple_type(tp: type) -> type: """Convert a type object into a type that can be dispatched on.""" + + def _resolve_aliases(tp: type) -> type: + tp = typing.get_origin(tp) or tp + if isinstance(tp, typing.TypeAliasType): + return _resolve_aliases(tp.__value__) + return tp + if isinstance(tp, typing.TypeVar): tp = ( tp.__bound__ @@ -317,7 +325,7 @@ def _simple_type(tp: type) -> type: tp = functools.reduce(operator.or_, (type(arg) for arg in args)) if isinstance(tp, types.UnionType): raise TypeError(f"Union types are not supported: {tp}") - return typing.get_origin(tp) or tp + return _resolve_aliases(tp) class _TypeofIntp(ObjectInterpretation): diff --git a/effectful/ops/syntax.py b/effectful/ops/syntax.py index d73546028..b5ef5bbba 100644 --- a/effectful/ops/syntax.py +++ b/effectful/ops/syntax.py @@ -5,7 +5,7 @@ import numbers import operator import typing -from collections.abc import Callable, Iterable, Mapping +from collections.abc import Callable, Iterable, KeysView, Mapping, ValuesView from typing import Annotated, Any from effectful.ops.types import ( @@ -606,11 +606,6 @@ def __init__( self._args = args self._kwargs = kwargs - def __eq__(self, other) -> bool: - from effectful.ops.syntax import syntactic_eq - - return syntactic_eq(self, other) - @property def op(self): return self._op @@ -790,6 +785,111 @@ def __next__(self: collections.abc.Iterator[T]) -> T: next_ = _IteratorTerm.__next__ +@defdata.register(collections.abc.Collection) +class _CollectionTerm[T](_IterableTerm[T], collections.abc.Collection[T]): + @defop + def __contains__(self: collections.abc.Collection[T], x: T) -> bool: + if not isinstance(self, Term) and not isinstance(x, Term): + return x in self + else: + raise NotHandled + + @defop + def __len__(self: collections.abc.Collection[T]) -> int: + if not isinstance(self, Term): + return len(self) + else: + raise NotHandled + + +@defdata.register(collections.abc.Sequence) +class _SequenceTerm[T](_CollectionTerm[T], collections.abc.Sequence[T]): + @Operation.define + def __getitem__(self: collections.abc.Sequence[T], key: int) -> T: + if not isinstance(self, Term) and not isinstance(key, Term): + return self[key] + else: + raise NotHandled + + @Operation.define + def __reversed__(self: collections.abc.Sequence[T]) -> collections.abc.Iterator[T]: + if not isinstance(self, Term): + return reversed(self) + else: + raise NotHandled + + @Operation.define + def index(self: collections.abc.Sequence[T], *args, **kwargs) -> int: + from effectful.ops.semantics import fvsof + + if not fvsof((self, *args, *kwargs.values())): + return self.index(*args, **kwargs) + else: + raise NotHandled + + @Operation.define + def count(self: collections.abc.Sequence[T], value: T) -> int: + from effectful.ops.semantics import fvsof + + if not fvsof((self, value)): + return self.count(value) + else: + raise NotHandled + + +@defdata.register(collections.abc.Mapping) +class _MappingTerm[K, V](_CollectionTerm[K]): + @defop + def __getitem__(self: collections.abc.Mapping[K, V], key: K) -> V: + from effectful.ops.semantics import fvsof + + if not isinstance(self, Term) and not fvsof(key): + return self[key] + else: + raise NotHandled + + @defop + def get( + self: collections.abc.Mapping[K, V], key: K, default: V | None = None + ) -> V | None: + from effectful.ops.semantics import fvsof + + if not isinstance(self, Term) and not fvsof(key): + return self.get(key, default) + else: + raise NotHandled + + @defop + def keys(self: collections.abc.Mapping[K, V]) -> KeysView[K]: + if not isinstance(self, Term): + return self.keys() + else: + raise NotHandled + + @defop + def values(self: collections.abc.Mapping[K, V]) -> ValuesView[V]: + if not isinstance(self, Term): + return self.values() + else: + raise NotHandled + + @defop + def __eq__(self: collections.abc.Mapping[K, V], other) -> bool: + if not isinstance(self, Term) and not isinstance(other, Term): + return self == other + else: + raise NotHandled + + +@Operation.define +def as_dict[K, V](*args: tuple[K, V]) -> Mapping[K, V]: + from effectful.ops.semantics import fvsof + + if not fvsof(args): + return dict(args) + raise NotHandled + + @_CustomSingleDispatchCallable def syntactic_eq( __dispatch: Callable[[type], Callable[[Any, Any], bool]], x, other @@ -891,7 +991,8 @@ def _(x: collections.abc.Sequence, other) -> bool: ) else: return ( - isinstance(other, collections.abc.Sequence) + not isinstance(other, Term) + and isinstance(other, collections.abc.Sequence) and len(x) == len(other) and all(syntactic_eq(a, b) for a, b in zip(x, other)) ) @@ -900,9 +1001,97 @@ def _(x: collections.abc.Sequence, other) -> bool: @syntactic_eq.register(object) @syntactic_eq.register(str | bytes) def _(x: object, other) -> bool: + if isinstance(other, Term): # Terms often override __eq__ + return False return x == other +@syntactic_eq.register(int | float) +def _(x: int | float, other) -> bool: + # Terms often override __eq__ + if isinstance(other, Term) or not isinstance(other, int | float): + return False + return x == other + + +@_CustomSingleDispatchCallable +def syntactic_hash(__dispatch: Callable[[type], Callable[[Any], int]], x) -> int: + """Structural hash compatible with :func:`syntactic_eq`. + + Guarantees that ``syntactic_eq(x, y)`` implies + ``syntactic_hash(x) == syntactic_hash(y)``. + + :param x: A term. + :returns: An integer hash. + """ + if dataclasses.is_dataclass(x) and not isinstance(x, type): + return hash( + ( + "dataclass", + type(x), + syntactic_hash( + { + field.name: getattr(x, field.name) + for field in dataclasses.fields(x) + } + ), + ) + ) + else: + return __dispatch(type(x))(x) + + +@syntactic_hash.register +def _(x: Term) -> int: + return hash( + ( + "term", + x.op, + len(x.args), + tuple(syntactic_hash(a) for a in x.args), + # sort kwargs so order doesn't affect the hash + tuple((k, syntactic_hash(x.kwargs[k])) for k in sorted(x.kwargs)), + ) + ) + + +@syntactic_hash.register +def _(x: collections.abc.Mapping) -> int: + # XOR over (key_hash, value_hash) pairs — order-independent, + # matching the set-based comparison in syntactic_eq's Mapping branch. + acc = 0 + for k in x: + acc ^= hash((hash(k), syntactic_hash(x[k]))) + return hash(("mapping", acc)) + + +@syntactic_hash.register +def _(x: collections.abc.Sequence) -> int: + if ( + isinstance(x, tuple) + and hasattr(x, "_fields") + and all(hasattr(x, f) for f in x._fields) + ): + return hash( + ( + "namedtuple", + type(x), + tuple(syntactic_hash(getattr(x, f)) for f in x._fields), + ) + ) + else: + # Use the abstract Sequence tag (not type(x)) because syntactic_eq + # treats any two Sequences of equal length and elementwise-equal + # contents as equal — e.g. [1,2] and (1,2) compare equal. + return hash(("sequence", len(x), tuple(syntactic_hash(a) for a in x))) + + +@syntactic_hash.register(object) +@syntactic_hash.register(str | bytes) +def _(x: object) -> int: + return hash(x) + + class ObjectInterpretation[T, V](collections.abc.Mapping): """A helper superclass for defining an ``Interpretation`` of many :class:`~effectful.ops.types.Operation` instances with shared state or behavior. @@ -1158,36 +1347,61 @@ def __index__(self) -> int: def __eq__(self, other) -> bool: if not isinstance(self, Term) and not isinstance(other, Term): return self.__eq__(other) - else: - return syntactic_eq(self, other) + + if not isinstance(other, numbers.Number): + return NotImplemented + + raise NotHandled + + @defop + def __ne__(self, other) -> bool: + if not isinstance(self, Term) and not isinstance(other, Term): + return self.__ne__(other) + + if not isinstance(other, numbers.Number): + return NotImplemented + + raise NotHandled @defop def __lt__(self, other) -> bool: if not isinstance(self, Term) and not isinstance(other, Term): return self.__lt__(other) - else: - raise NotHandled + + if not isinstance(other, numbers.Number): + return NotImplemented + + raise NotHandled @defop def __gt__(self, other) -> bool: if not isinstance(self, Term) and not isinstance(other, Term): return self.__gt__(other) - else: - raise NotHandled + + if not isinstance(other, numbers.Number): + return NotImplemented + + raise NotHandled @defop def __le__(self, other) -> bool: if not isinstance(self, Term) and not isinstance(other, Term): return self.__le__(other) - else: - raise NotHandled + + if not isinstance(other, numbers.Number): + return NotImplemented + + raise NotHandled @defop def __ge__(self, other) -> bool: - if not isinstance(self, Term) and not isinstance(other, Term): + if not isinstance(self, _NumberTerm) and not isinstance(other, _NumberTerm): return self.__ge__(other) - else: - raise NotHandled + + if not isinstance(other, numbers.Number): + return NotImplemented + + raise NotHandled @defop def __add__(self, other: T) -> T: @@ -1399,6 +1613,53 @@ class _BoolTerm[T: bool](_IntegralTerm[T]): # type: ignore pass +@Operation.define +def ite[T](cond, then: T, else_: T) -> T: + """If-then-else operation.""" + # Note: cond is specifically not annotated to allow this operation to take + # non-bool arguments (like boolean arrays). + if not isinstance(cond, Term): + return then if cond else else_ + raise NotHandled + + +@Operation.define +def range_(stop: int, *args: int) -> range: + if any(isinstance(x, Term) for x in (stop, *args)): + raise NotHandled + return range(stop, *args) + + +@defdata.register(range) +class _RangeTerm(_SequenceTerm[int]): + @property + @Operation.define + def start(self) -> int: + if not isinstance(self, Term): + return self.start + if self.op == range_: + return 0 if len(self.args) < 2 else self.args[0] + raise NotHandled + + @property + @Operation.define + def stop(self) -> int: + if not isinstance(self, Term): + return self.stop + if self.op == range_: + return self.args[0] if len(self.args) < 2 else self.args[1] + raise NotHandled + + @property + @Operation.define + def step(self) -> int: + if not isinstance(self, Term): + return self.step + if self.op == range_: + return 1 if len(self.args) < 3 else self.args[2] + raise NotHandled + + class ConstructorOperation[**Q, V](Operation[Q, V]): @classmethod @functools.cache diff --git a/effectful/ops/types.py b/effectful/ops/types.py index 96cda80ed..6f119e2fa 100644 --- a/effectful/ops/types.py +++ b/effectful/ops/types.py @@ -279,11 +279,17 @@ def _define_callable[**P, T]( cls, t: Callable[P, T], *, name: str | None = None ) -> "Operation[P, T]": if isinstance(t, Operation): + sig = inspect.signature(t) @functools.wraps(t) def func(*args, **kwargs): raise NotHandled + # functools.wraps does not copy the signature. Instead it points to + # the wrapped function. inspect.signature will traverse this chain + # unless it is removed. + del func.__wrapped__ + func.__signature__ = sig # type: ignore[attr-defined] op = cls.define(func, name=name) else: op = cls(t, name=name) # type: ignore[arg-type] @@ -319,6 +325,15 @@ def func(*args, **kwargs): return typing.cast(Operation[P, T], cls.define(func, **kwargs)) + @define.register(types.MethodType) + @classmethod + def _define_methodtype[**P, T]( + cls, t: Callable[P, T], *, name: str | None = None + ) -> "Operation[P, T]": + op = cls._define_callable(t, name=name) + op.__self__ = t.__self__ # type: ignore[attr-defined] + return typing.cast("Operation[P, T]", op) + @define.register(staticmethod) @classmethod def _define_staticmethod[**P, T](cls, t: "staticmethod[P, T]", **kwargs): @@ -496,7 +511,17 @@ def _instance_op(instance, *args, **kwargs): else: return default_result - instance_op = self.define(types.MethodType(_instance_op, instance)) + name: str = "" + if instance is not None and hasattr(instance, "__name__"): + assert isinstance(instance.__name__, str) + name = instance.__name__ + elif owner is not None: + name = owner.__name__ + name += f"_{self.__name__}" + + instance_op = self.define( + types.MethodType(_instance_op, instance), name=name + ) instance.__dict__[self._name_on_instance] = instance_op return instance_op elif instance is not None: diff --git a/pyproject.toml b/pyproject.toml index 0ae2f45b6..f1d4a672d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,22 +39,23 @@ Source = "https://github.com/BasisResearch/effectful" [project.optional-dependencies] torch = ["torch"] pyro = ["pyro-ppl>=1.9.1"] -jax = ["jax"] +jax = [ + "jax", + "opt_einsum" +] numpyro = [ "numpyro>=0.19", "jax<0.10" ] llm = [ - # 1.92.0 ships no wheel, so it builds from sdist; on Python 3.14 that build pulls - # pyo3-ffi 0.23.5, which doesn't support 3.14, so `uv sync` fails the whole matrix. - "litellm!=1.92.0", + "litellm<1.92", "tenacity", "mypy", "autoflake", "pillow", "pydantic", "typing_extensions", - "restrictedpython>=8.1" + "restrictedpython>=8.1", ] prettyprinter = ["prettyprinter"] docs = [ @@ -73,6 +74,7 @@ test = [ "pytest-cov", "pytest-xdist", "pytest-benchmark", + "hypothesis", "mypy", "ruff", "nbval", diff --git a/tests/_monoid_helpers.py b/tests/_monoid_helpers.py new file mode 100644 index 000000000..88b3b436d --- /dev/null +++ b/tests/_monoid_helpers.py @@ -0,0 +1,427 @@ +import builtins +import itertools +import typing +from abc import ABC, abstractmethod +from collections.abc import Callable, Mapping +from typing import Any, Literal, overload + +import jax +from hypothesis import given, settings +from hypothesis import strategies as st +from hypothesis.strategies import SearchStrategy + +import effectful.handlers.jax.numpy as _jnp +from effectful.internals.runtime import interpreter +from effectful.ops.monoid import ( + EvaluateIntp, + NormalizeIntp, + Stream, + _is_monoid_weighted, +) +from effectful.ops.semantics import apply, coproduct, evaluate, fvsof, handler +from effectful.ops.syntax import _BaseTerm, defdata, deffn, syntactic_eq +from effectful.ops.types import NotHandled, Operation, Term + + +def syntactic_eq_alpha(x, y) -> bool: + """Alpha-equivalence-respecting variant of ``syntactic_eq``. + + Walks each expression bottom-up with :func:`evaluate` and renames + every bound variable to a deterministic canonical Operation. The + canonical names are assigned by a counter that increments in + ``evaluate``'s natural traversal order, so two alpha-equivalent + expressions canonicalize to syntactically identical results. + """ + + _op_cache: dict[int, Operation] = {} + + def _canonical_op(idx: int, op: Operation) -> Operation: + """Cached canonical Operation, keyed by encounter index. + + Cached so that two independent canonicalize runs return the same + Operation object for the same index — letting ``syntactic_eq`` + compare canonical forms by Operation identity. + """ + if idx in _op_cache: + return _op_cache[idx] + + canon_op = Operation.define(op, name=f"__cv_{idx}") + _op_cache[idx] = canon_op + return canon_op + + cx = _canonicalize(x, _canonical_op) + cy = _canonicalize(y, _canonical_op) + return syntactic_eq(cx, cy) + + +def _canonicalize(expr, _canonical_op): + counter = itertools.count() + + def _substitute(arg, renaming): + """Apply a bound-variable renaming using ``evaluate`` for traversal.""" + if not renaming: + return arg + with interpreter({apply: _BaseTerm, **renaming}): + return evaluate(arg) + + def _bound_var_order(args, kwargs, bound_set: set[Operation]) -> list[Operation]: + """Return bound variables in deterministic encounter order.""" + seen: list[Operation] = [] + seen_set: set[Operation] = set() + + def _capture(op, *a, **kw): + if op in bound_set and op not in seen_set: + seen.append(op) + seen_set.add(op) + return defdata(op, *a, **kw) + + # ``evaluate`` walks Terms, lists, tuples, mappings, dataclasses, + # etc. for free; the apply handler captures bound vars used as + # ``x()`` anywhere in the body. + with interpreter({apply: _capture}): + evaluate((args, kwargs)) + + # Binders bypass the apply handler. Pick them up with a small structural + # walk that visits dict keys too. + def _walk_bare(obj): + if isinstance(obj, Operation): + if obj in bound_set and obj not in seen_set: + seen.append(obj) + seen_set.add(obj) + elif isinstance(obj, dict): + for k, v in obj.items(): + _walk_bare(k) + _walk_bare(v) + elif isinstance(obj, list | set | frozenset | tuple): + for v in obj: + _walk_bare(v) + + _walk_bare((args, kwargs)) + return seen + + def _apply_canonical(op, *args, **kwargs) -> Term: + bindings = op.__fvs_rule__(*args, **kwargs) + all_bound: set[Operation] = set().union( + *bindings.args, *bindings.kwargs.values() + ) + if not all_bound: + return _BaseTerm(op, *args, **kwargs) + + order = _bound_var_order(args, kwargs, all_bound) + canonical = {var: _canonical_op(next(counter), var) for var in order} + assert all_bound <= set(order) + + new_args = tuple( + _substitute( + arg, {v: canonical[v] for v in bindings.args[i] if v in canonical} + ) + for i, arg in enumerate(args) + ) + new_kwargs = { + k: _substitute( + v, + {var: canonical[var] for var in bindings.kwargs[k] if var in canonical}, + ) + for k, v in kwargs.items() + } + + # avoid the renaming from defdata + return _BaseTerm(op, *new_args, **new_kwargs) + + with interpreter({apply: _apply_canonical}): + return evaluate(expr) + + +class Backend(ABC): + """A value-domain spec used to share monoid tests across int and jax.Array + backends. Provides the concrete value type, the hypothesis strategy for + drawing scalars in property tests, and an equality predicate that works + for that domain. + """ + + name: str + scalar_typ: Any + stream_typ: Any + strategy_for_op: dict[Operation, st.SearchStrategy[Callable[..., Any]]] + + def __init__(self): + self.strategy_for_op = {} + + @abstractmethod + def eq(self, a: Any, b: Any) -> bool: + raise NotImplementedError + + @abstractmethod + def strategy( + self, + arg_types: tuple[type, ...] = (), + ret: Literal["scalar", "stream"] = "scalar", + ) -> SearchStrategy: + raise NotImplementedError + + def _fresh_op( + self, + name: str, + arg_types: tuple[type, ...] = (), + ret: Literal["scalar", "stream"] = "scalar", + ) -> Operation: + """Build a fresh, unhandled Operation whose parameter and return + annotations are derived from this backend. + + ``ret`` is ``"scalar"`` for a scalar return or ``"stream"`` for a + stream-of-scalar return. The operation has ``n_args`` parameters, + each of type ``scalar_typ``. + """ + scalar = self.scalar_typ + out = self.stream_typ if ret == "stream" else scalar + params = ", ".join(f"_a{i}" for i in range(len(arg_types))) + ns: dict[str, Any] = {"NotHandled": NotHandled} + exec(f"def _fn({params}):\n raise NotHandled\n", ns) + fn = ns["_fn"] + fn.__annotations__ = { + **{f"_a{i}": t for i, t in enumerate(arg_types)}, + "return": out, + } + op = Operation.define(fn, name=name) + self.strategy_for_op[op] = self.strategy(arg_types, ret) + return op + + @overload + def define_vars(self, name: str, /, **kwargs) -> Operation: ... + + @overload + def define_vars( + self, n1: str, n2: str, /, *names: str, **kwargs + ) -> tuple[Operation, ...]: ... + + def define_vars(self, *names: str, **kwargs) -> Operation | tuple[Operation, ...]: # type: ignore[misc] + if len(names) == 1: + return self._fresh_op(names[0], **kwargs) + return tuple(self._fresh_op(n, **kwargs) for n in names) + + def check_rewrite( + self, lhs, rhs, rule, *, max_examples: int = 25, deadline=None, normalize=None + ) -> None: + normalize = ( + normalize + if normalize is not None + else coproduct(EvaluateIntp, NormalizeIntp) + ) + + with handler(rule): + norm = evaluate(lhs) + assert syntactic_eq_alpha(norm, rhs) + + fvs = fvsof(lhs) | fvsof(rhs) + + @st.composite + def random_interpretation( + draw: st.DrawFn, + ) -> Mapping[Operation, Callable[..., Any]]: + """Draw an Interpretation binding every Operation in `free_vars` to + a randomly chosen value/callable. Keys are Operation identities. + """ + intp: dict[Operation, Callable[..., Any]] = {} + for op, strategy in self.strategy_for_op.items(): + if op in fvs: + intp[op] = draw(strategy) + return intp + + @given(intp=random_interpretation()) + @settings( + max_examples=max_examples, deadline=deadline, report_multiple_bugs=False + ) + def _check_semantics(intp): + with handler(normalize), handler(intp): + lhs_val = evaluate(lhs) + rhs_val = evaluate(rhs) + assert self.eq(lhs_val, rhs_val) + + _check_semantics() + + +def _is_weighted(x: Any) -> bool: + return isinstance(x, Term) and _is_monoid_weighted(x.op) + + +def _weight_pairs(x: Any, monoid: Any) -> list[tuple[Any, Any]] | None: + """Return ``(element, weight)`` pairs for a stream. + + A weighted-monoid Term yields each element paired with its weight. A plain + (unweighted) stream yields each element paired with ``monoid.identity`` -- + the no-op weight -- so an unweighted stream compares equal to a weighted one + exactly when every weight reduces to the identity (e.g. ``[()]`` vs a + weighted ``[()]`` whose single empty row reduces to the identity, and, more + generally, whenever both streams are empty). Returns ``None`` for a + non-stream Term, which never compares equal to a weighted stream. + """ + if isinstance(x, Term): + if not _is_monoid_weighted(x.op): + return None + stream, weight = x.args + assert not isinstance(stream, Term) + return [(e, typing.cast(Callable, weight)(e)) for e in stream] + return [(e, monoid.identity) for e in x] + + +def _weighted_stream_eq(a, b, leaf_eq: Callable[[Any, Any], bool]) -> bool: + monoids = {x.op.__self__ for x in (a, b) if _is_weighted(x)} + # distinct weight monoids can never be equal + if len(monoids) != 1: + return False + monoid = next(iter(monoids)) + + a_pairs = _weight_pairs(a, monoid) + b_pairs = _weight_pairs(b, monoid) + if a_pairs is None or b_pairs is None or len(a_pairs) != len(b_pairs): + return False + for (ea, wa), (eb, wb) in zip(a_pairs, b_pairs): + if not leaf_eq(ea, eb) or not leaf_eq(wa, wb): + return False + return True + + +class IntBackend(Backend): + name = "int" + scalar_typ = int + stream_typ = Stream[int] + + _unary_num_fns: list[Callable[[int], int]] = [ + lambda x: x, + lambda x: x + 1, + lambda x: x - 1, + lambda x: -x, + lambda x: 2 * x, + lambda x: 3 * x + 1, + ] + + _binary_num_fns: list[Callable[[int, int], int]] = [ + lambda x, y: x + y, + lambda x, y: x - y, + lambda x, y: x * y, + lambda x, y: x + 2 * y, + lambda x, y: 2 * x - y, + ] + + _unary_list_fns: list[Callable[[int], list[int]]] = [ + lambda _x: [], + lambda x: [x], + lambda x: [x, x + 1], + lambda x: [x, -x], + lambda x: [0, x, x + 1], + ] + + def strategy( + self, + arg_types: tuple[type, ...] = (), + ret: Literal["scalar", "stream"] = "scalar", + ) -> SearchStrategy: + match arg_types, ret: + case (), "scalar": + return st.integers(min_value=-100, max_value=100).map(deffn) + case (), "stream": + scalars = st.integers(min_value=-100, max_value=100) + return st.lists(scalars, max_size=2).map(deffn) + case (builtins.int,), "scalar": + return st.sampled_from(self._unary_num_fns) + case (builtins.int, builtins.int), "scalar": + return st.sampled_from(self._binary_num_fns) + case (builtins.int, builtins.int, builtins.int), "scalar": + return st.tuples( + st.sampled_from(self._binary_num_fns), + st.sampled_from(self._binary_num_fns), + ).map(lambda fg: lambda a, b, c: fg[0](a, fg[1](b, c))) + case (builtins.int,), "stream": + return st.sampled_from(self._unary_list_fns) + raise NotImplementedError( + f"No int strategy for op with return {ret!r} and {arg_types} args" + ) + + def eq(self, a: Any, b: Any) -> bool: + if _is_weighted(a) or _is_weighted(b): + return _weighted_stream_eq(a, b, self.eq) + return not isinstance(a, Term) and not isinstance(b, Term) and a == b + + +class JaxBackend(Backend): + name = "jax" + scalar_typ = jax.Array + stream_typ = jax.Array + + _unary_jax_scalar_fns: list[Callable[[jax.Array], jax.Array]] = [ + lambda a: a, + lambda a: a + 1, + lambda a: a - 1, + lambda a: -a, + lambda a: 2 * a, + ] + + _unary_jax_stream_fns: list[Callable[[jax.Array], Stream[jax.Array]]] = [ + lambda a: _jnp.stack([a, a + 1]), + lambda a: _jnp.stack([a, -a]), + lambda a: _jnp.stack([a, a + 1, 2 * a]), + ] + + _binary_jax_scalar_fns: list[Callable[[jax.Array, jax.Array], jax.Array]] = [ + lambda a, b: a + b, + lambda a, b: a - b, + lambda a, b: a * b, + ] + + def strategy( + self, + arg_types: tuple[type, ...] = (), + ret: Literal["scalar", "stream"] = "scalar", + ) -> st.SearchStrategy[Callable]: + match arg_types, ret: + case (), "scalar": + return ( + st.lists( + st.integers(min_value=-5, max_value=5), + min_size=2, + max_size=2, + ) + .map(lambda xs: jax.numpy.asarray(xs, dtype=jax.numpy.float32)) + .map(deffn) + ) + case (), "stream": + return ( + st.lists( + st.integers(min_value=-5, max_value=5), + min_size=1, + max_size=2, + ) + .map(lambda xs: jax.numpy.asarray(xs, dtype=jax.numpy.float32)) + .map(deffn) + ) + case (jax.Array,), "scalar": + return st.sampled_from(self._unary_jax_scalar_fns) + case (jax.Array, jax.Array), "scalar": + return st.sampled_from(self._binary_jax_scalar_fns) + case (jax.Array, jax.Array, jax.Array), "scalar": + return st.tuples( + st.sampled_from(self._binary_jax_scalar_fns), + st.sampled_from(self._binary_jax_scalar_fns), + ).map(lambda fg: lambda a, b, c: fg[0](a, fg[1](b, c))) + case (jax.Array,), "stream": + return st.sampled_from(self._unary_jax_stream_fns) + + raise NotImplementedError( + f"No jax strategy for op with return {ret!r} and {arg_types} args" + ) + + def eq(self, a: Any, b: Any) -> bool: + if _is_weighted(a) or _is_weighted(b): + return _weighted_stream_eq(a, b, self.eq) + + def _leaf_eq(x: Any, y: Any) -> bool: + return bool(jax.numpy.all(jax.numpy.isclose(x, y, equal_nan=True))) + + try: + leaves = jax.tree.leaves(jax.tree.map(_leaf_eq, a, b)) + except (ValueError, TypeError): + return False + return all(leaves) + + +__all__ = ["Backend", "IntBackend", "JaxBackend", "syntactic_eq_alpha"] diff --git a/tests/test_handlers_jax.py b/tests/test_handlers_jax.py index e07385725..4c825542c 100644 --- a/tests/test_handlers_jax.py +++ b/tests/test_handlers_jax.py @@ -2,7 +2,7 @@ import pytest import effectful.handlers.jax.numpy as jnp -from effectful.handlers.jax import bind_dims, jax_getitem, jit, sizesof +from effectful.handlers.jax import bind_dims, jax_getitem, jit, sizesof, unbind_dims from effectful.ops.semantics import evaluate, fvsof, handler from effectful.ops.syntax import defdata, defop, syntactic_eq from effectful.ops.types import Term @@ -80,6 +80,9 @@ def test_bind_dims(): t9 = evaluate(t8) assert not (fvsof(t9) & {i, j, k, w}) + t10 = bind_dims(jax_getitem(jnp.ones((5, 6)), [unbind_dims(w(), i), i()]), i) + assert isinstance(t10, Term) and t10.op == bind_dims + def test_tpe_1(): i, j = defop(jax.Array), defop(jax.Array) diff --git a/tests/test_handlers_jax_monoid.py b/tests/test_handlers_jax_monoid.py new file mode 100644 index 000000000..a538664d8 --- /dev/null +++ b/tests/test_handlers_jax_monoid.py @@ -0,0 +1,535 @@ +import sys + +import jax +import jax.dlpack +import numpy as np +import pyro.ops.contract +import pytest +import torch +from jax import random as random + +import effectful.handlers.jax.numpy as jnp +from effectful.handlers.jax import bind_dims, jax_getitem, unbind_dims +from effectful.handlers.jax.monoid import ( + ARRAY_REDUCTORS, + ReduceArray, + ReduceArrayGather, + ReduceDeltaSimpleRange, + ReduceSumProductContraction, + einsum, +) +from effectful.ops.monoid import ( + EliminateSingletonStreams, + EvaluateIntp, + NormalizeIntp, + Product, + Sum, +) +from effectful.ops.semantics import coproduct, evaluate, handler +from tests._monoid_helpers import JaxBackend + +MONOIDS = [ + pytest.param(monoid, reductor, id=monoid.__name__) + for (monoid, reductor) in ARRAY_REDUCTORS.items() +] + + +@pytest.fixture(scope="module") +def rng_key(): + return random.PRNGKey(0) + + +@pytest.fixture +def backend() -> JaxBackend: + return JaxBackend() + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_array_gather(monoid, reductor, backend: JaxBackend): + (x, k) = backend.define_vars("x", "k", ret="scalar") + X = jnp.arange(3) + + lhs = monoid.reduce(x(), {x: X}) + rhs = monoid.reduce(unbind_dims(X, k), {k: range(X.shape[0])}) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=coproduct(ReduceArrayGather(), EliminateSingletonStreams()), + ) + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_array_gather_step1(monoid, reductor, backend: JaxBackend): + """Step 1 alone: an array stream becomes an index range plus a length-1 + stream holding the gathered element. ``ReduceArrayGather`` does not perform + the gather substitution itself -- that is ``EliminateSingletonStreams``.""" + (x, k) = backend.define_vars("x", "k", ret="scalar") + X = jnp.arange(3) + + lhs = monoid.reduce(x(), {x: X}) + rhs = monoid.reduce(x(), {k: range(X.shape[0]), x: (unbind_dims(X, k),)}) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceArrayGather()) + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_array_gather_dep(monoid, reductor, backend: JaxBackend): + (x, y) = backend.define_vars("x", "y", ret="scalar") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="stream") + g = backend.define_vars( + "g", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) + X = jnp.arange(3) + + # The dependent stream ``y: f(x())`` gets the *gathered element* X[x] + # substituted for x -- i.e. ``f(X[x])`` -- not the bare index. + lhs = monoid.reduce(g(x(), y()), {y: f(x()), x: X}) + rhs = monoid.reduce( + g(unbind_dims(X, x), y()), + {y: f(unbind_dims(X, x)), x: range(X.shape[0])}, + ) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=coproduct(ReduceArrayGather(), EliminateSingletonStreams()), + ) + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_array_1(monoid, reductor, backend: JaxBackend): + x = backend.define_vars("x", ret="scalar") + arr = jnp.arange(4, 9) + X = jnp.arange(5) + + with handler(NormalizeIntp), handler(EvaluateIntp): + actual = monoid.reduce(unbind_dims(arr, x), {x: X}) + assert jnp.allclose(actual, reductor(arr)) + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_array_2(monoid, reductor, backend: JaxBackend): + (x, y) = backend.define_vars("x", "y", ret="scalar") + X = jnp.arange(5) + Y = jnp.arange(7) + arr_X = jnp.arange(4, 9) + arr_Y = jnp.arange(4, 11) + + with handler(NormalizeIntp), handler(EvaluateIntp): + actual = monoid.reduce( + monoid.plus(unbind_dims(arr_X, x), unbind_dims(arr_Y, y)), {x: X, y: Y} + ) + expected = reductor( + bind_dims(monoid.plus(unbind_dims(arr_X, x), unbind_dims(arr_Y, y)), x, y) + ) + + assert jnp.allclose(actual, expected) + + +SCALAR_PLUS = [ + pytest.param(Sum, 3.0, id="Sum"), + pytest.param(Product, 2.0, id="Product"), +] + + +@pytest.mark.parametrize("monoid,expected", SCALAR_PLUS) +def test_plus_scalars_stays_scalar(monoid, expected): + """``Monoid.plus`` of plain Python numbers must not become a ``jax.Array``. + + ``jax.typing.ArrayLike`` is a union that includes ``bool``, ``int``, + ``float`` and ``complex``, so the jax ``plus`` handlers -- which extend + ``EvaluateIntp`` after the scalar implementations and therefore take + precedence -- must not claim pure-Python scalar arithmetic. + """ + with handler(NormalizeIntp), handler(EvaluateIntp): + actual = monoid.plus(1.0, 2.0) + + assert not isinstance(actual, jax.Array) + assert isinstance(actual, float) + assert actual == expected + + +@pytest.mark.parametrize( + "monoid,expected", + [pytest.param(Sum, 6.0, id="Sum"), pytest.param(Product, 8.0, id="Product")], +) +def test_reduce_scalar_body_stays_scalar(monoid, expected, backend: JaxBackend): + """A reduction over a scalar body is likewise plain Python arithmetic.""" + i = backend.define_vars("i", ret="scalar") + + with handler(NormalizeIntp), handler(EvaluateIntp): + actual = monoid.reduce(2.0, {i: range(3)}) + + assert not isinstance(actual, jax.Array) + assert isinstance(actual, float) + assert actual == expected + + +@pytest.mark.parametrize("monoid,expected", SCALAR_PLUS) +def test_plus_mixed_array_and_scalar_is_array(monoid, expected): + """One genuine array is enough: the narrowing must not be over-applied.""" + with handler(NormalizeIntp), handler(EvaluateIntp): + actual = monoid.plus(jnp.asarray([1.0, 1.0]), 2.0) + + assert isinstance(actual, jax.Array) + assert jnp.allclose(actual, jnp.asarray([expected, expected])) + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_arange_reduce_indirect(monoid, reductor, backend: JaxBackend): + """When the range var is used both as a direct index and as a value + (``A[v()] + v()``), the direct use slices and the indirect use materializes + the range, both aligned on the same fresh dim.""" + (v, k) = backend.define_vars("v", "k", ret="scalar") + A = jnp.arange(10) + + lhs = monoid.reduce(jax_getitem(A, [v()]) + v(), {v: range(5)}) + rhs = reductor( + bind_dims( + jax_getitem(jax_getitem(A, [slice(0, 5, 1)]), [k()]) + + unbind_dims(jnp.arange(5), k), + k, + ), + axis=(0,), + ) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceArray(), ReduceDeltaSimpleRange()) + ) + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_arange_reduce_two_streams(monoid, reductor, backend: JaxBackend): + """Two arange streams indexing a 2-D array slice both axes and reduce over + both at once.""" + (u, w, k1, k2) = backend.define_vars("u", "w", "k1", "k2", ret="scalar") + A = jnp.arange(8 * 9).reshape((8, 9)) + + lhs = monoid.reduce(jax_getitem(A, [u(), w()]), {u: range(4), w: range(5)}) + rhs = reductor( + bind_dims( + jax_getitem(jax_getitem(A, [slice(0, 4, 1), slice(0, 5, 1)]), [k1(), k2()]), + k1, + k2, + ), + axis=(0, 1), + ) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceArray(), ReduceDeltaSimpleRange()) + ) + + +# --------------------------------------------------------------------------- +# Delta rules. All tests use the operation form ``delta(idx, body)`` rather +# than the ``Delta`` dataclass; the delta op is the user-facing surface. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_delta_empty_arange(monoid, reductor, backend: JaxBackend): + x = backend.define_vars("x", ret="scalar") + f = backend.define_vars("f", arg_types=[backend.scalar_typ], ret="scalar") + + lhs = monoid.reduce(monoid.delta((x(),), f(x())), {x: range(0)}) + rhs = bind_dims(f(unbind_dims(jnp.array([]), x)), x) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDeltaSimpleRange()) + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_delta_independent_one(monoid, reductor, backend: JaxBackend): + """One R1 step: peel the final preserved index off a delta. + + reduce(M, {y: Y()}, delta((y(),), f(y()))) ≡ bind_dims(f(unbind_dims(Y(), k)), k) + """ + (y, k) = backend.define_vars("y", "k", ret="scalar") + f = backend.define_vars("f", arg_types=[backend.scalar_typ], ret="scalar") + + # We use a concrete range here instead of an abstract one, because + # unbind_dims is undefined on empty arrays (and the rewrite produces a + # different rhs in this case) + lhs = monoid.reduce(monoid.delta((y(),), f(y())), {y: range(3)}) + rhs = bind_dims(f(unbind_dims(jnp.arange(3), k)), k) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDeltaSimpleRange()) + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_delta_independent_preserves_others( + monoid, reductor, backend: JaxBackend +): + """R1 peels only the final index. Streams not matching the peeled index op + stay untouched, as do earlier entries in the index tuple. + + reduce(M, {x: X(), y: Y()}, delta((x(), y()), f(x(), y()))) + ≡ reduce(M, {x: X()}, delta((x(),), bind_dims(f(x(), unbind_dims(Y(), k)), k))) + """ + (x, y, k) = backend.define_vars("x", "y", "k", ret="scalar") + f = backend.define_vars( + "f", arg_types=[backend.scalar_typ, backend.scalar_typ], ret="scalar" + ) + + lhs = monoid.reduce( + monoid.delta((x(), y()), f(x(), y())), {x: range(2), y: range(3)} + ) + rhs = bind_dims( + bind_dims(f(unbind_dims(jnp.arange(2), x), unbind_dims(jnp.arange(3), k)), k), x + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDeltaSimpleRange()) + + +@pytest.mark.parametrize("monoid,reductor", MONOIDS) +def test_reduce_delta_simple_dep(monoid, reductor, backend: JaxBackend): + (x, y) = backend.define_vars("x", "y", ret="scalar") + X = jnp.arange(3) + + lhs = monoid.reduce( + monoid.delta((x(),), unbind_dims(X, x) + y()), + {x: range(3), y: jnp.stack([x(), x() + 1])}, + ) + rhs = bind_dims( + monoid.reduce( + unbind_dims(X, x) + y(), + { + y: jnp.stack( + [unbind_dims(jnp.arange(3), x), unbind_dims(jnp.arange(3), x) + 1] + ) + }, + ), + x, + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDeltaSimpleRange()) + + +def test_reduce_contraction_single(backend: JaxBackend): + i = backend.define_vars("i", ret="scalar") + (A, B) = backend.define_vars( + "A", "B", arg_types=(backend.scalar_typ,), ret="scalar" + ) + + lhs = Sum.reduce(Product.plus(A(i()), B(i())), {i: range(5)}) + rhs = jnp.einsum( + "a...,a...->...", + Sum.reduce(Sum.delta((i(),), A(i())), {i: range(5)}), + Sum.reduce(Sum.delta((i(),), B(i())), {i: range(5)}), + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceSumProductContraction()) + + +def test_reduce_contraction_double(backend: JaxBackend): + i, j = backend.define_vars("i", "j", ret="scalar") + (A, B) = backend.define_vars( + "A", "B", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) + + lhs = Sum.reduce(Product.plus(A(i(), j()), B(i(), j())), {i: range(5), j: range(7)}) + rhs = jnp.einsum( + "ab...,ab...->...", + Sum.reduce(Sum.delta((i(), j()), A(i(), j())), {i: range(5), j: range(7)}), + Sum.reduce(Sum.delta((i(), j()), B(i(), j())), {i: range(5), j: range(7)}), + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceSumProductContraction()) + + +def test_reduce_matmul(backend: JaxBackend): + key = jax.random.PRNGKey(0) + # Define dimensions + B, I, J, K = 2, 3, 4, 5 + + # Create sample matrices + X = random.normal(key, (B, I, J)) + Y = random.normal(key, (B, J, K)) + (b, i, j, k) = backend.define_vars("b", "i", "j", "k", ret="scalar") + + with handler(NormalizeIntp): + norm = Sum.reduce( + Sum.delta( + (b(), i(), k()), unbind_dims(X, b, i, j) * unbind_dims(Y, b, j, k) + ), + {b: range(B), i: range(I), j: range(J), k: range(K)}, + ) + with handler(EvaluateIntp), handler(NormalizeIntp): + actual = evaluate(norm) + + assert isinstance(actual, jax.Array) + expected = jnp.einsum("bij,bjk->bik", X, Y) + assert jnp.allclose(actual, expected) + + +EINSUM_CASES = [ + pytest.param("ij,jk->ik", {"i": 64, "j": 64, "k": 64}, id="matmul"), + pytest.param( + "bij,bjk->bik", + {"b": 16, "i": 32, "j": 32, "k": 32}, + id="batched_matmul", + ), + pytest.param( + "a,abi,bcij,cdij->ij", + {"a": 4, "b": 4, "c": 4, "d": 4, "i": 8, "j": 8}, + id="mixed_rank", + ), + # ───────────────────────── single-operand reshuffles ───────────────────── + # No contraction across operands — these stress the diagonal/transpose/sum + # rewrites rather than any pairwise product ordering. + pytest.param("ij->ji", {"i": 256, "j": 256}, id="transpose"), + pytest.param("ijk->", {"i": 96, "j": 96, "k": 96}, id="full_reduce"), + pytest.param("ijk->k", {"i": 96, "j": 96, "k": 96}, id="partial_reduce"), + # Repeated index *within* one operand — exercises the implicit-diagonal path + # in ReduceDeltaSimpleRange (no explicit jnp.diagonal step). + pytest.param("ii->", {"i": 1024}, id="trace"), + pytest.param("ii->i", {"i": 1024}, id="diagonal"), + pytest.param("bii->b", {"b": 256, "i": 128}, id="batched_trace"), + pytest.param("iij->ij", {"i": 128, "j": 128}, id="diagonal_keep"), + # ───────────────────────── no-shared-index blowups ─────────────────────── + # Output is the full outer product — nothing contracts, so the result tensor + # is as large as the dense intermediate. Pure broadcast cost. + pytest.param("i,j->ij", {"i": 1024, "j": 1024}, id="outer_product"), + pytest.param("ij,kl->ijkl", {"i": 32, "j": 32, "k": 32, "l": 32}, id="outer_4d"), + # Element-wise: every index shared, none contracted. + pytest.param("ij,ij->ij", {"i": 512, "j": 512}, id="hadamard"), + # ───────────────────────── ordering-sensitive products ─────────────────── + # Skewed matrix chain: contracting middle-first (b,d small) is orders of + # magnitude cheaper than the left-to-right order, which materializes a big + # a×c intermediate. The classic "matrix chain order matters" case. + pytest.param( + "ab,bc,cd->ad", {"a": 256, "b": 2, "c": 256, "d": 2}, id="skewed_chain" + ), + pytest.param( + "ab,bc,cd,de->ae", + {"a": 50, "b": 40, "c": 30, "d": 20, "e": 10}, + id="chain_4", + ), + pytest.param( + "ab,bc,cd,de,ef->af", + {"a": 12, "b": 11, "c": 10, "d": 9, "e": 8, "f": 7}, + id="chain_5", + ), + # ───────────────────────── tensor-network shapes ───────────────────────── + # Cyclic / hyperedge contractions with no tree decomposition into matmuls; + # every operand shares indices with two others. + pytest.param("ij,jk,ki->", {"i": 64, "j": 64, "k": 64}, id="trace_of_product"), + pytest.param("ij,jk,ik->", {"i": 48, "j": 48, "k": 48}, id="triangle"), + pytest.param("ijk,jl,kl->il", {"i": 24, "j": 24, "k": 24, "l": 24}, id="hyperedge"), + # Star: many operands share one contracted index, fanning into a large + # outer-product output. + pytest.param( + "ai,bi,ci,di->abcd", + {"a": 8, "b": 8, "c": 8, "d": 8, "i": 32}, + id="star_contraction", + ), + # Bilinear / quadratic form over a batch (attention-score flavored). + pytest.param("bi,ij,bj->b", {"b": 128, "i": 64, "j": 64}, id="bilinear"), + # Batched matrix chain — batch axis rides through three contractions. + pytest.param( + "bij,bjk,bkl->bil", + {"b": 16, "i": 24, "j": 24, "k": 24, "l": 24}, + id="batched_chain", + ), + # Multi-index contraction surface: a whole axis-group (c) contracts at once. + pytest.param( + "abc,cde->abde", + {"a": 12, "b": 12, "c": 12, "d": 12, "e": 12}, + id="tensor_contraction", + ), + # Leading scalar factor plus an element-wise reduce — checks that the + # rank-0 operand threads through without spawning a degenerate axis. + pytest.param(",ij,ij->", {"i": 256, "j": 256}, id="scalar_scaled_reduce"), +] + + +def _make_operands(spec: str, sizes: dict[str, int], key: jax.Array) -> list[jax.Array]: + in_part = spec.split("->")[0] + in_specs = in_part.split(",") + keys = random.split(key, len(in_specs)) + return [ + random.normal(k, tuple(sizes[c] for c in s) if s else ()) + for k, s in zip(keys, in_specs, strict=True) + ] + + +@pytest.mark.parametrize( + "impl", [pytest.param(jnp.einsum, id="jax"), pytest.param(einsum, id="effectful")] +) +@pytest.mark.parametrize("spec,sizes", EINSUM_CASES) +@pytest.mark.benchmark(warmup=True, warmup_iterations=1) +def test_einsum_bench(benchmark, impl, spec, sizes, rng_key): + """Time one ``(spec, impl)`` pair. Group by ``spec`` to compare ``jnp`` + against ``effectful`` for the same subscript pattern (see module docstring). + """ + operands = _make_operands(spec, sizes, rng_key) + + @jax.jit + def f(*operands): + return impl(spec, *operands) + + @benchmark + def _run(): + return f(*operands).block_until_ready() + + +@pytest.mark.parametrize("spec,sizes", EINSUM_CASES) +def test_einsum_matches_jnp(spec: str, sizes, rng_key): + """``einsum`` returns the same result as ``jnp.einsum`` for every spec + in ``EINSUM_EXAMPLES``. + """ + operands = _make_operands(spec, sizes, rng_key) + actual = einsum(spec, *operands) + expected = jnp.einsum(spec, *operands) + assert actual.shape == expected.shape, ( + f"shape mismatch for {spec!r}: got {actual.shape}, expected {expected.shape}" + ) + assert jnp.allclose(actual, expected, atol=1e-4, rtol=1e-3), ( + f"value mismatch for {spec!r}" + ) + + +# see https://github.com/pyro-ppl/pyro/blob/dev/tests/ops/test_contract.py +# Let abcde be enum dims and ijk be plates. +PLATED_EINSUM_CASES = [ + ("abi,abi->", "i"), + ("abi,b->", "i"), + ("abi,b->i", "i"), + ("abi,b->b", "i"), + ("abi,b->ai", "i"), + ("aij,bi->", "ij"), + ("acij,bi->", "ij"), + ("aij,bi,c->", "ij"), + ("abij,bi->aij", "ij"), + ("abij,bci->acij", "ij"), + ("abij,bci->", "ij"), + ("ab,bcij->", "ij"), + ("ija,ika->", "ijk"), + ("ab,bcdi,defij,fgijk->", "ijk"), +] + + +@pytest.mark.parametrize("spec,plates", PLATED_EINSUM_CASES) +def test_plated_einsum(spec, plates, rng_key): + sys.setrecursionlimit(10_000) + + def _to_torch(arr): + return torch.from_dlpack(arr) + + def _to_mapping(arr): + return {i: float(arr[*i]) for i in np.ndindex(arr.shape)} + + dim_sizes = { + "a": 2, + "b": 3, + "c": 4, + "d": 5, + "e": 6, + "f": 7, + "g": 8, + "i": 2, + "j": 3, + "k": 4, + } + + operands = _make_operands(spec, dim_sizes, rng_key) + torch_operands = (_to_torch(op) for op in operands) + + try: + expected = pyro.ops.contract.naive_ubersum( + spec, *torch_operands, plates=plates, backend="torch" + )[0] + except NotImplementedError: + pytest.skip("Not implemented by pyro.ops.contract.einsum") + + actual = einsum(spec, *operands, plates=plates) + assert torch.allclose(torch.tensor(actual), expected, atol=1e-4, rtol=1e-4) diff --git a/tests/test_handlers_llm_provider.py b/tests/test_handlers_llm_provider.py index bd2b6ccd2..0f5dd5f70 100644 --- a/tests/test_handlers_llm_provider.py +++ b/tests/test_handlers_llm_provider.py @@ -250,7 +250,7 @@ def test_agent_tool_names_are_valid_integration(): agent = _ToolNameAgent() template = agent.ask tools = template.tools - expected_helper_tool_name = f"self__{agent.helper.__name__}" + expected_helper_tool_name = "self__helper" assert tools assert expected_helper_tool_name in tools assert all(re.fullmatch(r"[a-zA-Z0-9_-]+", name) for name in tools) diff --git a/tests/test_internals_disassembler.py b/tests/test_internals_disassembler.py new file mode 100644 index 000000000..22e484167 --- /dev/null +++ b/tests/test_internals_disassembler.py @@ -0,0 +1,2109 @@ +import ast +import collections.abc +import copy +import typing + +import pytest + +from effectful.internals.disassembly import ( + CompLambda, + DummyIterName, + disassemble, + ensure_ast, +) + + +def compile_and_eval( + node: ast.expr | ast.Expression, globals_dict: dict | None = None +) -> typing.Any: + """Compile an AST node and evaluate it.""" + if globals_dict is None: + globals_dict = {} + + # Wrap in an Expression node if needed + if not isinstance(node, ast.Expression): + node = ast.Expression(body=node) + + # Fix location info + ast.fix_missing_locations(node) + + # Compile and evaluate + code = compile(node, "", "eval") + return eval(code, globals_dict) + + +def materialize[T](genexpr: collections.abc.Generator[T, None, None]) -> list[T]: + """Materialize a nested generator expression to a nested list.""" + + def _materialize(genexpr): + if isinstance(genexpr, str | bytes): + return genexpr + elif isinstance(genexpr, collections.abc.Generator): + return [_materialize(item) for item in genexpr] + elif isinstance(genexpr, tuple): + # Kept as a tuple so that sets of tuples stay hashable + return tuple(_materialize(item) for item in genexpr) + elif isinstance(genexpr, collections.abc.Sequence): + return [_materialize(item) for item in genexpr] + elif isinstance(genexpr, collections.abc.Set): + return {_materialize(item) for item in genexpr} + elif isinstance(genexpr, collections.abc.Mapping): + return {_materialize(k): _materialize(v) for k, v in genexpr.items()} + else: + return genexpr + + return [_materialize(x) for x in genexpr] + + +def assert_ast_equivalent( + genexpr: collections.abc.Generator[typing.Any, None, None], + reconstructed_ast: ast.AST, + globals_dict: dict | None = None, +): + """Assert that a reconstructed AST produces the same results as the original generator.""" + # Check AST structure + assert isinstance(reconstructed_ast, ast.Expression) + assert hasattr(reconstructed_ast.body, "elt") # The expression part + assert hasattr(reconstructed_ast.body, "generators") # The comprehension part + assert len(reconstructed_ast.body.generators) > 0 + for comp in reconstructed_ast.body.generators: + assert hasattr(comp, "target") # Loop variable + assert hasattr(comp, "iter") # Iterator + assert hasattr(comp, "ifs") # Conditions + + # Save current globals to restore later + curr_globals = globals().copy() + globals().update(globals_dict or {}) + + # Materialize original generator to list for comparison + original_list = materialize(genexpr) + + # Clean up globals to avoid pollution + for key in globals_dict or {}: + if key not in curr_globals: + del globals()[key] + globals().update(curr_globals) + + # Compile and evaluate the reconstructed AST + reconstructed_gen = compile_and_eval(reconstructed_ast, globals_dict) + reconstructed_list = materialize(reconstructed_gen) + assert reconstructed_list == original_list, ( + f"AST produced {reconstructed_list}, expected {original_list}" + ) + + +# ============================================================================ +# BASIC GENERATOR EXPRESSION TESTS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Simple generator expressions + (x for x in range(5)), + (y for y in range(10)), + (item for item in [1, 2, 3]), + # Edge cases for simple generators + (i for i in range(0)), # Empty range + (n for n in range(1)), # Single item range + (val for val in range(100)), # Large range + (x for x in range(-5, 5)), # Negative range + (step for step in range(0, 10, 2)), # Step range + (rev for rev in range(10, 0, -1)), # Reverse range + ], +) +def test_simple_generators(genexpr): + """Test reconstruction of simple generator expressions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# ARITHMETIC AND EXPRESSION TESTS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Basic arithmetic operations + (x * 2 for x in range(5)), + (x + 1 for x in range(5)), + (x - 1 for x in range(5)), + (x**2 for x in range(5)), + (x % 2 for x in range(10)), + (x / 2 for x in range(1, 6)), + (x // 2 for x in range(10)), + # Complex expressions + (x * 2 + 1 for x in range(5)), + ((x + 1) * (x - 1) for x in range(5)), + (x**2 + 2 * x + 1 for x in range(5)), + # Unary operations + (-x for x in range(5)), + (+x for x in range(-5, 5)), + (~x for x in range(5)), + # More complex arithmetic edge cases + (x**3 for x in range(1, 5)), # Higher powers + (x * x * x for x in range(5)), # Repeated multiplication + (x + x + x for x in range(5)), # Repeated addition + (x - x + 1 for x in range(5)), # Operations that might simplify + (x / x for x in range(1, 5)), # Division by self + (x % (x + 1) for x in range(1, 10)), # Modulo with expression + # Nested arithmetic expressions + ((x + 1) ** 2 for x in range(5)), + ((x * 2 + 3) * (x - 1) for x in range(5)), + (x * (x + 1) * (x + 2) for x in range(5)), + # Mixed operations with precedence + (x + 3 * 2 for x in range(3)), + (x * 2 + 9 / 3 for x in range(1, 4)), + ((x + 2) * (x - 2) for x in range(1, 4)), + # Edge cases with zero and one + (x * 0 for x in range(5)), + (x * 1 for x in range(5)), + (x + 0 for x in range(5)), + (x**1 for x in range(5)), + (0 + x for x in range(5)), + (1 * x for x in range(5)), + ], +) +def test_arithmetic_expressions(genexpr): + """Test reconstruction of generators with arithmetic expressions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# FSTRING EXPRESSIONS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Basic f-string cases + (f"{x}" for x in range(5)), # Single value, no format + (f"{x} is {x**2}" for x in range(5)), # Multiple values + (f"{x:02d}" for x in range(10)), # Format spec + (f"{x:.2f}" for x in [1.2345, 2.3456, 3.4567]), # Float format spec + # Conversion specifiers + (f"{x!r}" for x in ["hello", "world"]), # repr conversion + (f"{x!s}" for x in [1, 2, 3]), # str conversion + (f"{x!a}" for x in ["hello\n", "world\t"]), # ascii conversion + # Conversion with format spec + (f"{x!r:>10}" for x in ["hello", "world"]), # repr with alignment + (f"{x!s:^15}" for x in [1, 2, 3]), # str with center align + # Empty and literal f-strings + ("" for x in range(3)), # Empty f-string + ("constant" for x in range(3)), # No formatting + (f"x={x}" for x in range(5)), # Literal prefix + (f"result: {x * 2}" for x in range(5)), # Literal with expression + # Complex expressions in f-strings + (f"{x + 1}" for x in range(5)), # Arithmetic + (f"{x * x}" for x in range(5)), # Multiplication + (f"{x % 2}" for x in range(10)), # Modulo + (f"{-x}" for x in range(-2, 3)), # Unary minus + # Nested formatting + (f"{x:0{2}d}" for x in range(5)), # Format spec with expression + (f"{x:>{3 * 2}}" for x in range(5)), # Expression in format spec + # Multiple formatted values + (f"{x} + {y} = {x + y}" for x in range(3) for y in range(3)), # Multiple vars + (f"({x}, {y})" for x in range(2) for y in range(2)), # Tuple display + # F-strings with various data types + (f"{s}" for s in ["hello", "world"]), # Strings + (f"{b}" for b in [True, False]), # Booleans + (f"{n}" for n in [None, None]), # None values + (f"{lst}" for lst in [[1, 2], [3, 4]]), # Lists + # Complex format specifications + (f"{x:+05d}" for x in range(-2, 3)), # Sign, zero pad, width + (f"{x:.2%}" for x in [0.1, 0.25, 0.333]), # Percentage format + (f"{x:.2e}" for x in [100, 1000, 10000]), # Scientific notation + (f"{x:#x}" for x in [10, 15, 255]), # Hex with prefix + (f"{x:b}" for x in [2, 7, 15]), # Binary format + # Edge cases + ("{x}" for x in range(3)), # Escaped braces + (f"{{x}} = {x}" for x in range(3)), # Mixed escaped/formatted + (f"{{{x}}}" for x in range(3)), # Brace around formatted + ], +) +def test_fstring_expressions(genexpr): + """Test reconstruction of generators with f-string expressions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# COMPARISON OPERATORS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # All comparison operators + (x for x in range(10) if x < 5), + (x for x in range(10) if x <= 5), + (x for x in range(10) if x > 5), + (x for x in range(10) if x >= 5), + (x for x in range(10) if x == 5), + (x for x in range(10) if x != 5), + # in/not in operators + (x for x in range(10) if x in [2, 4, 6, 8]), + (x for x in range(10) if x not in [2, 4, 6, 8]), + # is/is not operators (with None) + (x for x in [1, None, 3, None, 5] if x is not None), + (x for x in [1, None, 3, None, 5] if x is None), + # Boolean operations - these are complex cases that might need special handling + (x for x in range(10) if not x % 2), + (x for x in range(10) if not (x > 5)), + (x for x in range(10) if x > 2 and x < 8), + (x for x in range(10) if x < 3 or x > 7), + # More complex comparison edge cases + # Comparisons with expressions + (x for x in range(10) if x * 2 > 10), + (x for x in range(10) if x + 1 <= 5), + (x for x in range(10) if x**2 < 25), + (x for x in range(10) if (x + 1) * 2 != 6), + # Complex membership tests + (x for x in range(20) if x in range(5, 15)), + (x for x in range(10) if x not in range(3, 7)), + (x for x in range(10) if x % 2 in [0]), + (x for x in range(10) if x not in []), # Empty container + # Complex boolean combinations + (x for x in range(20) if not (x < 5 or x > 15)), + (x for x in range(20) if x > 5 and x < 15 and x % 2 == 0), + (x for x in range(20) if x < 5 or x > 15 or x == 10), + (x for x in range(20) if not (x > 5 and x < 15)), + # Mixed comparison and boolean operations + (x for x in range(20) if (x > 10 and x % 2 == 0) or (x < 5 and x % 3 == 0)), + (x for x in range(20) if not (x % 2 == 0 and x % 3 == 0)), + # Edge cases with identity comparisons + (x for x in [0, 1, 2, None, 4] if x is not None and x > 1), + (x for x in [True, False, 1, 0] if x is True), + (x for x in [True, False, 1, 0] if x is not False), + ], +) +def test_comparison_operators(genexpr): + """Test reconstruction of all comparison operators.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# CHAINED COMPARISON TESTS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Chained comparisons + (x for x in range(20) if 5 < x < 15), + (x for x in range(20) if 0 <= x <= 10), + ], +) +def test_chained_comparison_operators(genexpr): + """Test reconstruction of chained (ternary) comparison operators.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# FILTERED GENERATOR TESTS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Simple filters + (x for x in range(10) if x % 2 == 0), + (x for x in range(10) if x > 5), + (x for x in range(10) if x < 5), + (x for x in range(10) if x != 5), + # Complex filters + (x for x in range(20) if x % 2 == 0 if x % 3 == 0), + (x for x in range(100) if x > 10 if x < 90 if x % 5 == 0), + # Filters with expressions + (x * 2 for x in range(10) if x % 2 == 0), + (x**2 for x in range(10) if x > 3), + # Boolean operations in filters + (x for x in range(10) if not x % 2), + (x for x in range(10) if x > 2 and x < 8), + (x for x in range(10) if x < 3 or x > 7), + # More complex filter edge cases + (x for x in range(50) if x % 7 == 0), # Different modulo + (x for x in range(10) if x >= 0), # Always true condition + (x for x in range(10) if x < 0), # Always false condition + ( + x for x in range(20) if x % 2 == 0 and x % 3 == 0 + ), # Multiple conditions with and + ( + x for x in range(20) if x % 2 == 0 or x % 3 == 0 + ), # Multiple conditions with or + # Nested boolean operations + (x for x in range(20) if (x > 5 and x < 15) or x == 0), + (x for x in range(20) if not (x > 10 and x < 15)), + (x for x in range(50) if x > 10 and (x % 2 == 0 or x % 3 == 0)), + # Multiple consecutive filters + (x for x in range(100) if x > 20 if x < 80 if x % 10 == 0), + (x for x in range(50) if x % 2 == 0 if x % 3 != 0 if x > 10), + # Filters with complex expressions + (x + 1 for x in range(20) if (x * 2) % 3 == 0), + (x**2 for x in range(10) if x * (x + 1) > 10), + (x / 2 for x in range(1, 20) if x % (x // 2 + 1) == 0), + # Edge cases with truthiness + (x for x in range(10) if x), # Truthy filter + (x for x in range(-5, 5) if not x), # Falsy filter + (x for x in range(10) if bool(x % 2)), # Explicit bool conversion + ], +) +def test_filtered_generators(genexpr): + """Test reconstruction of generators with if conditions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# NESTED LOOP TESTS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Basic nested loops + ((x, y) for x in range(3) for y in range(3)), + (x + y for x in range(3) for y in range(3)), + (x * y for x in range(1, 4) for y in range(1, 4)), + # Nested with filters + ((x, y) for x in range(5) for y in range(5) if x < y), + (x + y for x in range(5) if x % 2 == 0 for y in range(5) if y % 2 == 1), + # Triple nested + (x + y + z for x in range(2) for y in range(3) for z in range(4)), + ((x, y, z) for x in range(2) for y in range(3) for z in range(4)), + # More complex nested loop edge cases + # Different sized ranges + ((x, y) for x in range(2) for y in range(5)), + ((x, y) for x in range(10) for y in range(2)), + # Asymmetric operations + (x - y for x in range(5) for y in range(3)), + (x / (y + 1) for x in range(1, 6) for y in range(3)), + (x**y for x in range(1, 4) for y in range(3)), + # Complex expressions with multiple variables + (x * y + x for x in range(3) for y in range(3)), + (x + y + x * y for x in range(1, 4) for y in range(1, 4)), + ((x + y) ** 2 for x in range(3) for y in range(3)), + # Filters on different loop levels + ((x, y) for x in range(10) if x % 2 == 0 for y in range(10) if y % 3 == 0), + (x * y for x in range(5) for y in range(5) if x != y), + (x + y for x in range(5) for y in range(5) if x + y < 5), + # Triple and quadruple nested with various patterns + (x + y + z for x in range(2) for y in range(2) for z in range(2)), + (x * y * z for x in range(1, 3) for y in range(1, 3) for z in range(1, 3)), + ( + (x, y, z, w) + for x in range(2) + for y in range(2) + for z in range(2) + for w in range(2) + ), + # Nested loops with complex filters + ((x, y) for x in range(5) if x > 1 for y in range(5) if x < y), + (x + y for x in range(3) if x > 0 for y in range(3)), + # Mixed range types + ((x, y) for x in range(-2, 2) for y in range(0, 4, 2)), + (x * y for x in range(5, 0, -1) for y in range(1, 6)), + # Dependent nested loops + ((x, y) for x in range(3) for y in range(x, 3)), + (x + y for x in range(3) for y in range(x + 1, 3)), + ], +) +def test_nested_loops(genexpr): + """Test reconstruction of generators with nested loops.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# =========================================================================== +# NESTED COMPREHENSIONS +# =========================================================================== + + +@pytest.mark.parametrize( + "genexpr", + [ + # nested generators + ((x for x in range(i + 1)) for i in range(5)), + ((x for j in range(i) for x in range(j)) for i in range(5)), + (((x for x in range(i + j)) for j in range(i)) for i in range(5)), + # nested generators with filters + ((x for x in range(i)) for i in range(5) if i > 0), + ((x for x in range(i) if x < i) for i in range(5) if i > 0), + (((x for x in range(i + j) if x < i + j) for j in range(i)) for i in range(5)), + # aggregation function call + (sum(x for x in range(i + 1)) for i in range(3)), + (max(x for x in range(i + 1)) for i in range(3)), + (dict((x, x + 1) for x in range(i + 1)) for i in range(3)), + (set(x for x in range(i + 1)) for i in range(3)), + # map + (list(map(abs, (x + 1 for x in range(i + 1)))) for i in range(3)), + (list(enumerate(x + 1 for x in range(i + 1))) for i in range(3)), + # nesting on both sides + ((y for y in range(x)) for x in (x_ + 1 for x_ in range(5))), + ((y for y in range(x)) for x in (x_ + 1 for x_ in range(5))), + ], +) +def test_nested_comprehensions(genexpr): + """Test reconstruction of nested comprehensions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +def test_nested_comprehensions_multiline(): + """The same filter reconstructs the same way however the source is laid out. + + On Python 3.12 these two spellings disassemble to different jump layouts -- + only the one-line form emits POP_JUMP_IF_TRUE -- which used to make the + multiline form come out negated. + """ + one_line = (x for x in range(5) if x > 1) + assert_ast_equivalent(one_line, disassemble(one_line)) + + multiline = ( + x + for x in range(5) # comment to avoid reformatting + if x > 1 + ) + assert_ast_equivalent(multiline, disassemble(multiline)) + + assert ast.unparse(disassemble(x for x in range(5) if x > 1)) == ast.unparse( + disassemble( + x + for x in range(5) # comment to avoid reformatting + if x > 1 + ) + ) + + +# ============================================================================ +# DIFFERENT COMPREHENSION TYPES +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Comprehensions as iterator constants + (x_ for x_ in [x for x in range(5)]), + (x_ for x_ in {x for x in range(5)}), + (x_ for x_ in {x: x**2 for x in range(5)}), + # Comprehensions as yield expressions + ([y * 2 for y in range(x + 1)] for x in range(3)), + ({y + 3 for y in range(x + 1)} for x in range(3)), + ({y: y**2 for y in range(x + 1)} for x in range(3)), + # nested non-generators + ([x for x in range(i)] for i in range(5)), + ([x for j in range(i) for x in range(j)] for i in range(5)), + ({x: x**2 for x in range(i)} for i in range(5)), + ([[x for x in range(i + j)] for j in range(i)] for i in range(5)), + # Nested comprehensions with filters inside + ([x for x in range(i)] for i in range(5) if i > 0), + ([x for x in range(i) if x < i] for i in range(5) if i > 0), + ([[x for x in range(i + j) if x < i + j] for j in range(i)] for i in range(5)), + ], +) +def test_different_comprehension_types(genexpr): + """Test reconstruction of different comprehension types.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# DICT DISPLAYS +# +# A dict display with dynamic keys is built by BUILD_MAP from key/value pairs +# that the compiler pushed in source order. Dicts compare equal whatever order +# they were built in, so these tests pin the order down directly: which of two +# equal keys wins, what `items()` yields, and when each subexpression runs. +# ============================================================================ + + +_EVAL_ORDER: list[str] = [] + + +def _note(tag: str, value: typing.Any) -> typing.Any: + """Record that this subexpression was evaluated, and pass its value along.""" + _EVAL_ORDER.append(tag) + return value + + +def test_dict_display_duplicate_keys(): + """The last of two equal keys wins, so the pairs must keep their order.""" + genexpr = ({x: "first", x: "second"} for x in range(1)) # noqa: F602 + reconstructed = disassemble(genexpr) + assert ast.unparse(reconstructed) == ( + "({x: 'first', x: 'second'} for x in range(0, 1, 1))" + ) + assert materialize(genexpr) == [{0: "second"}] + assert materialize(compile_and_eval(reconstructed)) == [{0: "second"}] + + +def test_dict_display_insertion_order(): + genexpr = ({x: "a", x + 1: "b", x + 2: "c"} for x in range(1)) + reconstructed = disassemble(genexpr) + expected = [[(0, "a"), (1, "b"), (2, "c")]] + assert [list(d.items()) for d in genexpr] == expected + assert [list(d.items()) for d in compile_and_eval(reconstructed)] == expected + + +def test_dict_display_evaluation_order(): + """Keys and values run left to right, key before its own value.""" + genexpr = ( + {_note("k1", "a"): _note("v1", 1), _note("k2", "b"): _note("v2", 2)} + for _ in range(1) + ) + reconstructed = disassemble(genexpr) + + _EVAL_ORDER.clear() + assert materialize(genexpr) == [{"a": 1, "b": 2}] + assert _EVAL_ORDER == ["k1", "v1", "k2", "v2"] + + _EVAL_ORDER.clear() + assert materialize(compile_and_eval(reconstructed, {"_note": _note})) == [ + {"a": 1, "b": 2} + ] + assert _EVAL_ORDER == ["k1", "v1", "k2", "v2"] + + +# ============================================================================ +# CONDITIONAL EXPRESSIONS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # simple conditional expressions without nesting + ((lambda x: x if x % 2 == 0 else -x)(xi) for xi in range(5)), + ((lambda x: (x + 1) if x < 5 else (x - 1))(xi) for xi in range(10)), + ((lambda x: (x * 2) if x > 0 else (x / 2))(xi) for xi in range(-5, 5)), + ((lambda x: (x**2) if x != 0 else 1)(xi) for xi in range(-3, 4)), + # simple conditional expressions with negation + ((lambda x: (x + 10) if not (x < 5) else (x - 10))(xi) for xi in range(20)), + ((lambda x: (x * 3) if not (x % 2 == 0) else (x // 3))(xi) for xi in range(10)), + ((lambda x: (x**3) if not (x < 0) else (x**0.5))(xi) for xi in range(-5, 15)), + # conditional expressions with lazy test + ( + (lambda x: (x + 10) if (x > 5 and x < 15) else (x - 10))(xi) + for xi in range(20) + ), + ( + (lambda x: (x * 3) if (x % 2 == 0 or x % 3 == 0) else (x // 3))(xi) + for xi in range(10) + ), + ( + (lambda x: (x**3) if not (x < 0 or x > 10) else (x**0.5))(xi) + for xi in range(-5, 15) + ), + ], +) +def test_conditional_expressions_simple_no_comprehension(genexpr): + """Test reconstruction of simple conditional expressions isolated from comprehension bodies.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # nested conditional expressions + ( + (lambda x: (x + 1) if x < 5 else ((x - 1) if x < 10 else (x * 2)))(xi) + for xi in range(15) + ), + ( + ( + lambda x: ( + (x * 2) if x % 2 == 0 else ((x // 2) if x % 3 == 0 else (x + 2)) + ) + )(xi) + for xi in range(10) + ), + ( + (lambda x: (x**2) if x > 0 else ((-x) ** 2 if x < -5 else 1))(xi) + for xi in range(-10, 5) + ), + ], +) +def test_conditional_expressions_nested_no_comprehension(genexpr): + """Test reconstruction of nested conditional expressions isolated from comprehension bodies.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # Basic conditional expressions in comprehension bodies + ((x if x % 2 == 0 else -x) for x in range(5)), + ((x * 2 if x > 0 else x / 2) for x in range(-3, 4)), + ((x**2 if x != 0 else 1) for x in range(-2, 3)), + # Conditional expressions with filters + ((x if x % 2 == 0 else -x) for x in range(10) if x > 2), + ((x * 3 if x > 5 else x + 1) for x in range(20) if x % 3 == 0), + # Nested loops with conditional expressions + ((x + y if x > y else x - y) for x in range(3) for y in range(3)), + ( + (x * y if x != 0 and y != 0 else 0) + for x in range(-2, 3) + for y in range(-2, 3) + ), + # Multiple conditional expressions + ( + (x if x > 0 else 0) + (y if y > 0 else 0) + for x in range(-2, 3) + for y in range(-2, 3) + ), + # Conditional expressions in different parts + ([x if x > 0 else -x for x in range(i)] for i in range(1, 4)), + ((x if x % 2 == 0 else -x) for x in (y if y > 2 else y + 10 for y in range(5))), + # Complex nested conditional expressions + ((x if x > 0 else (x + 5 if x > -3 else x * 2)) for x in range(-5, 5)), + ((x * 2 if x > 0 else (x / 2 if x < 0 else 1)) for x in range(-3, 4)), + # Conditional expressions with function calls + ((abs(x) if x < 0 else x) for x in range(-3, 4)), + ((max(x, 0) if x is not None else 0) for x in [None, -1, 0, 1, 2]), + # Mixed with other complex expressions + ((x + 1 if x % 2 == 0 else x - 1) * 2 for x in range(5)), + ((x, y, x + y if x > y else x - y) for x in range(3) for y in range(3)), + ], +) +def test_conditional_expressions_simple_comprehensions(genexpr): + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # `and` chains compile to a run of jumps that each fall through to the + # loop back-edge, which the disassembler folds into a single filter. + (x for x in range(10) if x > 2 and x < 8), + (x for x in range(20) if x > 5 and x % 2 == 0 and x < 15), + (x for x in range(-10, 10) if abs(x) > 3 and x % 2 == 0), + (x for x in ["hello", "world", "test"] if len(x) > 3 and x.startswith("h")), + (x for x in range(20) if x % 2 == 0 and x % 3 == 0 and x > 0 and x < 18), + ((x, y) for x in range(5) for y in range(5) if x < y and x + y > 2), + # `or` in filter position: reconstructed incorrectly, see the marker. + (x for x in range(10) if x < 3 or x > 7), + (x for x in range(20) if x < 5 or x > 15 or x == 10), + (x for x in range(20) if (x > 10 and x % 2 == 0) or (x < 5 and x % 3 == 0)), + (x for x in range(20) if x > 5 and (x < 10 or x > 15)), + (x for x in range(100) if (x > 10 and x < 50) and (x % 3 == 0 or x % 5 == 0)), + # `not (a and b)` is compiled exactly like `not a or not b`. + (x for x in range(100) if not (x > 30 and x < 70)), + # Chained comparisons in filter position. + (x for x in range(20) if 5 < x < 15), + (x for x in range(20) if 0 <= x <= 10), + (x for x in range(50) if 10 < x < 20 < x * 2), + (x for x in range(10) if 0 <= x <= 5 <= x + 5), + (x for x in range(50) if 5 < x < 15 and x % 2 == 0), + (x for x in range(50) if x > 20 or 5 < x < 15), + ], +) +def test_lazy_boolean_and_chained_comparisons_in_filters(genexpr): + """Lazy boolean operators and chained comparisons in *filter* position. + + This is the hard case: a filter's condition is recognised structurally, by + the jump falling through to the loop back-edge, so any condition CPython + compiles with an intermediate join point is misread. + """ + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # The same operators in *ternary* position all work: both arms produce a + # value, so the fork/merge machinery in _symbolic_exec applies directly. + ((x if x > 5 and x < 15 else 0) for x in range(20)), + ((x if x < 3 or x > 17 else -x) for x in range(20)), + ((x if 5 < x < 15 else 0) for x in range(20)), + ((x * 2 if 0 <= x <= 10 else x / 2) for x in range(-5, 15)), + ((x if x > 2 and x < 8 else -x) for x in range(10)), + ((x if x < 2 or x > 8 else -x) for x in range(10)), + ((x if not (x > 2 and x < 8) else -x) for x in range(10)), + ((x if 0 <= x <= 5 <= x + 5 else -x) for x in range(10)), + ((x if x > 1 and x < 9 or x == 0 else -x) for x in range(10)), + # ... including nested inside another ternary + ((x if x > 5 or x < 2 else (0 if x == 3 else 1)) for x in range(10)), + ((x if x > 5 else (0 if 2 < x < 4 else 1)) for x in range(10)), + ], +) +def test_lazy_boolean_and_chained_comparisons_in_ternaries(genexpr): + """Lazy boolean operators and chained comparisons in conditional expressions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +def test_short_circuit_filter_is_a_disjunction_of_paths(): + """A short-circuiting filter is rebuilt as the disjunction over its paths. + + `x < 3 or x > 7` used to be misread as a conditional expression, yielding an + AST that compiled and ran but computed `[0, 1, 2, False, False, ...]` instead + of `[0, 1, 2, 8, 9]`. Each path through the condition now contributes one + disjunct, so the reconstruction is equivalent rather than merely plausible. + """ + genexpr = (x for x in range(10) if x < 3 or x > 7) + reconstructed = disassemble(genexpr) + + assert isinstance(reconstructed.body, ast.GeneratorExp) + filters = reconstructed.body.generators[0].ifs + assert len(filters) == 1 + assert isinstance(filters[0], ast.BoolOp) and isinstance(filters[0].op, ast.Or) + assert materialize(compile_and_eval(reconstructed)) == [0, 1, 2, 8, 9] + + +@pytest.mark.parametrize( + "genexpr", + [ + # Simple conditional as function argument + (max(x if x > 0 else 0, 1) for x in range(-2, 3)), + (abs(x if x < 0 else -x) for x in range(-3, 3)), + (len(str(x) if x > 10 else "small") for x in range(15)), + # Multiple conditional arguments + ( + max(x if x > 0 else 0, y if y > 0 else 0) + for x in range(-1, 2) + for y in range(-1, 2) + ), + ( + pow(x if x != 0 else 1, y if y > 0 else 1) + for x in range(3) + for y in range(3) + ), + # Nested function calls with conditionals + (max(abs(x if x < 0 else -x), 1) for x in range(-3, 4)), + (int(str(x if x > 5 else x + 10)) for x in range(10)), + # Conditionals in keyword arguments (using dict constructor as example) + (dict(a=x if x > 0 else 0, b=x * 2 if x < 5 else x) for x in range(8)), + # Method calls with conditional arguments + ([1, 2, 3].index(x if x in [1, 2, 3] else 1) for x in range(5)), + ("hello".replace("l", x if isinstance(x, str) else "X") for x in ["a", 1, "b"]), + # Complex nested case: conditional in function argument, function call in conditional + (abs(x if len(str(x)) > 1 else x * 10) for x in range(15)), + # Mixed: conditional in function call within comprehension filter + (x for x in range(20) if max(x if x > 10 else 0, 5) > 8), + ], +) +def test_conditional_expressions_function_arguments(genexpr): + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# GENERATOR EXPRESSION WITH GLOBALS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr,globals_dict", + [ + # Using constants + ((x + a for x in range(5)), {"a": 10}), # type: ignore # noqa: F821 + ((data[i] for i in range(2)), {"data": [3, 4]}), # type: ignore # noqa: F821 + # Using global functions + ((abs(x) for x in range(-5, 5)), {"abs": abs}), + ((len(s) for s in ["a", "ab", "abc"]), {"len": len}), + ((max(x, 5) for x in range(10)), {"max": max}), + ((min(x, 5) for x in range(10)), {"min": min}), + ((round(x / 3, 2) for x in range(10)), {"round": round}), + ], +) +def test_variable_lookup(genexpr, globals_dict): + """Test reconstruction of expressions with globals.""" + ast_node = disassemble(genexpr) + + # Need to provide the same globals for evaluation + assert_ast_equivalent(genexpr, ast_node, globals_dict) + + +# ============================================================================ +# EDGE CASES AND COMPLEX SCENARIOS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr,globals_dict", + [ + # Using lambdas and functions + (((lambda y: y * 2)(x) for x in range(5)), {}), + (((lambda y: y + 1)(x) for x in range(5)), {}), + (((lambda y: y**2)(x) for x in range(5)), {}), + (((lambda a, b: a + b)(x, x) for x in range(5)), {}), + (((lambda: (x for x in range(i)))() for i in range(3)), {}), + ((f(x) for x in range(5)), {"f": lambda y: y * 3}), # type: ignore # noqa: F821 + # Attribute access + ((x.real for x in [1 + 2j, 3 + 4j, 5 + 6j]), {}), + ((x.imag for x in [1 + 2j, 3 + 4j, 5 + 6j]), {}), + ((x.conjugate() for x in [1 + 2j, 3 + 4j, 5 + 6j]), {}), + # slicing and indexing + ((s[:2] for s in ["hello", "world"]), {}), + ((s[1:3] for s in ["hello", "world"]), {}), + ((s[-1] for s in ["hello", "world"]), {}), + ((s[0:3] for s in ["hello", "world"]), {}), + ((s[::-1] for s in ["hello", "world"]), {}), + ((s[1:2:] for s in ["hello", "world"]), {}), + # Method calls + ((s.upper() for s in ["hello", "world"]), {}), + ((s.lower() for s in ["HELLO", "WORLD"]), {}), + ((s.strip() for s in [" hello ", " world "]), {}), + ((x.bit_length() for x in range(1, 10)), {}), + ((str(x).zfill(3) for x in range(10)), {"str": str}), + # Subscript operations + (((10, 20, 30)[i] for i in range(3)), {}), + (([10, 20, 30][i] for i in range(3)), {}), + (({"a": 1, "b": 2, "c": 3}[k] for k in ["a", "b", "c"]), {}), + (("hello"[i] for i in range(5)), {}), + ((data[i][j] for i in range(2) for j in range(2)), {"data": [[1, 2], [3, 4]]}), # type: ignore # noqa: F821 + # # More complex attribute chains + # ((obj.value.bit_length() for obj in [type('', (), {'value': x})() for x in range(1, 5)]), {}), + # Multiple function calls + ((abs(max(x, -x)) for x in range(-3, 4)), {"abs": abs, "max": max}), + ((len(str(x)) for x in range(100, 110)), {"len": len, "str": str}), + # Mixed operations + ( + (abs(x) + len(str(x)) for x in range(-10, 10)), + {"abs": abs, "len": len, "str": str}, + ), + ((s.upper().lower() for s in ["Hello", "World"]), {}), + # Edge cases with complex data structures + (((1, 2, 3)[x % 3] for x in range(10)), {}), + (([1, 2, 3][x % 3] for x in range(10)), {}), + (({1, 2, 3} for x in range(10)), {}), + # (({"even": x, "odd": x + 1}["even" if x % 2 == 0 else "odd"] for x in range(5)), {}), + # Function calls with multiple arguments + ((pow(x, 2, 10) for x in range(5)), {"pow": pow}), + ((divmod(x, 3) for x in range(10)), {"divmod": divmod}), + ], +) +def test_complex_scenarios(genexpr, globals_dict): + """Test reconstruction of complex generator expressions.""" + ast_node = disassemble(genexpr) + + # Need to provide the same globals for evaluation + assert_ast_equivalent(genexpr, ast_node, globals_dict) + + +# ============================================================================ +# UNPACKING LOOP TARGETS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Simple tuple targets + ((a, b) for a, b in [(1, 2), (3, 4)]), + (a + b for a, b in [(1, 2), (3, 4)]), + (a * b for a, b in [(2, 3), (4, 5)]), + ((b, a) for a, b in [(1, 2), (3, 4)]), + ((a, b, c) for a, b, c in [(1, 2, 3), (4, 5, 6)]), + ((a, b, c, d) for a, b, c, d in [(1, 2, 3, 4)]), + # Nested tuple targets + ((a, b, c) for a, (b, c) in [(1, (2, 3)), (4, (5, 6))]), + ((a, b, c) for (a, b), c in [((1, 2), 3)]), + ((a, b, c, d) for (a, b), (c, d) in [((1, 2), (3, 4))]), + ((a, b, c) for a, (b, (c,)) in [(1, (2, (3,)))]), + # Unpacking over dict views + ((k, v) for k, v in {1: "a", 2: "b"}.items()), + (v for k, v in {1: "a", 2: "b"}.items()), + # Unpacking combined with filters + ((a, b) for a, b in [(1, 2), (3, 1)] if a < b), + (a + b for a, b in [(1, 2), (3, 4)] if a % 2 == 0), + ((a, b) for a, b in [(1, 2), (3, 4)] if a > 0 if b > 3), + # Unpacking in nested loops, in either position + ((x, a, b) for x in range(2) for a, b in [(1, 2), (3, 4)]), + ((a, b, y) for a, b in [(1, 2)] for y in range(2)), + ((a, b, c, d) for a, b in [(1, 2)] for c, d in [(3, 4), (5, 6)]), + ((a, b, x) for a, b in [(1, 2), (3, 4)] for x in range(a)), + # Unpacking inside other comprehension types + ([a for a, b in [(1, 2), (3, 4)]] for _ in range(2)), + ({a for a, b in [(1, 2), (3, 4)]} for _ in range(2)), + ({a: b for a, b in [(1, 2), (3, 4)]} for _ in range(2)), + ((a for a, b in [(1, 2), (3, 4)]) for _ in range(2)), + # Unpacking with a conditional expression in the body + ((a if a > b else b) for a, b in [(1, 2), (4, 3)]), + ], +) +def test_unpacking_targets(genexpr): + """Test reconstruction of comprehensions that unpack their loop target.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + ((x, a, b) for x in range(2) for a, b in [(1, 2)]), + ((a, b, c, d) for a, b in [(1, 2)] for c, d in [(3, 4)]), + ((x, a, b) for x in range(2) for a, *b in [(1, 2, 3)]), + ((x, a, b, c) for x in range(2) for a, (b, c) in [(1, (2, 3))]), + ((x, a) for x in range(2) for (a,) in [(1,)]), + ((x, a, b) for x in range(2) for a, *b in [(1,)]), # type: ignore[var-annotated] + ((x, y, a) for x in range(2) for y in range(2) for a, b in [(1, 2)]), + ], +) +def test_unpacking_over_single_element_literal(genexpr): + """A one-element inner loop over a literal. + + Python 3.14 unrolls this: it assigns the targets outright and emits no + FOR_ITER, so the loop is not there to be recovered. The names it bound are + substituted at their uses instead, which reproduces the same elements. + """ + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # Starred last: UNPACK_EX with only a "before" count + (a for a, *b in [(1, 2, 3), (4, 5, 6)]), + (b for a, *b in [(1, 2, 3), (4, 5, 6)]), + ((a, b) for a, *b in [(1, 2, 3)]), + ((a, b, c) for a, b, *c in [(1, 2, 3, 4)]), + # Starred first: the "after" count lives in the high byte of the + # argument, so the instruction is prefixed with EXTENDED_ARG + (a for *a, b in [(1, 2, 3), (4, 5, 6)]), + (b for *a, b in [(1, 2, 3), (4, 5, 6)]), + ((a, b) for *a, b in [(1, 2, 3)]), + # Starred in the middle + ((a, b, c) for a, *b, c in [(1, 2, 3, 4)]), + ((a, b, c) for a, *b, c in [(1, 2, 3, 4, 5)]), + # Starred target that collects nothing + ((a, b) for a, *b in [(1,)]), # type: ignore[var-annotated] + # Combined with filters, nesting and other comprehension types + ((a, b) for a, *b in [(1, 2), (3, 4)] if a > 1), + ((x, a, b) for x in range(2) for a, *b in [(1, 2, 3), (4, 5, 6)]), + ([a for a, *b in [(1, 2, 3)]] for _ in range(2)), + ], +) +def test_unpacking_starred_targets(genexpr): + """Test reconstruction of starred loop targets (UNPACK_EX).""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# OUTERMOST ITERABLE TYPES +# +# The outermost iterable is not part of the comprehension's bytecode: it is a +# live object reachable through `gi_frame.f_locals[".0"]`, so `ensure_ast` has +# to rebuild an expression for it from the object alone. +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Strings and bytes + (c for c in "hello"), + (c.upper() for c in "hello" if c != "l"), + (c for c in "h\xe9llo"), # non-ASCII takes a different iterator type + (b for b in b"abc"), + (b for b in bytearray(b"abc")), + # Sequences + (x for x in [1, 2, 3]), + (x for x in (1, 2, 3)), + (x for x in range(3)), + # Sets and frozensets + (x for x in {1, 2, 3}), + (x for x in frozenset({1, 2, 3})), + # Dict views + (k for k in {1: "a", 2: "b"}), + (k for k in {1: "a", 2: "b"}.keys()), + (v for v in {1: "a", 2: "b"}.values()), + (kv for kv in {1: "a", 2: "b"}.items()), + # reversed() over each of the underlying sequence types + (x for x in reversed([1, 2, 3])), + (x for x in reversed((1, 2, 3))), + (c for c in reversed("abc")), + (x for x in reversed(range(3))), + # Comprehensions as the outermost iterable + (x for x in (y for y in range(3))), + (x for x in [y for y in range(3)]), + (x for x in {y for y in range(3)}), + # Nested/structured contents + (t for t in [(1, 2), (3, 4)]), + (d for d in [{"a": 1}, {"b": 2}]), + (x for x in [[1, 2], [3, 4]]), + ], +) +def test_outermost_iterable_types(genexpr): + """Test reconstruction of the outermost iterable from the live object.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # zip/enumerate/map/filter wrap other iterators rather than a concrete + # sequence, but pickle with their constituent parts. + (x for x in zip([1, 2], [3, 4])), + ((a, b) for a, b in zip([1, 2], [3, 4])), + (x for x in zip("ab", range(2), [7, 8])), + (x for x in enumerate("ab")), + ((i, c) for i, c in enumerate("abc")), + (x for x in enumerate("ab", 5)), + (x for x in map(abs, [-1, 2])), + (x for x in map(max, [1, 2], [3, 0])), + (x for x in filter(None, [0, 1, 2])), + (x for x in filter(bool, [0, 1, 2])), + # A strict zip over equal-length iterables behaves like a lax one, but + # its strictness has to survive anyway -- see the ragged case below. + ((a, b) for a, b in zip([1, 2], [3, 4], strict=True)), + (x for x in zip("ab", range(2), [7, 8], strict=True)), + (x for x in zip(range(2), map(abs, [-1, -2]), strict=True)), + # Nested adaptors, and adaptors over non-sequence iterables + (x for x in zip(range(2), map(abs, [-1, -2]))), + (x for x in enumerate(filter(None, [0, 1]))), + (x for x in map(abs, range(-2, 2))), + (x for x in zip("ab", (y for y in range(2)))), + # With a filter and a non-trivial element expression + (a + b for a, b in zip([1, 2], [3, 4]) if a > 1), + ], +) +def test_outermost_iterable_adaptors(genexpr): + """zip/enumerate/map/filter are rebuilt from the parts they pickle with.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # A conditional expression as an inner loop's iterable. Its value is + # consumed by FOR_ITER rather than by the yield, so the half-built + # IfExp has to be pulled back out of the element slot. + (y for x in range(4) for y in (range(x) if x % 2 == 0 else range(x, x + 2))), + (y for x in range(4) for y in ([x] if x else [0])), + (y for x in range(4) for y in (range(x) if x > 1 else range(1))), + # An *empty* list literal as an arm is misread: BUILD_LIST(0) is also + # how an inlined list comprehension starts, and nothing later + # disambiguates the two here. + pytest.param( + (y for x in range(4) for y in ([x] if x else [])), + marks=pytest.mark.xfail( + strict=True, + reason="an empty list literal is indistinguishable from the start of a list comprehension", + ), + ), + ((x, y) for x in range(3) for y in ([0] if x % 2 else [1, 2])), + # ... with a filter on the inner loop, and nested two deep + (y for x in range(4) for y in (range(x) if x % 2 == 0 else [9]) if y > 0), + pytest.param( + ( + z + for x in range(3) + for y in (range(x) if x else [0]) + for z in ([y] if y else [7]) + ), + marks=pytest.mark.xfail( + strict=True, + reason="two conditional iterables in one comprehension leave paths that do not pairwise merge", + ), + ), + # ... and one whose arms are comprehensions of different kinds + (y for x in range(3) for y in ([i for i in range(x)] if x else {8})), + ], +) +def test_conditional_expression_as_iterable(genexpr): + """Test a conditional expression in the iterable position of a for clause.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # An always-false filter lets the compiler drop the body entirely, so no + # element expression survives in the bytecode. The loop still runs. + (x for x in range(6) if False), + (x for x in range(6) if x and False), + (y for x in range(6) if False and (y := x)), # noqa: F821 + ([x] for x in range(4) if False), + ((x, y) for x in range(4) for y in range(3) if False), + (x for x in range(6) if False if x > 1), + ({x for x in range(3)} for _ in range(2) if False), + ], +) +def test_unreachable_comprehension_body(genexpr): + """A comprehension whose body the compiler proved unreachable yields nothing.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + assert materialize(compile_and_eval(ast_node)) == [] + + +def test_outermost_iterable_strict_zip_stays_strict(): + """A ragged strict zip raises; the reconstruction must raise there too.""" + genexpr = (a + b for a, b in zip([1, 2, 3], [4, 5], strict=True)) + reconstructed = disassemble(genexpr) + assert ast.unparse(reconstructed) == ( + "(a + b for a, b in zip([1, 2, 3], [4, 5], strict=True))" + ) + + with pytest.raises(ValueError): + materialize(compile_and_eval(reconstructed)) + with pytest.raises(ValueError): + materialize(genexpr) + + +def test_outermost_iterable_lax_zip_stays_lax(): + """A lax zip stops at the shortest iterable and must not become strict.""" + genexpr = (a + b for a, b in zip([1, 2, 3], [4, 5])) + reconstructed = disassemble(genexpr) + assert ast.unparse(reconstructed) == "(a + b for a, b in zip([1, 2, 3], [4, 5]))" + assert materialize(compile_and_eval(reconstructed)) == [5, 7] + assert materialize(genexpr) == [5, 7] + + +def test_outermost_iterable_partially_consumed_strict_zip(): + """Strictness survives even once the zip has been partly consumed.""" + zipped = zip([1, 2, 3], [4, 5], strict=True) + next(zipped) + + genexpr = (a + b for a, b in zipped) + reconstructed = disassemble(genexpr) # must precede consuming `genexpr` + assert "strict=True" in ast.unparse(reconstructed) + with pytest.raises(ValueError): + materialize(compile_and_eval(reconstructed)) + + +@pytest.mark.parametrize( + "genexpr", + [ + # "dict_item" was once an internal marker in the first slot of a tuple, + # which is a string a user's data is perfectly entitled to hold. + (x for x in (("dict_item", 1, 2),)), + (x for x in [("dict_item", "key", "value")]), + (x for x in (("dict_item",), ("dict_item", 1), ("dict_item", 1, 2, 3))), + (("dict_item", x) for x in range(2)), + (("dict_item", x, x + 1) for x in range(2)), + (x for x in {("dict_item", 1, 2): "v"}.items()), + ], +) +def test_tuples_starting_with_dict_item(genexpr): + """No element of a user's tuple is an internal marker to be stripped.""" + assert_ast_equivalent(genexpr, disassemble(genexpr)) + + +def test_outermost_iterable_partially_consumed_adaptor(): + """A consumed prefix is reflected in the adaptor's inner iterators.""" + zipped = zip([1, 2, 3], [4, 5, 6]) + next(zipped) + + genexpr = (a + b for a, b in zipped) + reconstructed = disassemble(genexpr) # must precede consuming `genexpr` + assert materialize(genexpr) == [7, 9] + assert materialize(compile_and_eval(reconstructed)) == [7, 9] + + +def test_outermost_iterable_partially_consumed(): + """Only the *unconsumed* remainder of the outermost iterator belongs in the AST.""" + iterator = iter([10, 20, 30, 40]) + next(iterator) + next(iterator) + + genexpr = (x + 1 for x in iterator) + assert ast.unparse(disassemble(genexpr)) == "(x + 1 for x in [30, 40])" + assert materialize(genexpr) == [31, 41] + + +def test_outermost_iterable_partially_consumed_str(): + iterator = iter("hello") + next(iterator) + + genexpr = (c for c in iterator) + assert ast.unparse(disassemble(genexpr)) == "(c for c in 'ello')" + assert materialize(genexpr) == ["e", "l", "l", "o"] + + +# ============================================================================ +# BINARY OPERATORS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Bitwise operators, which BINARY_OP folds in with the arithmetic ones + (x & 3 for x in range(8)), + (x | 3 for x in range(8)), + (x ^ 3 for x in range(8)), + (x << 2 for x in range(4)), + (x >> 1 for x in range(8)), + (~x & 7 for x in range(8)), + # Mixed precedence across the whole operator table + (x & 1 | x >> 2 ^ 3 for x in range(8)), + ((x | 1) & (x ^ 2) for x in range(8)), + (x + 1 & x - 1 for x in range(8)), + (x * 2 % 5 // 2 for x in range(8)), + (x**2 - x // 2 + x % 3 for x in range(1, 8)), + # Operators on non-numeric operands + (s + "!" for s in ["a", "b"]), + (s * 2 for s in ["a", "b"]), + (t + (9,) for t in [(1,), (2,)]), + (frozenset({x}) | frozenset({9}) for x in range(3)), + (frozenset({x, 1}) & frozenset({1}) for x in range(3)), + (frozenset({x, 1}) ^ frozenset({1}) for x in range(3)), + ({"a": x} | {"b": 0} for x in range(3)), + ], +) +def test_binary_operators(genexpr): + """Test reconstruction of the full BINARY_OP table.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +def test_matmul_operator(): + """BINARY_OP argument 4 is `@`, which no built-in type implements. + + The comprehension is disassembled but never evaluated, so this checks the + reconstructed source rather than the reconstructed values. + """ + genexpr = (a @ b for a in [1, 2]) # noqa: F821 + assert ast.unparse(disassemble(genexpr)) == "(a @ b for a in (1, 2))" + + +# ============================================================================ +# KEYWORD ARGUMENTS AT CALL SITES +# +# Python 3.12 compiles these as KW_NAMES followed by CALL; 3.13 replaced the +# pair with a single CALL_KW instruction. +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr,globals_dict", + [ + ((dict(a=x) for x in range(3)), {}), + ((dict(a=x, b=x * 2) for x in range(3)), {}), + ((dict(a=x if x > 0 else 0, b=x * 2 if x < 5 else x) for x in range(8)), {}), + # Mixed positional and keyword arguments + ((sorted([3, x], reverse=True) for x in range(3)), {}), + ((sorted([x, 1], key=abs) for x in range(3)), {}), + ((sorted([x, 1], key=abs, reverse=True) for x in range(3)), {}), + ((int(str(x), base=8) for x in range(8)), {}), + ((round(x / 3, ndigits=2) for x in range(5)), {}), + # Keyword arguments on a method call + (("a,b".split(sep=",") for x in range(2)), {}), + (("a-b".replace("-", "+") for x in range(2)), {}), + # Nested calls, each with keywords + ((dict(a=dict(b=x)) for x in range(3)), {}), + ((sorted(sorted([x, 1]), reverse=True) for x in range(3)), {}), + # Keyword arguments to a user-supplied callable + ((f(x, scale=2) for x in range(3)), {"f": lambda v, scale=1: v * scale}), # type: ignore # noqa: F821 + ], +) +def test_keyword_arguments(genexpr, globals_dict): + """Test reconstruction of calls with keyword arguments.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node, globals_dict) + + +@pytest.mark.parametrize( + "genexpr", + [ + # Starred positional arguments + (max(*[x, 1]) for x in range(3)), + (max(1, *[x, 2]) for x in range(3)), + (max(*[x, 1], *[2, 3]) for x in range(3)), + (max(*(x, 1)) for x in range(3)), + (sum([x, 1], *[0]) for x in range(3)), + # Double-starred keyword arguments + (dict(**{"a": x}) for x in range(3)), + (dict(a=x, **{"b": 1}) for x in range(3)), + (dict(**{"a": x}, **{"b": 2}) for x in range(3)), + (sorted([x, 1], **{"reverse": True}) for x in range(3)), # type: ignore[call-overload] + # Both at once + (max(*[[x, 1]], **{"default": 0}) for x in range(3)), # type: ignore[call-overload] + # On a method call, where the callable comes with a `self` + ("-".join(*[[str(x), "z"]]) for x in range(3)), + # Unpacking a comprehension, and unpacking into a nested call + (max(*[y for y in range(x + 2)]) for x in range(3)), + (max(*[abs(y) for y in range(-x - 1, 1)]) for x in range(3)), + (dict(**{str(k): k for k in range(x + 1)}) for x in range(3)), + ], +) +def test_star_argument_calls(genexpr): + """Test reconstruction of `*args`/`**kwargs` call sites (CALL_FUNCTION_EX). + + The arguments arrive already collected into a sequence and a mapping, so the + reconstruction spells every argument as unpacked; that evaluates identically + even where the source passed some of them plainly. + """ + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# LAMBDAS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + ((lambda y: y * 2)(x) for x in range(5)), + ((lambda y, z: y + z)(x, x) for x in range(5)), + ((lambda y: y)(x) for x in range(5)), + # Closing over the loop variable, and over an enclosing comprehension + ((lambda y: y + x)(x) for x in range(5)), + ((lambda: x)() for x in range(5)), + (((lambda y: lambda z: z + y)(x))(1) for x in range(5)), + # Lambdas whose body is itself a comprehension + ((lambda: (y for y in range(x)))() for x in range(3)), + ((lambda: [y for y in range(x)])() for x in range(3)), + ((lambda n: sum(y for y in range(n)))(x) for x in range(3)), + # Lambdas as arguments to other calls + (sorted([x, 1], key=lambda v: -v) for x in range(3)), + (list(map(lambda v: v * 2, [x, 1])) for x in range(3)), + # Conditional expressions inside a lambda body + ((lambda y: y if y % 2 else -y)(x) for x in range(5)), + ((lambda y: (y if y > 1 else 0) + 1)(x) for x in range(5)), + ], +) +def test_lambdas(genexpr): + """Test reconstruction of lambdas appearing inside comprehensions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # Positional defaults, which attach to the *trailing* parameters + ((lambda y, z=2: y * z)(x) for x in range(3)), # type: ignore[assignment] + ((lambda y=1: y)() for x in range(3)), # type: ignore[assignment] + ((lambda y, z=2, w=3: y * z * w)(x) for x in range(3)), # type: ignore[assignment] + ((lambda y, z=2: y * z)(x, 5) for x in range(3)), + # Keyword-only parameters, with and without defaults + ((lambda y, *, z=1: y + z)(x) for x in range(3)), # type: ignore[assignment] + ((lambda y, *, z=1, w=2: y + z + w)(x) for x in range(3)), # type: ignore[assignment] + ((lambda y, *, z: y + z)(x, z=4) for x in range(3)), + # Positional-only parameters + ((lambda y, /, z=2: y * z)(x) for x in range(3)), # type: ignore[assignment] + # *args and **kwargs + ((lambda *a: sum(a))(x, x) for x in range(3)), + ((lambda **k: sum(k.values()))(a=x) for x in range(3)), + ((lambda *a, **k: len(a) + len(k))(x, b=1) for x in range(3)), + ((lambda y, *a: y + len(a))(x, 1, 2) for x in range(3)), + ( + (lambda y, *a, z=3, **k: y + len(a) + z + len(k))(x, 1, w=2) + for x in range(3) + ), + # Defaults that are themselves non-trivial expressions + ((lambda y, z=(1, 2): y + len(z))(x) for x in range(3)), # type: ignore[assignment] + ((lambda y, z=[1]: y + len(z))(x) for x in range(3)), # type: ignore[assignment] + ], +) +def test_lambda_default_and_variadic_arguments(genexpr): + """Test reconstruction of lambda defaults and variadic parameters.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # A lambda reached as a live object -- through the outermost iterable -- + # rather than built inside the comprehension. Its defaults are on the + # function, not in the code object it was compiled to. + (fn() for fn in [lambda value=1: value]), + (fn(2) for fn in [lambda a, b=10: a * b]), + (fn(2) for fn in (lambda a, b=10: a * b,)), + (fn(1) for fn in [lambda a, *, k=5: a + k]), + (fn(1) for fn in [lambda a, /, b=2, *, k=5: a * b + k]), + (fn() for fn in [lambda x=(1, 2): sum(x)]), + (fn() for fn in [lambda x=[1, 2]: len(x)]), + (fn(1, 2) for fn in [lambda a, b=0, *rest, k=5, **kw: a + b + k + len(rest)]), + # ... several of them, and one with no defaults alongside + (fn() for fn in [lambda: 0, lambda v=1: v, lambda v=2: v]), # type: ignore[misc] + # ... and one passed through map rather than iterated directly + (fn(3) for fn in map(lambda f: f, [lambda a, b=4: a * b])), + ], +) +def test_lambda_object_defaults(genexpr): + """A lambda arriving as a live object keeps the defaults attached to it.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# ASSIGNMENT EXPRESSIONS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # A walrus binds in the *enclosing* scope, so at module level it + # compiles to COPY + STORE_GLOBAL rather than to a local. + (y for x in range(4) if (y := x * 2) > 1), # noqa: F821 + ((y := x) + 1 for x in range(3)), # noqa: F821 + (y * y for x in range(4) if (y := x + 1) > 2), # noqa: F821 + ((y := x) if x > 1 else -1 for x in range(4)), # noqa: F821 + # Bound in one clause and read in a later one + ((y, z) for x in range(3) if (y := x + 1) for z in range(y)), # noqa: F821 + # Inside a nested comprehension, and inside a lambda + ([(z := w) + z for w in range(x)] for x in range(4)), # noqa: F821 + ((lambda n: [(z := w) + z for w in range(n)])(x) for x in range(4)), # noqa: F821 + # Combined with a short-circuiting filter. The `or` must not re-evaluate + # the assignment, which is why the disjunction is absorbed. + (y for x in range(6) if (y := x * 2) > 6 or y == 0), # noqa: F821 + ], +) +def test_assignment_expressions(genexpr): + """A walrus in a comprehension binds in the *enclosing* scope (STORE_GLOBAL here).""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# COMPREHENSIONS NESTED IN EACH SYNTACTIC POSITION +# +# A comprehension can appear in the element, in the iterable, and inside a +# filter, and each of the four comprehension kinds can nest inside any other. +# The filter position is the interesting one: filters are reconstructed from +# control flow, so a comprehension inside a filter has to survive being treated +# as part of a boolean condition. +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # ... in filter position + (x for x in range(5) if any(y > 2 for y in range(x))), + (x for x in range(5) if all(y < 3 for y in range(x))), + (x for x in range(5) if [y for y in range(x)]), + (x for x in range(5) if {y for y in range(x)}), + (x for x in range(5) if {y: y for y in range(x)}), + (x for x in range(6) if len([y for y in range(x) if y % 2]) > 1), + (x for x in range(5) if any(y for y in range(x) if y % 2)), + ( + x + for x in range(6) + if all(y < x for y in range(2)) and any(z > 1 for z in range(x)) + ), # noqa: E501 # fmt: skip + # ... in filter position, inside a short-circuiting condition + (x for x in range(6) if sum(y for y in range(x)) > 3 or x == 0), + (x for x in range(6) if x == 0 or any(y > 2 for y in range(x))), + (x for x in range(6) if x == 0 or len([y for y in range(x)]) > 2), + ( + x + for x in range(6) + if any(y > 1 for y in range(x)) or all(z < 2 for z in range(x)) + ), # noqa: E501 # fmt: skip + # ... in iterable position + (x for x in [y for y in range(5) if y % 2]), + (x for x in {y for y in range(5) if y % 2}), + (x for x in {y: y for y in range(3)}), + (x for x in (y for y in range(5) if y > 1 or y == 0)), + (x for x in [y for y in [z for z in range(4)] if y % 2]), + (x for x in [y for y in range(4)] if x > 1 or x == 0), + ((a, b) for a in range(3) for b in [c for c in range(a)]), + # ... in element position + ([y for y in range(x) if y % 2 or y == 0] for x in range(4)), + ({y for y in range(x) if y > 1} for x in range(5)), + ({y: [z for z in range(y)] for y in range(x)} for x in range(4)), + (sum(y for y in range(x) if y % 2) for x in range(5)), + ([(z for z in range(y)) for y in range(x)] for x in range(3)), + # ... in several positions at once, with different kinds + ( + [y for y in {z for z in range(x)}] + for x in range(4) + if any(w > 1 for w in range(x)) + ), # noqa: E501 # fmt: skip + ({k: [v for v in range(k)] for k in {j for j in range(x)}} for x in range(4)), + ([y for y in range(x) if y or y == 0] for x in range(5) if x > 1 or x == 0), + ( + (y for y in range(x) if y % 2 or y == 0) + for x in (z for z in range(4)) + if x < 3 or x == 3 + ), # noqa: E501 # fmt: skip + # ... nesting the four kinds inside one another + ({k: {v for v in range(k)} for k in [j for j in range(x)]} for x in range(4)), + ([{y: y} for y in range(x)] for x in range(4)), + ({(y, y * 2) for y in range(x)} for x in range(4)), + ({y: y for y in range(x) if y % 2 or y == 0} for x in range(5)), + ({y for y in range(x) if y % 2 or y == 0} for x in range(5)), + ([y for y in range(x) if 1 < y < 4] for x in range(6)), + ], +) +def test_comprehensions_in_every_position(genexpr): + """Test comprehensions nested in the element, the iterable and the filter.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # A lambda body is a separate code object with no loop of its own, so + # every branch inside one is a conditional expression rather than a + # filter -- including the branches of a comprehension nested in it. + ((lambda n: [y for y in range(n) if y % 2])(x) for x in range(4)), + ((lambda n: {y for y in range(n)})(x) for x in range(4)), + ((lambda n: {y: y**2 for y in range(n)})(x) for x in range(4)), + ((lambda n: sum(y for y in range(n)))(x) for x in range(4)), + ((lambda n: [y for y in range(n) if y > 1 or y == 0])(x) for x in range(5)), + ((lambda n: (y for y in range(n) if y % 2 or y == 0))(x) for x in range(4)), + ((lambda n: [y if y > 1 else -y for y in range(n)])(x) for x in range(4)), + ((lambda n: [y for y in range(n) if 1 < y < 3])(x) for x in range(5)), + # Lambdas nested in lambdas, and lambdas inside the comprehension body + ((lambda n: (lambda m: [y for y in range(m)])(n))(x) for x in range(3)), + ([(lambda v: v * 2)(y) for y in range(x)] for x in range(4)), + (list(map(lambda n: [y for y in range(n)], range(x))) for x in range(3)), + ([(lambda v: v if v > 1 else -v)(y) for y in range(x)] for x in range(4)), + ], +) +def test_comprehensions_inside_lambdas(genexpr): + """Test comprehensions nested inside lambda bodies.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# CLOSURES +# +# A comprehension written inside a function reads the function's locals out of +# closure cells, not out of globals. Those cells belong to a scope the +# reconstruction does not reproduce, so leaving a free variable as a bare name +# would silently turn it into a global lookup -- picking up a different value, +# or none at all. The captured value is written into the tree instead. Each +# test below evaluates the reconstruction in a namespace that binds the same +# name to something else, so a lookup that leaked out would be caught. +# ============================================================================ + + +def test_closure_shadowed_by_global(): + def make(): + value = 1 + return (value for _ in range(1)) + + genexpr = make() + reconstructed = disassemble(genexpr) + assert ast.unparse(reconstructed) == "(1 for _ in range(0, 1, 1))" + assert materialize(genexpr) == [1] + assert materialize(compile_and_eval(reconstructed, {"value": 2})) == [1] + + +@pytest.mark.parametrize( + "make,shadow,expected", + [ + # In the element expression, the filter, and an inner iterable + (lambda: (lambda n: (x * n for x in range(4)))(3), {"n": 100}, [0, 3, 6, 9]), + ( + lambda: (lambda t: (x for x in range(6) if x > t))(3), + {"t": -1}, + [4, 5], + ), + ( + lambda: (lambda k: (y for x in range(2) for y in range(k)))(2), + {"k": 5}, + [0, 1, 0, 1], + ), + # A captured container, indexed and iterated + ( + lambda: (lambda d: (d[i] for i in range(2)))([10, 20]), + {"d": [0, 0]}, + [10, 20], + ), + ( + lambda: (lambda d: (v for v in d))({"a": 1, "b": 2}), + {"d": {}}, + ["a", "b"], + ), + # Captured by a lambda nested inside the comprehension + ( + lambda: (lambda n: ((lambda y: y * n)(x) for x in range(4)))(3), + {"n": 100}, + [0, 3, 6, 9], + ), + ( + lambda: (lambda n: ((lambda: (lambda: n)())() for _ in range(2)))(7), + {"n": 0}, + [7, 7], + ), + # Captured by a comprehension nested inside the comprehension + ( + lambda: (lambda n: ([y * n for y in range(x)] for x in range(3)))(2), + {"n": 100}, + [[], [0], [0, 2]], + ), + ( + lambda: (lambda n: (sum(y for y in range(x) if y < n) for x in range(4)))( + 2 + ), + {"n": 100}, + [0, 0, 1, 1], + ), + # Two free variables at once + ( + lambda: (lambda a, b: (x * a + b for x in range(3)))(2, 1), + {"a": 0, "b": 0}, + [1, 3, 5], + ), + ], +) +def test_closure_free_variables(make, shadow, expected): + genexpr = make() + reconstructed = disassemble(genexpr) + assert materialize(genexpr) == expected + assert materialize(compile_and_eval(reconstructed, dict(shadow))) == expected + # The name must not survive anywhere in the tree, or the shadowing binding + # above would have been the one that answered. + assert not ( + {node.id for node in ast.walk(reconstructed) if isinstance(node, ast.Name)} + & set(shadow) + ) + + +def test_closure_target_captured_by_nested_lambda_stays_a_name(): + """A cell this comprehension *creates* is bound in the reconstruction too.""" + genexpr = ((lambda: x)() for x in range(3)) + reconstructed = disassemble(genexpr) + assert materialize(genexpr) == [0, 1, 2] + assert materialize(compile_and_eval(reconstructed, {"x": 99})) == [0, 1, 2] + + +def test_closure_shadowed_by_a_nested_comprehension_target(): + """Only the free `n` is the captured one; the inner comprehension rebinds it.""" + + def make(): + n = 5 + return (n + sum(n for n in range(x)) for x in range(3)) + + genexpr = make() + reconstructed = disassemble(genexpr) + assert materialize(genexpr) == [5, 5, 6] + assert materialize(compile_and_eval(reconstructed, {"n": 100, "sum": sum})) == [ + 5, + 5, + 6, + ] + + +def test_closure_shadowed_by_a_nested_lambda_parameter(): + """Only the free `n` is the captured one; the lambda's parameter is its own.""" + + def make(): + n = 5 + return ((lambda n: n * 2)(x) + n for x in range(3)) + + genexpr = make() + reconstructed = disassemble(genexpr) + assert materialize(genexpr) == [5, 7, 9] + assert materialize(compile_and_eval(reconstructed, {"n": 100})) == [5, 7, 9] + + +def test_closure_inside_a_live_lambda(): + """A lambda reached as an object closes over cells of its own.""" + + def make(): + n = 3 + return (fn(2) for fn in [lambda a, b=10: a * n + b]) + + genexpr = make() + reconstructed = disassemble(genexpr) + assert materialize(genexpr) == [16] + assert materialize(compile_and_eval(reconstructed, {"n": 100})) == [16] + + +def test_closure_value_that_cannot_be_represented(): + """A capture with no AST spelling is refused rather than quietly dropped.""" + + class Opaque: + pass + + def make(): + obj = Opaque() + return (obj for _ in range(1)) + + with pytest.raises(TypeError, match="captured in free variable 'obj'"): + disassemble(make()) + + +def test_closure_captured_iterator_is_refused(): + """An iterator's remaining elements are not the iterator, so it is refused.""" + + def make(): + flags = iter([True, False, True, False]) + return (x for x in range(4) if next(flags)) + + with pytest.raises(TypeError, match="captured in free variable 'flags'"): + disassemble(make()) + + +# ============================================================================ +# STATEFUL EXPRESSIONS +# +# What comes back is syntax, so evaluating it runs the comprehension a second +# time. A filter or element expression that depends on state that has since +# moved on answers differently then -- faithfully reconstructed, but no longer +# in agreement with the generator it came from. +# ============================================================================ + + +_FLAGS = iter([True, False, True, False, False, False, True, False]) + + +def test_stateful_filter_is_re_evaluated(): + genexpr = (x for x in range(4) if next(_FLAGS)) + reconstructed = disassemble(genexpr) + assert ast.unparse(reconstructed) == "(x for x in range(0, 4, 1) if next(_FLAGS))" + + # The reconstruction is the same comprehension, but `_FLAGS` has advanced by + # the time it runs, so the two do not agree element for element. + assert materialize(genexpr) == [0, 2] + assert materialize( + compile_and_eval(reconstructed, {"_FLAGS": _FLAGS, "next": next}) + ) == [2] + + +@pytest.mark.parametrize( + "genexpr", + [ + # Short-circuiting conditions nested in one another + (x for x in range(20) if (x > 2 or x < 1) and (x < 10 or x > 15)), + (x for x in range(20) if ((x > 2 and x < 5) or (x > 10 and x < 15)) or x == 0), + (x for x in range(30) if not (x % 2 == 0 or x % 3 == 0)), + (x for x in range(30) if not (not (x > 5) or not (x < 20))), + (x for x in range(40) if (x > 5 and x < 35) and (x % 3 == 0 or x % 5 == 0)), + ( + x + for x in range(40) + if (x < 5 and x % 2 == 0) or (10 < x < 15) or (x > 35 and x % 3 == 0) + ), # noqa: E501 # fmt: skip + # Short-circuiting conditions spanning several generators + ( + (x, y) + for x in range(6) + if x < 2 or x > 4 + for y in range(6) + if y < 1 or y > 4 + ), + ( + (x, y) + for x in range(5) + if x % 2 == 0 or x == 1 + for y in range(x) + if y > 0 and y < 3 + ), # noqa: E501 # fmt: skip + # Conditional expressions and filters that are both lazy + ((x if (x > 2 or x < 1) else -x) for x in range(10) if x % 2 == 0 or x == 1), + (x for x in range(20) if (x if x > 5 else not x) or x == 3), + ( + (x if x > 5 or x < 2 else (0 if x % 2 == 0 or x == 3 else 1)) + for x in range(12) + ), # noqa: E501 # fmt: skip + # Chained comparisons combined with lazy operators + (x for x in range(30) if 5 < x < 15 or 20 < x < 25), + (x for x in range(30) if 5 < x < 15 and (x % 2 == 0 or x % 3 == 0)), + ((x if 5 < x < 15 else 0) for x in range(20) if 2 < x < 18), + # Lazy conditions inside a nested comprehension, and around it + ([y for y in range(x) if y > 1 or y == 0] for x in range(5) if x > 2 or x == 0), + ((y for y in range(x) if y % 2 or y == 0) for x in range(4) if x < 3 or x == 3), + ], +) +def test_nested_lazy_conditions(genexpr): + """Test short-circuiting conditions nested inside one another.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# STRUCTURAL STRESS CASES +# ============================================================================ + +# These two must stay on one line: on Python 3.12 `dis` mis-reports jumps for +# multiline comprehensions, which test_multiline_comprehensions covers directly. +_STRESS_MANY_FILTERS = ((x, y) for x in range(10) if x % 2 == 0 if x > 2 for y in range(10) if y % 3 == 0 if y < x) # fmt: skip +_STRESS_NESTED_TERNARY = ([y if y > 1 else -y for y in range(x)] for x in range(4) if (x if x % 2 == 1 else x % 2 == 0)) # fmt: skip + + +@pytest.mark.parametrize( + "genexpr", + [ + # Deep loop nesting + ( + a + b + c + d + e + for a in range(2) + for b in range(2) + for c in range(2) + for d in range(2) + for e in range(2) + ), + ( + (a, b, c, d) + for a in range(2) + for b in range(a + 1) + for c in range(b + 1) + for d in range(c + 1) + ), + # Many filters spread over many loops. Kept on one line: on Python 3.12 + # `dis` mis-reports jumps for multiline comprehensions, which is covered + # separately by test_multiline_comprehensions below. + (x for x in range(50) if x > 5 if x < 40 if x % 2 == 0 if x % 3 == 0), + _STRESS_MANY_FILTERS, + # Deep comprehension nesting + (((z for z in range(y)) for y in range(x)) for x in range(3)), + ([[z for z in range(y)] for y in range(x)] for x in range(3)), + ({y: [z for z in range(y)] for y in range(x)} for x in range(3)), + # Structured literals in the element position + ({x, x + 1} for x in range(3)), + ({x: x + 1} for x in range(3)), + (((x, x), x) for x in range(3)), + ([x, [x, [x]]] for x in range(3)), + ({"k": [x, {"j": (x,)}]} for x in range(3)), + # Ternaries interleaved with nesting + _STRESS_NESTED_TERNARY, + (((y if y else -1) for y in range(x)) for x in range(3)), + ], +) +def test_structural_stress(genexpr): + """Test reconstruction of deeply nested and heavily filtered comprehensions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# MULTILINE COMPREHENSIONS +# +# On Python 3.12 `dis` reports a different jump layout for a filter whose source +# spans several lines. The reconstruction used to come out negated, because the +# filter/conditional distinction was drawn from the *local* instruction order. +# Classifying branches from the control-flow graph instead is insensitive to +# that, so these now reconstruct identically on 3.12 and 3.13+. +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + ( + x + for x in range(5) # comment to avoid reformatting + if x > 1 + ), + ( + x + for x in range(10) # comment to avoid reformatting + if x % 2 == 0 + if x > 2 + ), + ( + (x, y) + for x in range(10) + if x % 2 == 0 + if x > 2 + for y in range(10) + if y % 3 == 0 + if y < x + ), + ( + [y if y > 1 else -y for y in range(x)] + for x in range(4) + if (x if x % 2 == 1 else x % 2 == 0) + ), + ], +) +def test_multiline_comprehensions(genexpr): + """Filters in multiline comprehensions are mis-disassembled on Python 3.12.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +def test_multiline_comprehensions_same_on_one_line(): + """The same expressions reconstruct correctly when written on one line.""" + one_line = (x for x in range(10) if x % 2 == 0 if x > 2) + assert_ast_equivalent(one_line, disassemble(one_line)) + + +# ============================================================================ +# HELPER FUNCTION TESTS +# ============================================================================ + + +@pytest.mark.parametrize( + "value,expected_str", + [ + # AST nodes should be returned as-is + (ast.Name(id="x", ctx=ast.Load()), "x"), + (ast.Constant(value=42), "42"), + (ast.List(elts=[], ctx=ast.Load()), "[]"), + ( + ast.BinOp( + left=ast.Constant(value=1), op=ast.Add(), right=ast.Constant(value=2) + ), + "1 + 2", + ), + # Constants should become ast.Constant nodes + (42, "42"), + (3.14, "3.14"), + (-42, "-42"), + (-3.14, "-3.14"), + ("hello", "'hello'"), + ("", "''"), + (b"bytes", "b'bytes'"), + (b"", "b''"), + (True, "True"), + (False, "False"), + (None, "None"), + # Complex numbers + (1 + 2j, "(1+2j)"), + (0 + 1j, "1j"), + (3 + 0j, "(3+0j)"), + (-1 - 2j, "(-1-2j)"), + # Tuples should become ast.Tuple nodes + ((), "()"), + ((1,), "(1,)"), + ((1, 2), "(1, 2)"), + (("a", "b", "c"), "('a', 'b', 'c')"), + # A tuple is a tuple: no element of one is a marker to be stripped + (("dict_item", "key", "value"), "('dict_item', 'key', 'value')"), + (("dict_item", 42, "answer"), "('dict_item', 42, 'answer')"), + # Nested tuples + ((1, (2, 3)), "(1, (2, 3))"), + (((1, 2), (3, 4)), "((1, 2), (3, 4))"), + ((1, 2, (3, (4, 5))), "(1, 2, (3, (4, 5)))"), + # Lists should become ast.List nodes + ([1, 2, 3], "[1, 2, 3]"), + (["hello", "world"], "['hello', 'world']"), + ([True, False, None], "[True, False, None]"), + # Nested lists + ([[1, 2], [3, 4]], "[[1, 2], [3, 4]]"), + ([1, [2, [3, 4]], 5], "[1, [2, [3, 4]], 5]"), + # Mixed nested structures + ([(1, 2), (3, 4)], "[(1, 2), (3, 4)]"), + (([1, 2], [3, 4]), "([1, 2], [3, 4])"), + # Dicts should become ast.Dict nodes + ({"a": 1}, "{'a': 1}"), + ({"x": 10, "y": 20}, "{'x': 10, 'y': 20}"), + ({1: "one", 2: "two"}, "{1: 'one', 2: 'two'}"), + # Nested dicts + ({"a": {"b": 1}}, "{'a': {'b': 1}}"), + ( + {"nums": [1, 2, 3], "strs": ["a", "b"]}, + "{'nums': [1, 2, 3], 'strs': ['a', 'b']}", + ), + # Range objects + (range(5), "range(0, 5, 1)"), + (range(1, 10), "range(1, 10, 1)"), + (range(0, 10, 2), "range(0, 10, 2)"), + (range(10, 0, -1), "range(10, 0, -1)"), + (range(-5, 5), "range(-5, 5, 1)"), + # Empty collections + ([], "[]"), + ((), "()"), + ({}, "{}"), + # Complex nested structures + ([1, [2, 3], 4], "[1, [2, 3], 4]"), + ({"a": [1, 2], "b": {"c": 3}}, "{'a': [1, 2], 'b': {'c': 3}}"), + ([(1, {"a": [2, 3]}), ({"b": 4}, 5)], "[(1, {'a': [2, 3]}), ({'b': 4}, 5)]"), + # Edge cases with special values + ([None, True, False, 0, ""], "[None, True, False, 0, '']"), + ( + {"": "empty", None: "none", 0: "zero"}, + "{'': 'empty', None: 'none', 0: 'zero'}", + ), + # Large numbers + (999999999999999999999, "999999999999999999999"), + (1.7976931348623157e308, "1.7976931348623157e+308"), # Close to float max + # Sets - note unparse equivalence may fail for unordered collections + ({1, 2, 3}, "{1, 2, 3}"), + ], +) +def test_ensure_ast(value, expected_str): + """Test that ensure_ast correctly converts various values to AST nodes.""" + + result = ensure_ast(value) + + # Compare the unparsed strings + result_str = ast.unparse(result) + assert result_str == expected_str, ( + f"ensure_ast({repr(value)}) produced '{result_str}', expected '{expected_str}'" + ) + + +def test_error_handling(): + """Test that appropriate errors are raised for unsupported cases.""" + # Test with non-generator input + with pytest.raises(ValueError): + disassemble([1, 2, 3]) # Not a generator + + # Test with consumed generator + gen = (x for x in range(5)) + list(gen) # Consume it + with pytest.raises(ValueError): + disassemble(gen) + + # Test with a generator that has been started but not consumed + gen = (x for x in range(5)) + next(gen) + with pytest.raises(ValueError): + disassemble(gen) + + +def test_comp_lambda_copy(): + """Test that CompLambda is compatible with copy.copy and copy.deepcopy.""" + # Create a test generator expression AST + genexpr_ast = ast.GeneratorExp( + elt=ast.Name(id="x", ctx=ast.Load()), + generators=[ + ast.comprehension( + target=ast.Name(id="x", ctx=ast.Store()), + iter=DummyIterName(), + ifs=[], + is_async=0, + ) + ], + ) + + # Create a CompLambda instance + comp_lambda = CompLambda(genexpr_ast) + + # Test copy.copy + copied = copy.copy(comp_lambda) + assert isinstance(copied, CompLambda) + assert ast.unparse(copied.body) == ast.unparse(comp_lambda.body) + assert copied.body is comp_lambda.body # Shallow copy shares the body + + # Test copy.deepcopy + deep_copied = copy.deepcopy(comp_lambda) + assert isinstance(deep_copied, CompLambda) + assert ast.unparse(deep_copied.body) == ast.unparse(comp_lambda.body) + assert deep_copied.body is not comp_lambda.body # Deep copy creates new body + + # Test that deep copied version works the same way + iterator = ast.Call( + func=ast.Name(id="range", ctx=ast.Load()), + args=[ast.Constant(value=5)], + keywords=[], + ) + + original_result = comp_lambda.inline(iterator) + deep_copied_result = deep_copied.inline(iterator) + + assert ast.unparse(original_result) == ast.unparse(deep_copied_result) + assert type(original_result) == type(deep_copied_result) + + +# ============================================================================ +# AST TRANSFORMER TESTS +# ============================================================================ diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py new file mode 100644 index 000000000..4dffba967 --- /dev/null +++ b/tests/test_ops_monoid.py @@ -0,0 +1,1555 @@ +import functools +import math +import random +import sys +import typing +from collections.abc import Iterable, Mapping + +import pytest +from hypothesis import HealthCheck, given, settings +from hypothesis import strategies as st + +import effectful.handlers.jax.monoid # noqa: F401 +from effectful.ops.monoid import ( + And, + CartesianProduct, + CartesianProductPlus, + EliminateSingletonStreams, + EvaluateIntp, + Factor, + Max, + Min, + Monoid, + MonoidOverMapping, + MonoidOverSequence, + NormalizeIntp, + Or, + PlusAssoc, + PlusConsecutiveDups, + PlusDistr, + PlusEmpty, + PlusOrder, + PlusSingle, + Product, + ReduceDependentRangeMask, + ReduceDisjunctiveDisequalityMask, + ReduceDistributeCartesianProduct, + ReduceEmpty, + ReduceEqualityMaskRange, + ReduceFusion, + ReduceGroundCartesianProduct, + ReduceMaskHoist, + ReducePartial, + ReduceSplit, + ReduceUnfactor, + ReduceUnion, + ReduceWeightedStream, + ReduceWhereEqualityPeel, + ReduceWhereToMasks, + Sum, + Union, + WhereHoist, + distributes_over, + is_commutative, +) +from effectful.ops.semantics import coproduct, evaluate, fvsof, handler +from effectful.ops.syntax import as_dict, ite, range_, syntactic_eq +from effectful.ops.types import NotHandled, Operation, Term +from tests._monoid_helpers import Backend, IntBackend, JaxBackend, syntactic_eq_alpha + +sys.setrecursionlimit(10_000) + + +@pytest.fixture(params=[IntBackend, JaxBackend], ids=["int", "jax"]) +def backend(request) -> Backend: + return request.param() + + +ALL_MONOIDS = [ + pytest.param(Sum, id="Sum"), + pytest.param(Product, id="Product"), + pytest.param(Min, id="Min"), + pytest.param(Max, id="Max"), +] + +COMMUTATIVE = [ + pytest.param(Sum, id="Sum"), + pytest.param(Product, id="Product"), + pytest.param(Min, id="Min"), + pytest.param(Max, id="Max"), +] + +IDEMPOTENT = [ + pytest.param(Min, id="Min"), + pytest.param(Max, id="Max"), +] + +WITH_ZERO = [ + pytest.param(Product, id="Product"), +] + +# Pairs (outer, inner) such that inner distributes over outer — i.e. the lifting +# identity ``outer(inner(body, A), CartesianProduct...) == inner(outer(body, D), ...)`` +# is valid for that semiring pair. +MONOID_PAIRS = [ + pytest.param(o.values[0], i.values[0], id=f"{o.id}-{i.id}") + for o in ALL_MONOIDS + for i in ALL_MONOIDS + if distributes_over( + typing.cast(Monoid, i.values[0]), typing.cast(Monoid, o.values[0]) + ) +] + +COMMUTATIVE_MONOID_PAIRS = [ + v + for v in MONOID_PAIRS + if all(is_commutative(typing.cast(Monoid, vv)) for vv in v.values) +] + + +def test_plus_ite_hoist(backend: Backend): + """Monoid addition distributes into both branches of an ``ite``.""" + cond_lhs, cond_rhs, a, b, c = backend.define_vars( + "cond_lhs", "cond_rhs", "a", "b", "c", ret="scalar" + ) + cond = cond_lhs() == cond_rhs() + lhs = Product.plus(ite(cond, a(), b()), c()) + rhs = ite(cond, Product.plus(a(), c()), Product.plus(b(), c())) + + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=WhereHoist()) + + +def test_reduce_ite_hoist(backend: Backend): + """An ``ite`` with a stream-independent condition hoists.""" + i, out_i, out_j = backend.define_vars("i", "out_i", "out_j", ret="scalar") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") + streams = {i: range(3)} + cond = out_i() == out_j() + + lhs = Sum.reduce(ite(cond, f(i()), g(i())), streams) + rhs = ite(cond, Sum.reduce(f(i()), streams), Sum.reduce(g(i()), streams)) + + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=WhereHoist()) + + +def test_reduce_ite_hoist_dependent_noop(backend: Backend): + """An ``ite`` whose condition uses the reduced stream remains in place.""" + i = backend.define_vars("i", ret="scalar") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") + term = Sum.reduce(ite(i() == 0, f(i()), g(i())), {i: range(3)}) + + backend.check_rewrite(lhs=term, rhs=term, rule=WhereHoist()) + + +def test_reduce_ite_equality_peel(backend: Backend): + """An independent conjunct peels off a stream-selecting equality guard.""" + i, j = backend.define_vars("i", "j", ret="scalar") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") + streams = {j: range(3)} + outer = 0 == i() + inner = 0 == j() + + lhs = Product.reduce(ite(And.plus(inner, outer), f(j()), g(j())), streams) + rhs = ite( + And.plus(outer), + Product.reduce(ite(And.plus(inner), f(j()), g(j())), streams), + Product.reduce(g(j()), streams), + ) + + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceWhereEqualityPeel()) + + +def test_reduce_ite_equality_peel_requires_stream_equality(backend: Backend): + """A mixed guard without a stream equality is left unchanged.""" + j, out_i, i = backend.define_vars("j", "out_i", "i", ret="scalar") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") + term = Product.reduce( + ite(And.plus(j() < 2, out_i() == i()), f(j()), g(j())), + {j: range(3)}, + ) + + backend.check_rewrite(lhs=term, rhs=term, rule=ReduceWhereEqualityPeel()) + + +def test_reduce_disjunctive_disequality_mask(): + backend = IntBackend() + i, j, out_i, out_j, value = backend.define_vars( + "i", "j", "out_i", "out_j", "value", ret="scalar" + ) + streams = {i: range(3), j: range(4)} + + lhs = Sum.reduce( + Sum.mask(value(), Or.plus(i() != out_i(), j() != out_j())), streams + ) + rhs = Sum.plus( + Sum.reduce(Sum.mask(value(), And.plus(i() != out_i())), streams), + Sum.reduce( + Sum.mask(value(), And.plus(i() == out_i(), j() != out_j())), streams + ), + ) + + with handler(ReduceDisjunctiveDisequalityMask()): + actual = evaluate(lhs) + assert syntactic_eq_alpha(actual, rhs) + + +def test_reduce_where_to_masks(): + backend = IntBackend() + """A conjunctive equality where partitions into complementary masks.""" + i, j = backend.define_vars("i", "j", ret="scalar") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") + streams = {i: range(2), j: range(3)} + + lhs = Product.reduce(ite(And.plus(i() == 0, j() == 0), f(i()), g(j())), streams) + rhs = Product.plus( + Product.reduce(Product.mask(f(i()), And.plus(i() == 0, j() == 0)), streams), + Product.reduce(Product.mask(g(j()), Or.plus(i() != 0, j() != 0)), streams), + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceWhereToMasks()) + + +def test_reduce_where_to_masks_requires_stream_equalities(backend: JaxBackend): + """Independent equality conjuncts are left for WhereEqualityPeel.""" + i, out_i, x, y = backend.define_vars("i", "out_i", "x", "y", ret="scalar") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") + term = Product.reduce( + ite(And.plus(i() == out_i(), x() == y()), f(i()), g(i())), + {i: range(2)}, + ) + + with handler(ReduceWhereToMasks()): + actual = evaluate(term) + assert syntactic_eq_alpha(actual, term) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +@given(data=st.data()) +@settings( + max_examples=50, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +def test_associativity(monoid, backend: Backend, data): + a = data.draw(backend.strategy(ret="scalar"))() + b = data.draw(backend.strategy(ret="scalar"))() + c = data.draw(backend.strategy(ret="scalar"))() + with handler(NormalizeIntp): + left = monoid.plus(monoid.plus(a, b), c) + right = monoid.plus(a, monoid.plus(b, c)) + assert syntactic_eq(left, right) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +@given(data=st.data()) +@settings( + max_examples=50, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +def test_identity(monoid, backend: Backend, data): + a = data.draw(backend.strategy(ret="scalar"))() + with handler(NormalizeIntp): + assert backend.eq(monoid.plus(monoid.identity, a), a) + assert backend.eq(monoid.plus(a, monoid.identity), a) + + +@pytest.mark.parametrize("monoid", COMMUTATIVE) +@given(data=st.data()) +@settings( + max_examples=50, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +def test_commutativity(monoid, backend: Backend, data): + a = data.draw(backend.strategy(ret="scalar"))() + b = data.draw(backend.strategy(ret="scalar"))() + with handler(NormalizeIntp): + assert syntactic_eq(monoid.plus(a, b), monoid.plus(b, a)) + + +@pytest.mark.parametrize("monoid", IDEMPOTENT) +@given(data=st.data()) +@settings( + max_examples=50, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +def test_idempotence(monoid, backend: Backend, data): + a = data.draw(backend.strategy(ret="scalar"))() + with handler(NormalizeIntp): + assert backend.eq(monoid.plus(a, a), a) + + +@pytest.mark.parametrize("monoid", WITH_ZERO) +@given(data=st.data()) +@settings( + max_examples=50, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +def test_zero_absorbs(monoid, backend: Backend, data): + a = data.draw(backend.strategy(ret="scalar"))() + with handler(NormalizeIntp): + assert backend.eq(monoid.plus(monoid.zero, a), monoid.zero) + assert backend.eq(monoid.plus(a, monoid.zero), monoid.zero) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_plus_empty(monoid, backend: Backend): + backend.check_rewrite(lhs=monoid.plus(), rhs=monoid.identity, rule=PlusEmpty()) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_plus_single(monoid, backend: Backend): + x = backend.define_vars("x", ret="scalar") + backend.check_rewrite(lhs=monoid.plus(x()), rhs=x(), rule=PlusSingle()) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_plus_identity_right(monoid, backend: Backend): + x = backend.define_vars("x", ret="scalar") + + lhs = monoid.plus(x(), monoid.identity) + rhs = monoid.plus(x()) + + backend.check_rewrite(lhs=lhs, rhs=rhs, rule={}) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_plus_identity_left(monoid, backend: Backend): + x = backend.define_vars("x", ret="scalar") + + lhs = monoid.plus(monoid.identity, x()) + rhs = monoid.plus(x()) + + backend.check_rewrite(lhs=lhs, rhs=rhs, rule={}) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_plus_assoc_right(monoid, backend: Backend): + x, y, z = backend.define_vars("x", "y", "z", ret="scalar") + backend.check_rewrite( + lhs=monoid.plus(x(), monoid.plus(y(), z())), + rhs=monoid.plus(x(), y(), z()), + rule=PlusAssoc(), + ) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_plus_assoc_left(monoid, backend: Backend): + x, y, z = backend.define_vars("x", "y", "z", ret="scalar") + backend.check_rewrite( + lhs=monoid.plus(monoid.plus(x(), y()), z()), + rhs=monoid.plus(x(), y(), z()), + rule=PlusAssoc(), + ) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_plus_sequence(monoid, backend: Backend): + a, b, c, d = backend.define_vars("a", "b", "c", "d", ret="scalar") + backend.check_rewrite( + lhs=monoid.plus((a(), b()), (c(), d())), + rhs=(monoid.plus(a(), c()), monoid.plus(b(), d())), + rule=MonoidOverSequence(), + ) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_plus_mapping(monoid, backend: Backend): + a, b, c, d = backend.define_vars("a", "b", "c", "d", ret="scalar") + + lhs = monoid.plus({0: a(), 1: b()}, {0: c(), 2: d()}) + rhs = {0: monoid.plus(a(), c()), 1: monoid.plus(b()), 2: monoid.plus(d())} + + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=MonoidOverMapping()) + + +def test_plus_distributes_1(backend: Backend): + a, b, c = backend.define_vars("a", "b", "c", ret="scalar") + lhs = Product.plus(c(), Sum.plus(a(), b())) + rhs = Sum.plus(Product.plus(c(), a()), Product.plus(c(), b())) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=coproduct(PlusDistr(), coproduct(PlusSingle(), PlusAssoc())), + ) + + +def test_plus_distributes_2(backend: Backend): + a, b, c = backend.define_vars("a", "b", "c", ret="scalar") + lhs = Product.plus(Sum.plus(a(), b()), c()) + rhs = Sum.plus(Product.plus(a(), c()), Product.plus(b(), c())) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=coproduct(PlusDistr(), coproduct(PlusSingle(), PlusAssoc())), + ) + + +def test_plus_distributes_3(backend: Backend): + a, b, c, d = backend.define_vars("a", "b", "c", "d", ret="scalar") + lhs = Product.plus(Sum.plus(a(), b()), Sum.plus(c(), d())) + rhs = Sum.plus( + Product.plus(a(), c()), + Product.plus(a(), d()), + Product.plus(b(), c()), + Product.plus(b(), d()), + ) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=coproduct(PlusDistr(), coproduct(PlusSingle(), PlusAssoc())), + ) + + +def test_plus_distributes_4(backend: Backend): + a, b, c, d = backend.define_vars("a", "b", "c", "d", ret="scalar") + lhs = Product.plus(Sum.plus(a(), b()), c(), d()) + rhs = Sum.plus(Product.plus(a(), c(), d()), Product.plus(b(), c(), d())) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=coproduct(PlusDistr(), coproduct(PlusSingle(), PlusAssoc())), + ) + + +def test_plus_distributes_constant(backend: Backend): + a, b, c, d, e = backend.define_vars("a", "b", "c", "d", "e", ret="scalar") + lhs = Product.plus(Sum.plus(a(), b()), Sum.plus(c(), d()), e()) + rhs = Sum.plus( + Product.plus(a(), c(), e()), + Product.plus(a(), d(), e()), + Product.plus(b(), c(), e()), + Product.plus(b(), d(), e()), + ) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=coproduct(PlusDistr(), coproduct(PlusSingle(), PlusAssoc())), + ) + + +def test_plus_distributes_multiple(backend: Backend): + a, b, c, d = backend.define_vars("a", "b", "c", "d", ret="scalar") + lhs = Sum.plus(Min.plus(a(), b()), Max.plus(c(), d())) + rhs = Min.plus( + Max.plus(Sum.plus(a(), c()), Sum.plus(a(), d())), + Max.plus(Sum.plus(b(), c()), Sum.plus(b(), d())), + ) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=coproduct(PlusDistr(), coproduct(PlusSingle(), PlusAssoc())), + ) + + +@pytest.mark.parametrize("monoid", IDEMPOTENT) +def test_plus_idempotent_consecutive(monoid, backend: Backend): + """``a, a, b → a, b`` — only consecutive duplicates collapse.""" + a, b = backend.define_vars("a", "b", ret="scalar") + lhs = monoid.plus(a(), a(), b()) + return backend.check_rewrite( + lhs=lhs, rhs=monoid.plus(a(), b()), rule=PlusConsecutiveDups() + ) + + +@pytest.mark.parametrize("monoid", IDEMPOTENT) +def test_plus_idempotent_non_consecutive(monoid, backend: Backend): + a, b = backend.define_vars("a", "b", ret="scalar") + lhs = monoid.plus(a(), b(), a()) + rhs = monoid.plus(a(), b(), a()) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=PlusConsecutiveDups()) + + +@pytest.mark.parametrize("monoid", [Min, Max]) +def test_plus_commutative_idempotent_long(monoid, backend: Backend): + """Long alternation collapses via commutative dedup (Min/Max only).""" + lhs = monoid.plus(0, 1, 0, 1, 1, 0, 0) + rhs = monoid.plus(0, 1) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(PlusOrder(), PlusConsecutiveDups()) + ) + + +@pytest.mark.parametrize("monoid", WITH_ZERO) +def test_plus_zero(monoid, backend: Backend): + a = backend.define_vars("a", ret="scalar") + lhs_right = monoid.plus(a(), monoid.zero) + lhs_left = monoid.plus(monoid.zero, a()) + rhs = monoid.zero + backend.check_rewrite(lhs=lhs_right, rhs=rhs, rule={}) + backend.check_rewrite(lhs=lhs_left, rhs=rhs, rule={}) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_partial_1(monoid, backend: Backend): + x = backend.define_vars("x", ret="scalar") + lhs = monoid.reduce(x(), {x: []}) + rhs = monoid.plus() + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReducePartial()) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_partial_2(monoid, backend: Backend): + x, y = backend.define_vars("x", "y", ret="scalar") + Y = backend.define_vars("Y", ret="stream") + + lhs = monoid.reduce(x(), {y: Y(), x: []}) + rhs = monoid.plus() + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReducePartial()) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_partial_3(monoid, backend: Backend): + x, y, a, b = backend.define_vars("x", "y", "a", "b", ret="scalar") + Y = backend.define_vars("Y", ret="stream") + + lhs = monoid.reduce(x(), {y: Y(), x: [a(), b()]}) + rhs = monoid.plus(monoid.reduce(a(), {y: Y()}), monoid.reduce(b(), {y: Y()})) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReducePartial()) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_partial_4(monoid, backend: Backend): + x, y, a, b = backend.define_vars("x", "y", "a", "b", ret="scalar") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="stream") + + lhs = monoid.reduce(x(), {y: f(x()), x: [a(), b()]}) + rhs = monoid.plus(monoid.reduce(a(), {y: f(a())}), monoid.reduce(b(), {y: f(b())})) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReducePartial()) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_eliminate_singleton_into_sibling(monoid, backend: Backend): + """A length-1 stream substitutes its element into the body *and* into a + sibling stream's definition, then drops out of the nest.""" + x, y, a = backend.define_vars("x", "y", "a", ret="scalar") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="stream") + g = backend.define_vars( + "g", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) + + lhs = monoid.reduce(g(x(), y()), {x: (a(),), y: f(x())}) + rhs = monoid.reduce(g(a(), y()), {y: f(a())}) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=EliminateSingletonStreams()) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_eliminate_singleton_only_stream(monoid, backend: Backend): + """When the length-1 stream is the only stream, reducing over the now-empty + nest yields the substituted body itself (not the monoid identity).""" + x, a = backend.define_vars("x", "a", ret="scalar") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + + lhs = monoid.reduce(f(x()), {x: (a(),)}) + rhs = monoid.reduce(f(a()), {}) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=EliminateSingletonStreams()) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_body_sequence(monoid, backend: Backend): + x = backend.define_vars("x", ret="scalar") + X = backend.define_vars("X", ret="stream") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") + + lhs = monoid.reduce((f(x()), g(x())), {x: X()}) + rhs = (monoid.reduce(f(x()), {x: X()}), monoid.reduce(g(x()), {x: X()})) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=MonoidOverSequence()) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_body_sequence_2(monoid, backend: Backend): + x, y = backend.define_vars("x", "y", ret="scalar") + X, Y = backend.define_vars("X", "Y", ret="stream") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") + + lhs = monoid.reduce((f(x()), g(y())), {x: X(), y: Y()}) + rhs = ( + monoid.reduce(f(x()), {x: X(), y: Y()}), + monoid.reduce(g(y()), {x: X(), y: Y()}), + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=MonoidOverSequence()) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_body_mapping(monoid, backend: Backend): + x = backend.define_vars("x", ret="scalar") + X = backend.define_vars("X", ret="stream") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") + + lhs = monoid.reduce({0: f(x()), 1: g(x())}, {x: X()}) + rhs = { + 0: monoid.reduce(f(x()), {x: X()}), + 1: monoid.reduce(g(x()), {x: X()}), + } + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=MonoidOverMapping()) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_no_streams(monoid, backend: Backend): + a = backend.define_vars("a", ret="scalar") + + lhs = monoid.reduce(a(), {}) + rhs = a() + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceEmpty()) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_reduce(monoid, backend: Backend): + a, b = backend.define_vars("a", "b", ret="scalar") + A, B = backend.define_vars("A", "B", ret="stream") + f = backend.define_vars( + "f", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) + + lhs = monoid.reduce(monoid.reduce(f(a(), b()), {a: A()}), {b: B()}) + rhs = monoid.reduce(f(a(), b()), {a: A(), b: B()}) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceFusion()) + + +@pytest.mark.parametrize("monoid", COMMUTATIVE) +def test_reduce_split_subset(monoid, backend: Backend): + """ReduceSplit confines a stream to the summands that use it: ``a`` is used + only by the first summand, so it is pushed into a reduce over just that + summand (the constant summand keeps its -- now innermost -- ``a`` reduce). + """ + a = backend.define_vars("a", ret="scalar") + A = backend.define_vars("A", ret="stream") + c = backend.define_vars("c", ret="scalar") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + + lhs = monoid.reduce(monoid.plus(f(a()), c()), {a: A()}) + rhs = monoid.plus( + monoid.reduce(f(a()), {a: A()}), + monoid.reduce(c(), {a: A()}), + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceSplit()) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_mask_hoist(monoid): + """A reduce-stream-independent mask condition lifts out of the reduce: + ``M.reduce(M.mask(v, c), {a: A}) == M.mask(M.reduce(v, {a: A}), c)``. + + Pinned to ``IntBackend``: ``Monoid.mask`` evaluates symbolically (via + ``MaskConcrete``) under the ops stack used by ``check_rewrite``; jax-array + masks are exercised separately in ``test_handlers_jax_monoid``. + """ + backend = IntBackend() + a, c, d = backend.define_vars("a", "c", "d", ret="scalar") + A = backend.define_vars("A", ret="stream") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + + cond = c() == d() # independent of the reduced stream `a` + lhs = monoid.reduce(monoid.mask(f(a()), cond), {a: A()}) + rhs = monoid.mask(monoid.reduce(f(a()), {a: A()}), cond) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceMaskHoist()) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_mask_hoist_dependent_noop(monoid): + """The mask does NOT hoist when its condition depends on the reduced + stream (that case is gather territory for ``ReduceEqualityMaskRange``). + """ + backend = IntBackend() + a, c = backend.define_vars("a", "c", ret="scalar") + A = backend.define_vars("A", ret="stream") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + + term = monoid.reduce(monoid.mask(f(a()), a() == c()), {a: A()}) + backend.check_rewrite(lhs=term, rhs=term, rule=ReduceMaskHoist()) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_equality_mask_range_simple(backend: Backend, monoid): + """Best case: a single equality on the reduced range stream becomes a + singleton-stream gather guarded by the corresponding bounds check. + """ + a, c = backend.define_vars("a", "c", ret="scalar") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + + lhs = monoid.reduce(monoid.mask(f(a()), a() == c()), {a: range(3)}) + rhs = monoid.reduce( + monoid.mask( + monoid.mask(f(a()), And.plus(0 <= c(), c() < 3)), + And.plus(), + ), + {a: (c(),)}, + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceEqualityMaskRange()) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_equality_mask_range_residual_conjuncts(backend: Backend, monoid): + """Non-equality conjuncts are preserved inside the gathered singleton + reduce; only the equality on the reduced range stream is discharged. + """ + a, c, d, e = backend.define_vars("a", "c", "d", "e", ret="scalar") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + + lhs = monoid.reduce( + monoid.mask(f(a()), And.plus(d() < e(), a() == c(), c() < e())), + {a: range(4)}, + ) + rhs = monoid.reduce( + monoid.mask( + monoid.mask(f(a()), And.plus(0 <= c(), c() < 4)), + And.plus(d() < e(), c() < e()), + ), + {a: (c(),)}, + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceEqualityMaskRange()) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_equality_mask_range_noncanonical_range_noop(backend: Backend, monoid): + """The rule only handles ``range(0, N, 1)`` streams; for other ranges it + should leave the term unchanged. + """ + a, c = backend.define_vars("a", "c", ret="scalar") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + + term = monoid.reduce(monoid.mask(f(a()), a() == c()), {a: range(1, 4)}) + backend.check_rewrite(lhs=term, rhs=term, rule=ReduceEqualityMaskRange()) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_equality_mask_plus(backend: Backend, monoid): + """ReduceEqualityMaskRange distributes over a plus body, discharging an + equality on the reduced stream in one summand via a singleton-stream gather + while leaving the other summand as an ordinary masked reduce. + """ + a, c = backend.define_vars("a", "c", ret="scalar") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") + + body = monoid.plus( + monoid.mask(f(a()), a() == c()), # eliminable: a == c over range + monoid.mask(g(a()), c() == 0), # not eliminable (no reduced-stream eq) + ) + lhs = monoid.reduce(body, {a: range(3)}) + rhs = monoid.plus( + monoid.reduce( + monoid.mask( + monoid.mask(f(a()), And.plus(0 <= c(), c() < 3)), + And.plus(), + ), + {a: (c(),)}, + ), + monoid.reduce(monoid.mask(g(a()), c() == 0), {a: range(3)}), + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceEqualityMaskRange()) + + +def test_reduce_independent_1(backend: Backend): + a, b = backend.define_vars("a", "b", ret="scalar") + A, B = backend.define_vars("A", "B", ret="stream") + + lhs = Sum.reduce(Product.plus(a(), b()), {a: A(), b: B()}) + rhs = Product.plus( + Sum.reduce(Product.plus(a()), {a: A()}), Sum.reduce(Product.plus(b()), {b: B()}) + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=Factor()) + + +def test_reduce_independent_2(backend: Backend): + a, b, c = backend.define_vars("a", "b", "c", ret="scalar") + A, B, C = backend.define_vars("A", "B", "C", ret="stream") + f = backend.define_vars( + "f", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) + + lhs = Sum.reduce(Product.plus(a(), b(), f(b(), c())), {a: A(), b: B(), c: C()}) + rhs = Product.plus( + Sum.reduce(Product.plus(a()), {a: A()}), + Sum.reduce( + Product.plus(b(), Sum.reduce(Product.plus(f(b(), c())), {c: C()})), + {b: B()}, + ), + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=Factor()) + + +def test_reduce_independent_3_negative(backend: Backend): + """Stream `b` depends on `a` (b: g(a())), so the proposed factorization + is unsound — the normalizer must NOT apply it.""" + a, b, c = backend.define_vars("a", "b", "c", ret="scalar") + A, C = backend.define_vars("A", "C", ret="stream") + f = backend.define_vars( + "f", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) + g = backend.define_vars("g", arg_types=(backend.scalar_typ,), ret="stream") + + with handler(Factor()): # ty:ignore[invalid-argument-type] + lhs = Sum.reduce( + Product.plus(a(), b(), f(b(), c())), {a: A(), b: g(a()), c: C()} + ) + bogus_rhs = Product.plus( + Sum.reduce(a(), {a: A()}), + Sum.reduce(Product.plus(b(), f(b(), c())), {b: g(a()), c: C()}), + ) + assert fvsof(bogus_rhs) != fvsof(lhs) + assert not syntactic_eq_alpha(lhs, bogus_rhs) + + +def test_reduce_independent_4(backend: Backend): + a, b, c, d = backend.define_vars("a", "b", "c", "d", ret="scalar") + A, B, C = backend.define_vars("A", "B", "C", ret="stream") + f = backend.define_vars( + "f", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) + + lhs = Sum.reduce(Product.plus(a(), b(), f(b(), c()), d()), {a: A(), b: B(), c: C()}) + rhs = Product.plus( + d(), + Sum.reduce(Product.plus(a()), {a: A()}), + Sum.reduce( + Product.plus(b(), Sum.reduce(Product.plus(f(b(), c())), {c: C()})), + {b: B()}, + ), + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=Factor()) + + +def test_reduce_chain(backend: Backend): + x, y = backend.define_vars("x", "y", ret="scalar") + X, Y = backend.define_vars("X", "Y", ret="stream") + f, h = backend.define_vars("f", "h", arg_types=(backend.scalar_typ,), ret="scalar") + g = backend.define_vars( + "g", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) + + lhs = Sum.reduce(Product.plus(f(x()), g(x(), y()), h(y())), {x: X(), y: Y()}) + rhs = Sum.reduce( + Product.plus(h(y()), Sum.reduce(Product.plus(f(x()), g(x(), y())), {x: X()})), + {y: Y()}, + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=Factor()) + + +@pytest.mark.parametrize("outer,inner", MONOID_PAIRS) +def test_reduce_lift_shared(outer, inner, backend: Backend): + """A stream free in every factor is hoisted into an outer reduce: + Sum.reduce(f(a, c) * g(b, c), {a: A, b: B, c: C}) + = Sum.reduce(Sum.reduce(f(a, c), {a: A}) * Sum.reduce(g(b, c), {b: B}), {c: C}) + """ + a, b, c = backend.define_vars("a", "b", "c", ret="scalar") + A, B, C = backend.define_vars("A", "B", "C", ret="stream") + f, g = backend.define_vars( + "f", "g", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) + + lhs = outer.reduce(inner.plus(f(a(), c()), g(b(), c())), {a: A(), b: B(), c: C()}) + rhs = outer.reduce( + inner.plus( + outer.reduce(inner.plus(f(a(), c())), {a: A()}), + outer.reduce(inner.plus(g(b(), c())), {b: B()}), + ), + {c: C()}, + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=Factor()) + + +@pytest.mark.parametrize("outer,inner", MONOID_PAIRS) +def test_reduce_lift_shared_deps(outer, inner, backend: Backend): + """A shared stream is lifted together with its dependencies: both ``c`` + and ``d = h(c)`` appear in every factor, so both are hoisted.""" + a, b, c, d = backend.define_vars("a", "b", "c", "d", ret="scalar") + A, B, C = backend.define_vars("A", "B", "C", ret="stream") + h = backend.define_vars("h", arg_types=(backend.scalar_typ,), ret="stream") + f, g = backend.define_vars( + "f", + "g", + arg_types=(backend.scalar_typ, backend.scalar_typ, backend.scalar_typ), + ret="scalar", + ) + + lhs = outer.reduce( + inner.plus(f(a(), c(), d()), g(b(), c(), d())), + {a: A(), b: B(), c: C(), d: h(c())}, + ) + rhs = outer.reduce( + inner.plus( + outer.reduce(inner.plus(f(a(), c(), d())), {a: A()}), + outer.reduce(inner.plus(g(b(), c(), d())), {b: B()}), + ), + {c: C(), d: h(c())}, + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=Factor()) + + +def test_cartesian_union(): + i, a = Operation.define(int), Operation.define(int) + lhs = CartesianProduct.reduce( + Union.reduce([as_dict(((i(),), a()))], {a: range(2)}), {i: range(2)} + ) + with handler(NormalizeIntp), handler(EvaluateIntp): + rhs = evaluate(lhs) + assert syntactic_eq( + rhs, + [ + {(0,): 0, (1,): 0}, + {(0,): 0, (1,): 1}, + {(0,): 1, (1,): 0}, + {(0,): 1, (1,): 1}, + ], + ) + + +_LiftingIntp = functools.reduce( + coproduct, # type: ignore + ( + ReduceDistributeCartesianProduct(), + ReduceUnion(), + EliminateSingletonStreams(), + ReduceEmpty(), + PlusSingle(), + ), +) + + +@pytest.mark.parametrize("outer,inner", MONOID_PAIRS) +def test_reduce_lifted_1(outer, inner, backend: Backend): + a, i = backend.define_vars("a", "i", ret="scalar") + N, A_domain = backend.define_vars("N", "A_domain", ret="stream") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + A = Operation.define(Mapping[tuple, backend.scalar_typ]) # type: ignore[name-defined] + + lhs = outer.reduce( + inner.reduce(f(A()[(i(),)]), {i: range(3)}), + { + A: CartesianProduct.reduce( + Union.reduce([as_dict(((i(),), a()))], {a: A_domain()}), {i: range(3)} + ) + }, + ) + rhs = inner.reduce(outer.reduce(f(a()), {a: A_domain()}), {i: range(3)}) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=_LiftingIntp) + + +def test_reduce_lifted_aligns_equal_ranges_by_row_position(): + """Equal-sized axes are aligned by row position, not stream order.""" + backend = JaxBackend() + a, p, q, i, j, k, l = backend.define_vars( + "a", "p", "q", "i", "j", "k", "l", ret="scalar" + ) + A = Operation.define(Mapping[tuple, backend.scalar_typ], name="A") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") + plate = range(2) + + lhs = Sum.reduce( + Product.plus( + Product.reduce(f(A()[(i(), j())]), {j: plate, i: plate}), + Product.reduce(g(A()[(k(), l())]), {k: plate, l: plate}), + ), + { + A: CartesianProduct.reduce( + Union.reduce([as_dict(((p(), q()), a()))], {a: range(3)}), + {p: plate, q: plate}, + ) + }, + ) + rhs = Product.reduce( + Product.reduce( + Sum.reduce(Product.plus(f(a()), g(a())), {a: range(3)}), {q: plate} + ), + {p: plate}, + ) + + norm = handler(_LiftingIntp)(evaluate)(lhs) + assert syntactic_eq_alpha(norm, rhs) + + +def test_reduce_lifted_unfactors_single_product_factor(): + """A unary product body can be exposed without an explicit Product.plus.""" + backend = JaxBackend() + a, b, i, p, h = backend.define_vars("a", "b", "i", "p", "h", ret="scalar") + A = Operation.define(Mapping[tuple, backend.scalar_typ], name="A") + f = backend.define_vars( + "f", + arg_types=(backend.scalar_typ, backend.scalar_typ), + ret="scalar", + ) + plate = range(2) + + lhs = Sum.reduce( + Product.plus( + Sum.reduce(Product.reduce(f(A()[(i(),)], b()), {i: plate}), {b: range(4)}), + h(), + ), + { + A: CartesianProduct.reduce( + Union.reduce([as_dict(((p(),), a()))], {a: range(3)}), {p: plate} + ) + }, + ) + rhs = Sum.reduce( + Product.plus( + Product.reduce(Sum.reduce(f(a(), b()), {a: range(3)}), {p: plate}), h() + ), + {b: range(4)}, + ) + + norm = handler(_LiftingIntp)(evaluate)(lhs) + assert syntactic_eq_alpha(norm, rhs) + + +def test_reduce_lifted_3(backend: Backend): + a, b, i, j = backend.define_vars("a", "b", "i", "j", ret="scalar") + N, A_domain, B_domain = backend.define_vars( + "N", "A_domain", "B_domain", ret="stream" + ) + f = backend.define_vars( + "f", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) + A = Operation.define(Mapping[tuple, backend.scalar_typ], name="A") # type: ignore[name-defined] + B = Operation.define(Mapping[tuple, backend.scalar_typ], name="B") # type: ignore[name-defined] + + lhs = Sum.reduce( + Product.reduce(f(A()[(i(), j())], B()[(i(), j())]), {i: range(2), j: range(3)}), + { + A: CartesianProduct.reduce( + Union.reduce([as_dict(((i(), j()), a()))], {a: A_domain()}), + {i: range(2), j: range(3)}, + ), + B: CartesianProduct.reduce( + Union.reduce([as_dict(((i(), j()), b()))], {b: B_domain()}), + {i: range(2), j: range(3)}, + ), + }, + ) + rhs = Product.reduce( + Product.reduce( + Sum.reduce(Sum.reduce(f(a(), b()), {a: A_domain()}), {b: B_domain()}), + {j: range(3)}, + ), + {i: range(2)}, + ) + norm = handler(_LiftingIntp)(evaluate)(lhs) + assert syntactic_eq_alpha(norm, rhs) + + +@pytest.mark.parametrize("outer,inner", MONOID_PAIRS) +def test_reduce_lifted_multi_index(outer, inner, backend: Backend): + a, i, j = backend.define_vars("a", "i", "j", ret="scalar") + N, M, A_domain = backend.define_vars("N", "M", "A_domain", ret="stream") + A = Operation.define(Mapping[tuple, backend.scalar_typ]) # type: ignore[name-defined] + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + + lhs = outer.reduce( + inner.reduce(f(A()[(i(), j())]), {i: range(3), j: range(2)}), + { + A: CartesianProduct.reduce( + Union.reduce([as_dict(((i(), j()), a()))], {a: A_domain()}), + {i: range(3), j: range(2)}, + ) + }, + ) + rhs = inner.reduce( + inner.reduce(outer.reduce(f(a()), {a: A_domain()}), {j: range(2)}), + {i: range(3)}, + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=_LiftingIntp) + + +@pytest.mark.parametrize("outer,inner", MONOID_PAIRS) +def test_reduce_lifted_2(outer, inner, backend: Backend): + """The worked example on page 396 of 'Lifted Variable Elimination: + Decoupling the Operators from the Constraint Language'. + + """ + a, i, s, t = backend.define_vars("a", "i", "s", "t", ret="scalar") + N, T = backend.define_vars("N", "T", ret="stream") + A = Operation.define(Mapping[tuple, backend.scalar_typ]) # type: ignore[name-defined] + A_domain = backend.define_vars( + "A_domain", arg_types=(backend.scalar_typ,), ret="stream" + ) + f1, f2 = backend.define_vars( + "f1", "f2", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) + + lhs = outer.reduce( + inner.reduce( + inner.plus(f1(A()[(i(),)], s()), f2(t(), A()[(i(),)])), {i: range(3)} + ), + { + A: CartesianProduct.reduce( + Union.reduce([as_dict(((i(),), a()))], {a: A_domain(i())}), + {i: range(3)}, + ), + t: T(), + }, + ) + rhs = outer.reduce( + inner.reduce( + outer.reduce(inner.plus(f1(a(), s()), f2(t(), a())), {a: A_domain(i())}), + {i: range(3)}, + ), + {t: T()}, + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=_LiftingIntp) + + +# --------------------------------------------------------------------------- +# Weighted streams +# --------------------------------------------------------------------------- + + +def test_reduce_single_weighted_stream(backend: Backend): + """Single weighted stream desugars: + Sum.reduce(body, {a: WS(A, w, Product)}) + = Sum.reduce(Product.plus(w(a), body), {a: A}) + """ + a = backend.define_vars("a", ret="scalar") + A = backend.define_vars("A", ret="stream") + body, w = backend.define_vars( + "body", "w", arg_types=(backend.scalar_typ,), ret="scalar" + ) + + lhs = Sum.reduce(body(a()), {a: Product.weighted(A(), w)}) + rhs = Sum.reduce(Product.plus(w(a()), body(a())), {a: A()}) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceWeightedStream()) + + +def test_reduce_weighted_factorization(backend: Backend): + """Two independent weighted streams under Sum with Product weights factor: + Sum.reduce(f(a)*g(b), {a: Product.weighted(A, a, w_a), b: Product.weighted(B, b, w_b)}) + = (Sum.reduce(w_a(a)*f(a), {a: A})) * (Sum.reduce(w_b(b)*g(b), {b: B})) + + Exercises chaining of ``ReduceWeightedStream`` with ``Factor`` + inside ``NormalizeIntp``. + """ + a, b = backend.define_vars("a", "b", ret="scalar") + A, B = backend.define_vars("A", "B", ret="stream") + f, g, w_a, w_b = backend.define_vars( + "f", "g", "w_a", "w_b", arg_types=(backend.scalar_typ,), ret="scalar" + ) + + lhs = Sum.reduce( + Product.plus(f(a()), g(b())), + {a: Product.weighted(A(), w_a), b: Product.weighted(B(), w_b)}, + ) + rhs = Product.plus( + Sum.reduce(Product.plus(w_a(a()), Product.plus(f(a()))), {a: A()}), + Sum.reduce(Product.plus(w_b(b()), Product.plus(g(b()))), {b: B()}), + ) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceWeightedStream(), Factor()) + ) + + +def test_weighted_expectation_demo(): + """Demo: compute E[f(X)] = Σ_x w(x)·f(x) via a weighted reduce. + + X ranges over [1, 2, 3, 4] with weights w(x) = x/10 (a valid distribution + since the weights sum to 1) and f(x) = x*x. Expected value: + 0.1·1 + 0.2·4 + 0.3·9 + 0.4·16 = 10.0 + """ + weights = {1: 0.1, 2: 0.2, 3: 0.3, 4: 0.4} + + def _w(v: int) -> float: + if isinstance(v, Term): + raise NotHandled + return weights[v] + + def _f(v: int) -> float: + if isinstance(v, Term): + raise NotHandled + return float(v * v) + + a = Operation.define(int, name="a") + w = Operation.define(_w, name="w") + f = Operation.define(_f, name="f") + + with handler(NormalizeIntp), handler(EvaluateIntp): + result = evaluate(Sum.reduce(f(a()), {a: Product.weighted([1, 2, 3, 4], w)})) + + assert math.isclose(result, 10.0) + + +# --------------------------------------------------------------------------- +# CartesianProduct.plus (pure-Python ``CartesianProductPlus`` implementation) +# --------------------------------------------------------------------------- +# +# A ``CartesianProduct`` value is a list of "rows", each row a ``dict`` mapping +# index variables to values. ``plus`` takes the cartesian product of its +# argument lists, disjoint-merging the dicts of each combination into a single +# row. The identity is ``[{}]`` (one empty row) and the zero is ``[]`` (no +# rows). + + +@pytest.fixture +def cprod(): + """A handler scope in which ``CartesianProduct.plus`` is concrete.""" + with handler(CartesianProductPlus()): + yield + + +def test_cprod_plus_two_singletons(cprod): + """Two single-row lists merge into one row with the union of their keys.""" + assert CartesianProduct.plus([{"a": 1}], [{"b": 2}]) == [{"a": 1, "b": 2}] + + +def test_cprod_plus_multi_key_rows(cprod): + """Every key of a multi-key row survives the merge (regression: the merge + used to keep only the last key of each dict).""" + assert CartesianProduct.plus([{"a": 1, "b": 2}], [{"c": 3, "d": 4}]) == [ + {"a": 1, "b": 2, "c": 3, "d": 4} + ] + + +def test_cprod_plus_cartesian_expansion(cprod): + """The result enumerates the full cartesian product of the input rows.""" + result = CartesianProduct.plus([{"a": 1}, {"a": 2}, {"a": 3}], [{"b": 4}, {"b": 5}]) + assert result == [ + {"a": 1, "b": 4}, + {"a": 1, "b": 5}, + {"a": 2, "b": 4}, + {"a": 2, "b": 5}, + {"a": 3, "b": 4}, + {"a": 3, "b": 5}, + ] + + +def test_cprod_plus_cardinality(cprod): + """|plus(A, B, C)| == |A| * |B| * |C|.""" + a = [{"a": i} for i in range(2)] + b = [{"b": i} for i in range(3)] + c = [{"c": i} for i in range(4)] + assert len(CartesianProduct.plus(a, b, c)) == 2 * 3 * 4 + + +def test_cprod_plus_three_args(cprod): + """``plus`` is variadic and merges across all arguments at once.""" + assert CartesianProduct.plus([{"a": 1}], [{"b": 2}], [{"c": 3}]) == [ + {"a": 1, "b": 2, "c": 3} + ] + + +def test_cprod_plus_single_arg(cprod): + """A single argument is returned row-for-row (product of one factor).""" + assert CartesianProduct.plus([{"a": 1}, {"a": 2}]) == [{"a": 1}, {"a": 2}] + + +def test_cprod_plus_identity_right(cprod): + """The identity ``[{}]`` is a right unit: merging an empty row changes + nothing.""" + assert CartesianProduct.plus([{"a": 1}, {"a": 2}], CartesianProduct.identity) == [ + {"a": 1}, + {"a": 2}, + ] + + +def test_cprod_plus_identity_left(cprod): + """The identity ``[{}]`` is a left unit.""" + assert CartesianProduct.plus(CartesianProduct.identity, [{"a": 1}, {"a": 2}]) == [ + {"a": 1}, + {"a": 2}, + ] + + +def test_cprod_plus_identity_with_identity(cprod): + """Identity ⊕ identity == identity.""" + assert ( + CartesianProduct.plus(CartesianProduct.identity, CartesianProduct.identity) + == CartesianProduct.identity + ) + + +def test_cprod_plus_zero_right(cprod): + """The zero ``[]`` absorbs on the right (empty cartesian product).""" + assert CartesianProduct.plus([{"a": 1}], CartesianProduct.zero) == [] + + +def test_cprod_plus_zero_left(cprod): + """The zero ``[]`` absorbs on the left.""" + assert CartesianProduct.plus(CartesianProduct.zero, [{"a": 1}]) == [] + + +def test_cprod_plus_zero_among_many(cprod): + """A single zero factor anywhere collapses the whole product to ``[]``.""" + assert CartesianProduct.plus([{"a": 1}], [], [{"c": 3}]) == [] + + +def test_cprod_plus_empty_rows_preserved(cprod): + """Rows that are themselves empty dicts merge cleanly.""" + assert CartesianProduct.plus([{}], [{"a": 1}]) == [{"a": 1}] + assert CartesianProduct.plus([{}], [{}]) == [{}] + + +def test_cprod_plus_associative(cprod): + """plus(plus(A, B), C) == plus(A, plus(B, C)) == plus(A, B, C).""" + a = [{"a": 1}, {"a": 2}] + b = [{"b": 3}] + c = [{"c": 4}, {"c": 5}] + flat = CartesianProduct.plus(a, b, c) + assert CartesianProduct.plus(CartesianProduct.plus(a, b), c) == flat + assert CartesianProduct.plus(a, CartesianProduct.plus(b, c)) == flat + + +def test_cprod_plus_duplicate_key_raises(cprod): + """Merging rows that share a key is ill-defined and rejected.""" + with pytest.raises(ValueError, match="Duplicate key found: 'a'"): + CartesianProduct.plus([{"a": 1}], [{"a": 2}]) + + +def test_cprod_plus_duplicate_key_in_multi_key_row_raises(cprod): + """The duplicate-key check sees every key of a multi-key row, not just the + last (regression for the same merge bug).""" + with pytest.raises(ValueError, match="Duplicate key found: 'a'"): + CartesianProduct.plus([{"a": 1, "x": 9}], [{"a": 2}]) + + +def test_cprod_plus_does_not_mutate_inputs(cprod): + """The input rows are left untouched; merges build fresh dicts.""" + left = [{"a": 1}] + right = [{"b": 2}] + CartesianProduct.plus(left, right) + assert left == [{"a": 1}] + assert right == [{"b": 2}] + + +def test_cprod_plus_forwards_on_term(): + """A symbolic (``Term``) argument cannot be enumerated, so ``plus`` forwards + and the call builds an unevaluated ``Term`` instead of a value.""" + x = Operation.define(Iterable, name="x") + with handler(CartesianProductPlus()): + result = CartesianProduct.plus(x(), [{"a": 1}]) + assert isinstance(result, Term) + assert result.op is CartesianProduct.plus + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_dependent_range_mask(monoid, backend: Backend): + """A dependent range stream gets rewritten to the referent's bbox stream, + with the original constraint folded into the body as a where-guard. + + reduce(M, {u: range(0, N, 1), v: range(0, u(), 1)}, body) + ≡ reduce(M, {u: range(0, N, 1), v: range(0, N, 1)}, where(v() < u(), body, M.identity)) + """ + (u, v) = backend.define_vars("u", "v", ret="scalar") + N = 5 + f = backend.define_vars( + "f", arg_types=[backend.scalar_typ, backend.scalar_typ], ret="scalar" + ) + + body = f(u(), v()) + + lhs = monoid.reduce(body, {u: range_(N), v: range_(u())}) + rhs = monoid.reduce(monoid.mask(body, v() < u()), {u: range_(N), v: range_(N)}) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceDependentRangeMask()) + + +@pytest.mark.parametrize("Sum,Product", COMMUTATIVE_MONOID_PAIRS) +def test_reduce_unfactor_simple(Sum, Product, backend: Backend): + x, y, g = backend.define_vars("x", "y", "g", ret="scalar") + X, Y = backend.define_vars("X", "Y", ret="stream") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + lhs = Sum.reduce(Product.plus(Sum.reduce(f(x()), {x: X()}), g()), {y: Y()}) + rhs = Sum.reduce(Product.plus(f(x()), g()), {x: X(), y: Y()}) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceUnfactor()) + + +@pytest.mark.parametrize("Sum,Product", COMMUTATIVE_MONOID_PAIRS) +def test_reduce_unfactor_reduces(Sum, Product, backend: Backend): + x, y, z = backend.define_vars("x", "y", "z", ret="scalar") + X, Y, Z = backend.define_vars("X", "Y", "Z", ret="stream") + f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") + lhs = Sum.reduce( + Product.plus(Sum.reduce(f(x()), {x: X()}), Sum.reduce(g(y()), {y: Y()})), + {z: Z()}, + ) + rhs = Sum.reduce(Product.plus(f(x()), g(y())), {x: X(), y: Y(), z: Z()}) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceUnfactor()) + + +def test_ground_cartesian_product_substitutes_dependent_plate_domain( + backend: Backend, +): + """Each grounded row value retains the domain of its plate assignment.""" + row_value, plate, t = backend.define_vars("row_value", "plate", "t", ret="scalar") + domain = backend.define_vars( + "domain", arg_types=(backend.scalar_typ,), ret="stream" + ) + factor = backend.define_vars( + "factor", + arg_types=( + backend.scalar_typ, + backend.scalar_typ, + backend.scalar_typ, + ), + ret="scalar", + ) + row = Operation.define(Mapping[tuple, backend.scalar_typ], name="row") # type: ignore[name-defined] + + lhs = Sum.reduce( + Product.reduce( + factor(t(), row()[(t(),)], row()[(t() + 1,)]), + {t: range(2)}, + ), + { + row: CartesianProduct.reduce( + Union.reduce( + [as_dict(((plate(),), row_value()))], + {row_value: domain(plate())}, + ), + {plate: range(3)}, + ) + }, + ) + + grounded = [Operation.define(row_value, name=f"row_value_{i}") for i in range(3)] + rhs = Sum.reduce( + Product.plus( + factor(0, grounded[0](), grounded[1]()), + factor(1, grounded[1](), grounded[2]()), + ), + {grounded[i]: domain(i) for i in range(3)}, + ) + + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=ReduceGroundCartesianProduct(), + ) + + +def test_ground_cartesian_product_declines_diagonal_chain(backend: Backend): + """A diagonal does not consume the off-diagonal cells of a two-dimensional row.""" + row_value, p, q, t = backend.define_vars("row_value", "p", "q", "t", ret="scalar") + factor = backend.define_vars( + "factor", + arg_types=(backend.scalar_typ, backend.scalar_typ), + ret="scalar", + ) + row = Operation.define(Mapping[tuple, backend.scalar_typ], name="row") # type: ignore[name-defined] + + lhs = Sum.reduce( + Product.reduce( + factor( + row()[(t(), t())], + row()[(t() + 1, t() + 1)], + ), + {t: range(1)}, + ), + { + row: CartesianProduct.reduce( + Union.reduce( + [as_dict(((p(), q()), row_value()))], {row_value: range(2)} + ), + {p: range(2), q: range(2)}, + ) + }, + ) + + backend.check_rewrite( + lhs=lhs, + rhs=lhs, + rule=ReduceGroundCartesianProduct(), + ) + + +def test_ground_cartesian_product_extra_value_streams(backend: Backend): + """Extra row value streams are retained.""" + row_value, row_value1, row_value2, p, q, t = backend.define_vars( + "row_value", "row_value1", "row_value2", "p", "q", "t", ret="scalar" + ) + row = Operation.define(Mapping[tuple, backend.scalar_typ], name="row") # type: ignore[name-defined] + + lhs = Sum.reduce( + Product.reduce(row()[(t(),)], {t: range(2)}), + { + row: CartesianProduct.reduce( + Union.reduce( + [as_dict(((p(),), row_value()))], {row_value: range(2), q: range(3)} + ), + {p: range(2)}, + ) + }, + ) + + rhs = Sum.reduce( + Product.plus(row_value1(), row_value2()), + { + row_value1: Union.reduce([row_value()], {q: range(3), row_value: range(2)}), + row_value2: Union.reduce([row_value()], {q: range(3), row_value: range(2)}), + }, + ) + + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=ReduceGroundCartesianProduct(), + ) + + +@pytest.mark.parametrize("T,K", [(10, 3)]) +def test_ground_cartesian_product_chain(T, K): + """A chain over a cartesian product grounds into a forward pass. + + ``ReduceDistributeCartesianProduct`` cannot invert this reduction: the body + subscripts the row at ``t`` and at ``t + 1``, so no single plate index + factors out. ``ReduceGroundCartesianProduct`` instead mints one variable + per plate assignment the body actually asks for, leaving an ordinary + variable-elimination problem whose elimination order is the forward + algorithm. + + Written out longhand; the generator-comprehension spelling + + .. code-block:: python + + Sum( + Product(phi()[t][ixs[t]][ixs[t + 1]] for t in range(T - 1)) + for ixs in CartesianProduct(range(K) for _ in range(T)) + ) + + desugars to exactly the reduction below, but the desugaring itself lands + separately. + + ``T`` is 10 rather than the 20 this scales to because normalization is + superlinear in the chain length until the ``evaluate`` fast path lands; at + 10 the whole test takes a few seconds, at 20 it takes over an hour. + """ + fs = [ + [[random.uniform(0, 1) for _ in range(K)] for _ in range(K)] for _ in range(T) + ] + + @Operation.define + def phi() -> list[list[list[float]]]: + raise NotHandled + + t = Operation.define(int, name="t") + plate = Operation.define(int, name="plate") + row_value = Operation.define(int, name="row_value") + ixs = Operation.define(Mapping[tuple, int], name="ixs") + + with handler(NormalizeIntp): + zf_normal = Sum.reduce( + Product.reduce( + phi()[t()][ixs()[(t(),)]][ixs()[(t() + 1,)]], + {t: range(T - 1)}, + ), + { + ixs: CartesianProduct.reduce( + Union.reduce( + [as_dict(((plate(),), row_value()))], {row_value: range(K)} + ), + {plate: range(T)}, + ) + }, + ) + + # Grounding is what does the work here: no other rule applies to this + # reduction, so without it normalization is the identity and the answer is + # only recovered by enumerating all ``K ** T`` rows during evaluation. + assert isinstance(zf_normal, Term) and zf_normal.op is Sum.reduce + assert list(zf_normal.args[1].values()) == [range(K)] + + with handler(EvaluateIntp), handler({phi: lambda: fs}): + zf_actual = evaluate(zf_normal) + + alpha = [1.0] * K + for t_ in range(T - 1): + alpha = [sum(alpha[i] * fs[t_][i][j] for i in range(K)) for j in range(K)] + zf_expected = sum(alpha) + assert isinstance(zf_actual, float) and math.isclose(zf_actual, zf_expected) diff --git a/tests/test_ops_semantics.py b/tests/test_ops_semantics.py index 04e595f9e..66c391631 100644 --- a/tests/test_ops_semantics.py +++ b/tests/test_ops_semantics.py @@ -24,6 +24,7 @@ deffn, defop, implements, + syntactic_eq, ) from effectful.ops.types import Interpretation, NotHandled, Operation, Term @@ -469,7 +470,9 @@ def Nested(*args, **kwargs): t = Nested([{"a": y()}, x(), (x(), y())], x(), arg1={"b": x()}) with handler({x: lambda: 1, y: lambda: 2}): - assert evaluate(t) == Nested([{"a": 2}, 1, (1, 2)], 1, arg1={"b": 1}) + assert syntactic_eq( + evaluate(t), Nested([{"a": 2}, 1, (1, 2)], 1, arg1={"b": 1}) + ) def test_memoized_interpretation(): diff --git a/tests/test_ops_syntax.py b/tests/test_ops_syntax.py index a5fdb749c..ccfd2dae0 100644 --- a/tests/test_ops_syntax.py +++ b/tests/test_ops_syntax.py @@ -489,7 +489,6 @@ def _(self, x: bool) -> bool: ) assert isinstance(term_float, Term) - assert term_float.op.__name__ == "my_singledispatch" assert term_float.args == (1.5,) assert term_float.kwargs == {}