arborist/tests/test_pi_star_arithmetic.py

"""arithmetic@v1 π* tests.

Exact-rational arithmetic canonicalizer over Σ* → ℚ ∪ {⊥}. KATs
mirror the docstring's equivalence-class examples, plus negative
cones (identifier, function call, non-integer exponent, division
by zero, boolean literal, empty, non-bytes, syntax error) and
round-trip idempotence.

Backfills the test-coverage gap for arborist/pi_star/arithmetic.py
identified during the 2026-05-10 zero-coverage sweep.
"""

from __future__ import annotations

import pytest

from arborist.pi_star import get
from arborist.pi_star.protocol import (
    PiStarError,
    assert_round_trip,
    equivalence_class_id,
)


@pytest.fixture
def ps():
    return get("arithmetic@v1")


# --- registry presence + metadata ------------------------------------


def test_registry_contains_arithmetic():
    ps = get("arithmetic@v1")
    assert ps.name == "arithmetic"
    assert ps.version == "v1"
    assert ps.domain == "arithmetic"


# --- positive KATs (docstring equivalence classes) -------------------


@pytest.mark.parametrize(
    "expr,canonical",
    [
        # Integer literal collapse
        ("1+2", "3/1"),
        ("(1+2)*3", "9/1"),
        ("100", "100/1"),
        ("-2", "-2/1"),
        # Decimal → rational, exact
        ("0.1", "1/10"),
        ("0.5", "1/2"),
        ("0.1+0.2", "3/10"),
        # Lowest-terms reduction
        ("6/4", "3/2"),
        ("10/5", "2/1"),
        # Unary
        ("+7", "7/1"),
        ("-(3+4)", "-7/1"),
        # Integer exponent
        ("2**10", "1024/1"),
        ("(-2)**3", "-8/1"),
        # Mixed
        ("3/10 + 0.1", "2/5"),
    ],
)
def test_positive_kats(ps, expr, canonical):
    assert ps.canonicalize(expr.encode("utf-8")) == canonical.encode("utf-8")


# --- equivalence classes preserved -----------------------------------


@pytest.mark.parametrize(
    "a,b",
    [
        ("1+2", "3"),
        ("0.1", "1/10"),
        ("0.5", "1/2"),
        ("0.1+0.2", "0.3"),
        ("(1+2)*3", "9"),
        ("6/4", "3/2"),
        ("-2", "-2/1"),
        ("2*3", "(1+2)*2"),
    ],
)
def test_equivalent_inputs_collapse(ps, a, b):
    assert ps.canonicalize(a.encode()) == ps.canonicalize(b.encode())


@pytest.mark.parametrize(
    "a,b",
    [
        ("1+2", "3"),
        ("0.1+0.2", "0.3"),
        ("6/4", "3/2"),
    ],
)
def test_equivalence_class_id_matches(ps, a, b):
    assert equivalence_class_id(ps, a.encode()) == equivalence_class_id(
        ps, b.encode()
    )


# --- equivalence classes kept distinct -------------------------------


@pytest.mark.parametrize(
    "a,b",
    [
        ("1+2", "4"),
        ("0.1", "0.2"),
        ("0.1+0.2", "0.30000000001"),  # the float-drift form is distinct
        ("2/3", "3/2"),
        ("-2", "2"),
    ],
)
def test_distinct_inputs_stay_distinct(ps, a, b):
    assert ps.canonicalize(a.encode()) != ps.canonicalize(b.encode())


# --- round-trip idempotence ------------------------------------------


@pytest.mark.parametrize(
    "expr",
    ["0", "1", "-1", "3/10", "0.1+0.2", "(1+2)*3", "2**10", "-(3+4)"],
)
def test_round_trip(ps, expr):
    assert_round_trip(ps, expr.encode("utf-8"))


def test_output_is_valid_input(ps):
    """canonicalize output is itself valid input that re-canonicalizes."""
    first = ps.canonicalize(b"6/4")
    second = ps.canonicalize(first)
    assert first == second == b"3/2"


# --- invalid-input cone ----------------------------------------------


@pytest.mark.parametrize(
    "expr,fragment",
    [
        ("x", "identifier"),
        ("x + 1", "identifier"),
        ("sqrt(2)", "function calls"),
        ("max(1, 2)", "function calls"),
        ("2 ** 0.5", "non-integer exponent"),
        ("1/0", "division by zero"),
        ("(2 + 3)/0", "division by zero"),
        ("True", "boolean"),
        ("False", "boolean"),
        ("", "empty"),
        ("   ", "empty"),
    ],
)
def test_invalid_raises(ps, expr, fragment):
    with pytest.raises(PiStarError, match=fragment):
        ps.canonicalize(expr.encode("utf-8"))


def test_syntax_error_raises(ps):
    with pytest.raises(PiStarError, match="not valid expression"):
        ps.canonicalize(b"1 + ")


def test_unsupported_op_raises(ps):
    """Comparison / bitwise / logical ops are rejected."""
    for expr in ("1 < 2", "1 & 2", "1 | 2"):
        with pytest.raises(PiStarError):
            ps.canonicalize(expr.encode("utf-8"))


def test_non_bytes_input_raises(ps):
    with pytest.raises(PiStarError, match="expects bytes"):
        ps.canonicalize("3/10")  # type: ignore[arg-type]
    with pytest.raises(PiStarError, match="expects bytes"):
        ps.canonicalize(42)  # type: ignore[arg-type]


def test_bytearray_accepted(ps):
    """bytearray is an accepted subtype per the bytes/bytearray guard."""
    assert ps.canonicalize(bytearray(b"6/4")) == b"3/2"


# --- determinism -----------------------------------------------------


def test_repeated_canonicalize_is_stable(ps):
    expr = b"(0.1 + 0.2) * 3 - 1/2 + 7/14"
    a = ps.canonicalize(expr)
    b = ps.canonicalize(expr)
    c = ps.canonicalize(expr)
    assert a == b == c