arborist/tests/test_pi_star_compositions.py

"""Tests for π* compositions across the SymPy substrate (#000030).

The compose API auto-registers composite kernels under the key
``<inner-name>-then-<outer-name>@v1``. These tests exercise pairs
that compose naturally — kernels that emit canonical bytes the
next kernel can re-parse.

What composes:

- ``algebra-symbolic@v1 ∘ algebra-symbolic@v1`` — idempotency.
  Inner produces ``sp.expand(sp.sympify(text))`` srepr; outer
  sympifies the srepr back, expands again. Should equal the
  one-pass output.

- ``algebra-symbolic@v1 ∘ algebra-symbolic-simplified@v1`` —
  expand then simplify. Catches polynomial identity AND trig /
  exp / log identity that single-pass `expand` misses.

- ``calculus-derivative@v1 ∘ algebra-symbolic@v1`` — derivative
  composed with re-canonicalizer. Output of derivative is already
  re-canonicalized through algebra-symbolic internally, so the
  composite is idempotent over the derivative's output.

What doesn't compose (documented gaps):

- Most cross-π* pairs require shape-compatible round-trip; e.g.
  ``arithmetic@v1 ∘ algebra-symbolic@v1`` would require arithmetic
  to accept symbolic-srepr bytes, which it doesn't. Compose
  raises a clean error path on incompatible bytes — see
  ``test_composition_propagates_pi_star_error``.
"""

from __future__ import annotations

import pytest


sympy = pytest.importorskip("sympy")


def _safe_compose(inner_key: str, outer_key: str):
    """Compose without re-registering. The registry rejects dup
    keys (#000015 invariant), so test ordering would otherwise
    matter. Using register_in_registry=False sidesteps the
    invariant for tests that only need .canonicalize() access."""
    from arborist.pi_star import compose
    return compose(inner_key, outer_key, register_in_registry=False)


def test_compose_auto_registers_composite_key():
    """Calling compose() registers the result under
    ``<inner>-then-<outer>@v1`` so subsequent get() calls work.
    Re-registration is rejected by the registry — operators that
    just want to construct the composite (without re-registering)
    pass ``register_in_registry=False``."""
    from arborist.pi_star import compose, get, list_keys
    key = "algebra-symbolic-then-algebra-symbolic@v1"
    if key not in list_keys():
        composite = compose(
            "algebra-symbolic@v1", "algebra-symbolic@v1",
        )
    else:
        # Already registered by an earlier test; build without
        # re-registering, then look up the registered one.
        composite = compose(
            "algebra-symbolic@v1", "algebra-symbolic@v1",
            register_in_registry=False,
        )
    assert composite.name == "algebra-symbolic-then-algebra-symbolic"

    # Lookup via get() works after registration.
    looked_up = get(key)
    assert looked_up.name == composite.name


def test_algebra_symbolic_idempotent_under_self_composition():
    """Running expand twice should equal expand once. The composite
    key bytes must equal the one-pass key bytes for any expression."""
    from arborist.pi_star import compose, get
    one_pass = get("algebra-symbolic@v1")
    twice = _safe_compose("algebra-symbolic@v1", "algebra-symbolic@v1")

    for expr in (b"(x+1)**2", b"x**2 + 2*x + 1", b"a*b - c", b"sin(x)"):
        single = one_pass.canonicalize(expr)
        composed = twice.canonicalize(expr)
        assert single == composed, f"divergence on {expr!r}"


def test_simplify_composition_collapses_trig_identity():
    """`expand` alone leaves sin²+cos² untouched; simplify catches it.
    Composing expand-then-simplify should produce ``Integer(1)`` for
    the Pythagorean identity."""
    from arborist.pi_star import compose
    composite = _safe_compose(
        "algebra-symbolic@v1", "algebra-symbolic-simplified@v1",
    )
    out = composite.canonicalize(b"sin(x)**2 + cos(x)**2")
    # SymPy's srepr of Integer(1) is exactly "Integer(1)"
    assert out == b"Integer(1)"


def test_simplify_composition_does_not_break_polynomial_identity():
    """Expand-then-simplify on a polynomial that's already in
    expanded form should round-trip stably."""
    from arborist.pi_star import compose
    composite = _safe_compose(
        "algebra-symbolic@v1", "algebra-symbolic-simplified@v1",
    )
    a = composite.canonicalize(b"(x+1)**2")
    b = composite.canonicalize(b"x**2 + 2*x + 1")
    # Both should collapse to the same canonical bytes.
    assert a == b


def test_composition_propagates_pi_star_error():
    """When the inner kernel raises PiStarError, the composite
    surfaces it rather than swallowing or converting."""
    from arborist.pi_star import PiStarError, compose
    composite = _safe_compose(
        "algebra-symbolic@v1", "algebra-symbolic@v1",
    )
    with pytest.raises(PiStarError):
        composite.canonicalize(b"this is not !#$ a valid !#$ sympy expression")


def test_composition_is_listed_in_registry_after_compose():
    """``list_keys()`` shows composite kernels alongside the rest of
    the registry once they've been compose()'d. Auto-registered
    composites have keys with ``-then-`` infix."""
    from arborist.pi_star import compose, list_keys
    compose("algebra-symbolic@v1", "algebra-symbolic-simplified@v1")
    keys = list_keys()
    assert any(
        k.startswith("algebra-symbolic-then-") for k in keys
    )


def test_canonical_composition_id_is_stable():
    """The composition manifest fingerprint depends only on the
    (inner, outer) keys; same pair → same id."""
    from arborist.pi_star import canonical_composition_id
    a = canonical_composition_id(
        "algebra-symbolic@v1", "algebra-symbolic-simplified@v1",
    )
    b = canonical_composition_id(
        "algebra-symbolic@v1", "algebra-symbolic-simplified@v1",
    )
    assert a == b
    assert len(a) == 64  # sha256 hex


def test_canonical_composition_id_distinguishes_order():
    """Order matters in composition: A∘B ≠ B∘A, so the manifest id
    must distinguish them."""
    from arborist.pi_star import canonical_composition_id
    forward = canonical_composition_id(
        "algebra-symbolic@v1", "algebra-symbolic-simplified@v1",
    )
    reversed_ = canonical_composition_id(
        "algebra-symbolic-simplified@v1", "algebra-symbolic@v1",
    )
    assert forward != reversed_


def test_compose_unknown_inner_raises_keyerror():
    """compose() guards both ends — an unknown inner π* fails
    cleanly with KeyError."""
    from arborist.pi_star import compose
    with pytest.raises(KeyError):
        compose("does-not-exist@v1", "algebra-symbolic@v1")


def test_compose_unknown_outer_raises_keyerror():
    from arborist.pi_star import compose
    with pytest.raises(KeyError):
        compose("algebra-symbolic@v1", "does-not-exist@v1")


def test_composite_domain_is_inner_domain():
    """Composition's domain is the inner π*'s domain (the composite
    accepts what the inner accepts)."""
    from arborist.pi_star import compose
    composite = _safe_compose(
        "algebra-symbolic@v1", "algebra-symbolic@v1",
    )
    assert composite.domain == "symbolic-algebra"


def test_composition_bytes_match_manual_application():
    """A composite kernel must produce byte-equal output to a manual
    inner.canonicalize() → outer.canonicalize() chain."""
    from arborist.pi_star import compose, get
    inner = get("algebra-symbolic@v1")
    outer = get("algebra-symbolic-simplified@v1")
    composite = _safe_compose(
        "algebra-symbolic@v1", "algebra-symbolic-simplified@v1",
    )
    for expr in (
        b"sin(x)**2 + cos(x)**2",
        b"(a + b)**2",
        b"x**3 - x**3 + 2*x",
    ):
        manual = outer.canonicalize(inner.canonicalize(expr))
        assert composite.canonicalize(expr) == manual, (
            f"composite ≠ manual chain on {expr!r}"
        )