From 26c81cff0a70833af27929af4a044757f434919b Mon Sep 17 00:00:00 2001 From: Eli Date: Tue, 28 Jul 2026 11:21:44 -0400 Subject: [PATCH] Simplify cartesian-product inversion with a peeled row Standing a `_PeelRow` in for the row variable puts it in the mapping position of every `row[..]`, so a subscript is recognised by what it is indexing rather than by comparing against the variable it came from. One evaluation pass peels every summand and records the plate each folds over, replacing the per-summand `row_substitute` closure and its `InvalidIndexError` control flow. Nested same-monoid reductions are absorbed inline at the top level rather than by renormalizing the whole candidate under a dedicated interpretation, so `CartesianProductNormalizeIntp` is no longer needed. Co-Authored-By: Claude Opus 5 (1M context) --- effectful/ops/monoid.py | 283 +++++++++++++++++++++++++--------------- 1 file changed, 177 insertions(+), 106 deletions(-) diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index fec2cdce8..5ef625722 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -11,7 +11,6 @@ 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, @@ -21,6 +20,7 @@ as_dict, defdata, deffn, + defop, implements, ite, range_, @@ -905,6 +905,74 @@ def reduce(self, monoid: Monoid, body, streams): return fwd() +class _CannotPeel(Exception): + """A subscript of a row does not fold over the plate being peeled.""" + + +class _PeelRow(collections.abc.Mapping): + """A row with one index position peeled off. + + Standing this in for a row variable puts it in the mapping position of + every ``row[..]``, so a subscript can be recognised by what it is indexing + rather than by comparing against the variable it came from. Each subscript + becomes one of the narrowed row, and the variable it used at the peeled + position -- the plate that subscript folds over -- is recorded. + + Row indices are symbolic, so the subscripts never reach here on their own: + :meth:`_MappingTerm.__getitem__` only hands a mapping a *concrete* key. + The rule routes them here instead. + + ``row[p]``, with nothing left once the position is dropped, is the row's + own value: the narrowed row has no index left to take. + """ + + def __init__(self, position: int, peeled: Operation, value): + self._position = position + self._peeled = peeled + self._value = value + # How often each plate was indexed at the peeled position, against how + # often it was used at all. A plate used no more often than it was + # peeled does not survive peeling, so nothing refers to it afterwards. + self.plates: dict[Operation, int] = {} + self.uses: dict[Operation, int] = {} + + def used(self, plate: Operation) -> None: + self.uses[plate] = self.uses.get(plate, 0) + 1 + + def survives(self, plate: Operation) -> bool: + return self.uses.get(plate, 0) > self.plates.get(plate, 0) + + def __getitem__(self, key: Any) -> Any: + key = tuple(key) if isinstance(key, Sequence) else (key,) + if len(key) <= self._position: + raise _CannotPeel + folded = key[self._position] + if not (isinstance(folded, Term) and not folded.args and not folded.kwargs): + raise _CannotPeel + self.plates[folded.op] = self.plates.get(folded.op, 0) + 1 + + rest = tuple(k for (i, k) in enumerate(key) if i != self._position) + return self._peeled()[rest] if rest else self._value + + def __iter__(self): + return iter(()) + + def __len__(self) -> int: + return 0 + + +@evaluate.register +def _(expr: _PeelRow): + return expr + + +def _peel_subscript(mapping, key): + """Route a subscript of a row being peeled to the row itself.""" + if isinstance(mapping, _PeelRow): + return mapping[key] + return fwd() + + class ReduceDistributeCartesianProduct(ObjectInterpretation): """Eliminates a reduce over a cartesian product. ∑_x₁ ∑_x₂ ... ∑_xₙ ∏_i f(xᵢ) = ∏_i ∑_xᵢ f(xᵢ) @@ -940,15 +1008,28 @@ def reduce(self, monoid: Monoid, body, streams): ): 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 + # ``Factor`` may have moved product factors into nested reductions of + # this same monoid, which hides the product this rule is looking for. + # Absorbing their streams back into this one exposes it again. This is + # what ``ReduceUnfactor`` does, applied here to the body's own factors + # rather than by renormalizing the whole candidate: the shape being + # looked for is a product of factors, so only the top level matters. + if isinstance(body, Term) and _is_monoid_plus(body.op): + absorbed = dict(streams) + factors = [] + for factor in body.args: + if ( + isinstance(factor, Term) + and factor.op is monoid.reduce + and not (set(typing.cast(Mapping, factor.args[1])) & set(absorbed)) + ): + nested_body, nested_streams = factor.args + absorbed |= typing.cast(Mapping, nested_streams) + factors.append(nested_body) + else: + factors.append(factor) + if absorbed != streams: + body, streams = body.op(*factors), absorbed inner_reduces: tuple[Term, ...] if isinstance(body, Term) and _is_monoid_reduce(body.op): @@ -975,8 +1056,6 @@ def reduce(self, monoid: Monoid, body, streams): 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) @@ -1033,79 +1112,98 @@ def drop_elem(ls, index): ) 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)) + # Standing a peeled row in for the row variable puts it in the + # mapping position of every ``row[..]``, and routing subscripts of + # it back to it rewrites each one and records the plate it folds + # over. All that is left to check per summand is that the plate it + # named is the one that summand folds over. + peeled_var = defop(stream_key) + row = _PeelRow(plate_index, peeled_var, union_body) + + def _count(plate): + def use(*a, **k): + row.used(plate) 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 + return use + + peeling = coproduct( + typing.cast(Interpretation, {stream_key: lambda: row}), + typing.cast( + Interpretation, + {_MappingTerm.__getitem__: _peel_subscript} + | { + plate: _count(plate) + for reduce in row_inner_reduces + for plate in typing.cast(Mapping, reduce.args[1]) + }, + ), + ) try: - substituted = [ - row_substitute(*reduce.args) for reduce in row_inner_reduces - ] - except InvalidIndexError: + # One pass peels every summand. Which plate each one folds over + # need not be tracked as they go: a summand binds its own plate + # variable, so it is the one of those the row was indexed by + # that this summand's bundle binds. + peeled_bodies = typing.cast( + tuple, + handler(peeling)(evaluate)( + tuple(reduce.args[0] for reduce in row_inner_reduces) + ), + ) + + substituted = [] + for subst, inner_reduce in zip(peeled_bodies, row_inner_reduces): + inner_streams = typing.cast(Mapping, inner_reduce.args[1]) + if stream_key in fvsof(inner_streams): + raise _CannotPeel + named = set(inner_streams) & set(row.plates) + if len(named) > 1: + raise _CannotPeel + inner_plate_op = next(iter(named), None) + if inner_plate_op is None: + # 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. + if inner_streams: + raise _CannotPeel + elif not syntactic_eq(inner_streams[inner_plate_op], plate_range): + raise _CannotPeel + substituted.append((subst, inner_plate_op)) + except _CannotPeel: 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 - ) + renaming = { + inner_plate_op: shared_plate_op + for (_, inner_plate_op) in substituted + if inner_plate_op is not shared_plate_op + and row.survives(typing.cast(Operation, inner_plate_op)) + } + unified = tuple(subst_body for (subst_body, _) in substituted) + if renaming: + unified = typing.cast( + tuple, + handler(typing.cast(Interpretation, renaming))(evaluate)(unified), + ) - 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_factors = [ + inner_monoid.reduce( + subst_body, + { + k: v + for (k, v) in typing.cast(Mapping, inner_reduce.args[1]).items() + if k != plate + }, + ) + for subst_body, (_, plate), inner_reduce in zip( + unified, substituted, row_inner_reduces + ) + ] combined = ( combined_factors[0] @@ -1118,16 +1216,9 @@ def _getitem(mapping, idx1): 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) + # Nothing is left to index, and the peeled row already answered + # each subscript with the row's value. + inner_reduce_body = monoid.reduce(combined, union_streams) else: peeled_body = Union.reduce( [as_dict((peeled_idx, union_body))], union_streams @@ -1135,7 +1226,7 @@ def _to_body(mapping, key): peeled_cprod = CartesianProduct.reduce( peeled_body, peeled_cprod_streams ) - inner_reduce_body = monoid.reduce(combined, {stream_key: peeled_cprod}) + inner_reduce_body = monoid.reduce(combined, {peeled_var: peeled_cprod}) peeled_reduce = inner_monoid.reduce( inner_reduce_body, {shared_plate_op: plate_range} @@ -1853,26 +1944,6 @@ def extend(self, *intps: Interpretation) -> typing.Self: 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(),