arborist/tests/test_anchor_prg.py

"""Tests for arborist.substrate.anchor_prg per ticket #000035 §3.2.

Coverage matches the ticket's acceptance criteria:

- Determinism: same (seed, hard_hash, dim_h) → byte-identical output.
- Distinguishing-from-random sanity: chi² test on a sample of outputs
  (loose threshold; just catches gross PRG bugs like cycling on the
  counter or HMAC mis-keying).
- Boundary: dim_h=1 and large dim_h both produce sensible outputs.
- Seed-change avalanche: flipping one bit of the seed yields a result
  uncorrelated with the original (Hamming distance ≈ output_size / 2).
- Hash-input avalanche: flipping one bit of the hard hash same.
- Input validation: short hashes / non-positive dim_h raise ValueError.
- Range invariant: every output float is in [-1, 1).
- KAT (known-answer-test) vectors regression-pinned via the
  ``bench/fixtures/phi-prg/known-answer-tests.jsonl`` fixture.
"""

from __future__ import annotations

import hashlib
import json
from pathlib import Path

import pytest

from arborist.substrate.anchor_prg import (
    PHI_PRG_VERSION,
    PLACEHOLDER_SEED,
    _bytes_to_floats,
    _expand,
    phi_prg,
)


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


def test_phi_prg_deterministic_same_inputs():
    h = hashlib.sha256(b"deterministic-input").digest()
    v1 = phi_prg(h, dim_h=64)
    v2 = phi_prg(h, dim_h=64)
    assert v1 == v2


def test_phi_prg_deterministic_with_explicit_seed():
    h = hashlib.sha256(b"x").digest()
    seed = hashlib.sha256(b"my-seed").digest()
    v1 = phi_prg(h, dim_h=32, seed=seed)
    v2 = phi_prg(h, dim_h=32, seed=seed)
    assert v1 == v2


# ----------------------------------------------------------- range invariant


def test_phi_prg_outputs_in_unit_interval():
    h = hashlib.sha256(b"range-check").digest()
    v = phi_prg(h, dim_h=2048)
    for x in v:
        # Spec: [-1, 1) — strict upper bound is critical because
        # 2 * (2^32 - 1) / 2^32 - 1 = (2^33 - 2 - 2^32) / 2^32
        # = (2^32 - 2) / 2^32 < 1.
        assert -1.0 <= x < 1.0, f"out of range: {x}"


def test_bytes_to_floats_zero_maps_to_minus_one():
    assert _bytes_to_floats(b"\x00\x00\x00\x00") == [-1.0]


def test_bytes_to_floats_max_uint32_just_below_one():
    # 0xFFFFFFFF → 2 * (1 - 2^-32) - 1 = 1 - 2^-31, strictly < 1.
    out = _bytes_to_floats(b"\xff\xff\xff\xff")
    assert len(out) == 1
    assert out[0] < 1.0
    assert out[0] > 1.0 - 1e-9


def test_bytes_to_floats_midpoint_maps_to_zero():
    # 0x80000000 → 2 * 0.5 - 1 = 0.0.
    assert _bytes_to_floats(b"\x80\x00\x00\x00") == [0.0]


# ----------------------------------------------------------- chi² sanity


def test_phi_prg_chi2_loose_uniformity():
    """Bin a 4096-float sample into 16 buckets on [-1, 1); expect
    counts within a generous chi² acceptance region. Threshold is
    deliberately wide — this catches catastrophic PRG bugs (e.g.
    counter cycling, all-zero output) but doesn't claim
    cryptographic-grade evidence.
    """
    h = hashlib.sha256(b"uniformity-sample").digest()
    n = 4096
    nbuckets = 16
    samples = phi_prg(h, dim_h=n)
    counts = [0] * nbuckets
    for x in samples:
        # Map [-1, 1) → [0, nbuckets)
        idx = int((x + 1.0) * 0.5 * nbuckets)
        if idx == nbuckets:  # x just below 1.0 due to float
            idx = nbuckets - 1
        counts[idx] += 1
    expected = n / nbuckets
    chi2 = sum((c - expected) ** 2 / expected for c in counts)
    # df = 15, 99.9th percentile ≈ 37.7. We accept up to 60 to leave
    # headroom against single-sample tail behaviour without making the
    # test useless.
    assert chi2 < 60.0, f"χ² = {chi2:.2f}; counts = {counts}"


# ----------------------------------------------------------- boundary


def test_phi_prg_dim_h_one():
    h = hashlib.sha256(b"boundary-1").digest()
    v = phi_prg(h, dim_h=1)
    assert len(v) == 1
    assert -1.0 <= v[0] < 1.0


def test_phi_prg_dim_h_large_consistent_length():
    # 2^14 = 16384 floats → 64 KB output → 1024 HMAC blocks.
    # Cheap enough for a unit test.
    h = hashlib.sha256(b"boundary-large").digest()
    v = phi_prg(h, dim_h=16384)
    assert len(v) == 16384


# ----------------------------------------------------------- avalanche


def _hamming_bits(a: bytes, b: bytes) -> int:
    assert len(a) == len(b)
    return sum(bin(x ^ y).count("1") for x, y in zip(a, b))


def test_phi_prg_seed_bit_flip_avalanches():
    """Flip one bit of the seed; expect the byte-output Hamming
    distance to be ≈ output_size_bits / 2 ± noise. This is the
    standard PRF avalanche property; failure indicates the seed
    isn't actually keying HMAC (e.g. constant collision)."""
    h = hashlib.sha256(b"avalanche-seed").digest()
    seed_a = bytes(32)
    seed_b = bytes([0x80]) + bytes(31)  # flip top bit of byte 0
    raw_a = _expand(seed_a, h, 256)
    raw_b = _expand(seed_b, h, 256)
    bits_total = len(raw_a) * 8
    diff = _hamming_bits(raw_a, raw_b)
    # 256 bytes = 2048 bits; expected ≈ 1024 ± few sigma.
    # Conservative: 35-65% of bits flipped.
    assert 0.35 * bits_total < diff < 0.65 * bits_total, (
        f"avalanche failed: {diff}/{bits_total} bits differ "
        f"({100*diff/bits_total:.1f}%)"
    )


def test_phi_prg_hash_bit_flip_avalanches():
    """Flip one bit of the hard hash input; same avalanche property."""
    seed = hashlib.sha256(b"av-seed").digest()
    h_a = bytes(32)
    h_b = bytes([0x01]) + bytes(31)
    raw_a = _expand(seed, h_a, 256)
    raw_b = _expand(seed, h_b, 256)
    bits_total = len(raw_a) * 8
    diff = _hamming_bits(raw_a, raw_b)
    assert 0.35 * bits_total < diff < 0.65 * bits_total, (
        f"avalanche failed: {diff}/{bits_total} bits differ "
        f"({100*diff/bits_total:.1f}%)"
    )


# ----------------------------------------------------------- monotonicity / closure


def test_phi_prg_output_length_exactly_dim_h():
    """``phi_prg(h, n)`` must produce exactly ``n`` floats.

    Closure check pattern (cf. fox's `test_total_equals_sum_of_three_contributions`
    in test_t3_bound_calculator.py): the output length is the
    contract; off-by-one or truncation bugs in `_expand` would
    surface here. Parametrized to widen the cone."""
    h = hashlib.sha256(b"len-check").digest()
    for n in (1, 2, 4, 7, 16, 17, 64, 1024):
        v = phi_prg(h, dim_h=n)
        assert len(v) == n, f"dim_h={n}: expected {n} floats, got {len(v)}"


def test_phi_prg_output_is_prefix_extending():
    """``phi_prg(h, n)`` must equal the first ``n`` entries of
    ``phi_prg(h, n+k)``.

    Streaming-counter-mode invariant: the HMAC-SHA-512 expansion
    is deterministic counter-based, so increasing dim_h adds
    strictly more bytes at the tail without re-deriving the
    head. A bug that re-keyed HMAC per-call (e.g. seed mutation)
    would surface here.

    Pattern from fox's `test_monotone_in_window_length`: scaling
    one input dimension while holding others fixed is a
    closed-form invariant the function MUST satisfy.
    """
    h = hashlib.sha256(b"prefix-extend").digest()
    short = phi_prg(h, dim_h=8)
    long = phi_prg(h, dim_h=24)
    assert long[:8] == short, (
        "phi_prg should be prefix-stable: "
        f"long[:8] = {long[:8][:3]!r}... vs short = {short[:3]!r}..."
    )


# ----------------------------------------------------------- hand-formula


def test_phi_prg_first_block_matches_direct_hmac():
    """First HMAC-SHA-512 block of phi_prg's output should match
    a direct hmac.new(seed, hard_hash + b'\\x00\\x00\\x00\\x00',
    sha512).digest() invocation. Hand-computed against the
    function's spec (#000035 §3.1 + module §1).

    Pattern from fox's `test_b1_exact_formula`: don't rely on
    KAT regression alone — compute the first-principles math in
    the test file and assert exact agreement. Catches algorithm
    drift that KAT regenerated against a buggy version would
    miss.
    """
    import hmac

    seed = hashlib.sha256(b"hand-formula-seed").digest()
    h = hashlib.sha256(b"hand-formula-hash").digest()

    # Spec: out = HMAC-SHA-512(seed, hard_hash || counter_be_4) for
    # counter = 0, 1, 2, ...; concatenated; truncated to dim_h * 4 bytes.
    expected_block_0 = hmac.new(
        seed, h + (0).to_bytes(4, "big"), hashlib.sha512
    ).digest()
    # First block is 64 bytes = 16 uint32s = 16 floats. dim_h=16
    # consumes exactly the first block.
    floats = phi_prg(h, dim_h=16, seed=seed)
    # Convert expected_block_0 to floats per §2.3 spec.
    expected_floats = []
    for i in range(16):
        u32 = int.from_bytes(expected_block_0[4 * i : 4 * (i + 1)], "big")
        expected_floats.append(2.0 * (u32 / 2 ** 32) - 1.0)
    for j, (got, exp) in enumerate(zip(floats, expected_floats)):
        assert got == pytest.approx(exp, abs=1e-12), (
            f"first-block float mismatch at index {j}: got {got}, "
            f"expected {exp}"
        )


def test_phi_prg_seed_changes_every_byte_independently():
    """Pattern from fox's monotone tests: scaling one input
    independently shouldn't bleed into other parts of the output.
    For phi_prg this is harder to assert directly — HMAC mixes
    everything — but we can pin: changing the seed should change
    EVERY output float (not zero of them)."""
    h = hashlib.sha256(b"seed-bleed-check").digest()
    seed_a = hashlib.sha256(b"seed-A").digest()
    seed_b = hashlib.sha256(b"seed-B").digest()
    out_a = phi_prg(h, dim_h=64, seed=seed_a)
    out_b = phi_prg(h, dim_h=64, seed=seed_b)
    # Every position should differ — under HMAC-SHA-512 PRF,
    # the probability that any specific 32-bit float matches by
    # chance is 2^-32, so 64 positions × 2^-32 ≈ 2^-26 false
    # positives expected. None expected in practice.
    matches = sum(1 for a, b in zip(out_a, out_b) if a == b)
    assert matches == 0, (
        f"{matches}/64 positions matched between different seeds; "
        "PRF bleed check failed"
    )


# ----------------------------------------------------------- validation (parametrized)


@pytest.mark.parametrize("bad_hash", [
    b"too short",                # too few bytes
    b"\x00" * 33,                # too many bytes
    b"",                         # empty
    b"\x00" * 31,                # off by one short
])
def test_phi_prg_rejects_wrong_size_hash(bad_hash):
    """Pattern from fox: parametrize over the invalid-input cone
    rather than spawn a separate test function per case."""
    with pytest.raises(ValueError, match="32 bytes"):
        phi_prg(bad_hash, dim_h=8)


@pytest.mark.parametrize("bad_dim", [0, -1, -100])
def test_phi_prg_rejects_non_positive_dim_h(bad_dim):
    h = hashlib.sha256(b"x").digest()
    with pytest.raises(ValueError, match="positive"):
        phi_prg(h, dim_h=bad_dim)


def test_phi_prg_rejects_non_bytes_hash():
    with pytest.raises(ValueError):
        phi_prg("not bytes", dim_h=8)  # type: ignore[arg-type]


def test_phi_prg_rejects_non_int_dim_h():
    h = hashlib.sha256(b"x").digest()
    with pytest.raises(ValueError):
        phi_prg(h, dim_h=8.5)  # type: ignore[arg-type]


# ----------------------------------------------------------- module shape


def test_module_exports_version_string():
    assert PHI_PRG_VERSION == "phi-prg-v1-hmac-sha512"


def test_placeholder_seed_is_32_bytes():
    assert isinstance(PLACEHOLDER_SEED, bytes)
    assert len(PLACEHOLDER_SEED) == 32


# ----------------------------------------------------------- KAT regression


KAT_FIXTURE = (
    Path(__file__).parent.parent
    / "bench"
    / "fixtures"
    / "phi-prg"
    / "known-answer-tests.jsonl"
)


@pytest.mark.skipif(
    not KAT_FIXTURE.exists(),
    reason="KAT fixture not yet generated; run scripts/generate_phi_prg_kat.py",
)
def test_phi_prg_known_answer_tests():
    """Re-run every (seed, hard_hash, dim_h) triple in the pinned KAT
    fixture; bytes-output SHA-256 must match the recorded value.

    The fixture is generated once and committed; any future change to
    the algorithm (e.g. switching from HMAC-SHA-512 to a different
    construction) MUST bump ``PHI_PRG_VERSION`` and produce a new
    fixture file under ``bench/fixtures/phi-prg/`` rather than
    overwrite this one. Old runs replay against the old fixture.
    """
    for line in KAT_FIXTURE.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        kat = json.loads(line)
        seed = bytes.fromhex(kat["seed_hex"])
        hard_hash = bytes.fromhex(kat["hard_hash_hex"])
        dim_h = int(kat["dim_h"])
        out_floats = phi_prg(hard_hash, dim_h=dim_h, seed=seed)
        # Pin the bytes-form so the fixture is independent of any
        # future float-format choice (list[float] vs array.array vs
        # numpy.ndarray). bytes-output SHA-256 is the durable
        # contract. We reconstruct the bytes via _expand to keep the
        # fixture format stable.
        raw = _expand(seed, hard_hash, dim_h * 4)
        digest = hashlib.sha256(raw).hexdigest()
        assert digest == kat["output_sha256"], (
            f"KAT mismatch for label={kat.get('label')!r}: "
            f"expected {kat['output_sha256']}, got {digest}"
        )
        # Also assert the float list has the expected length so a
        # bug in _bytes_to_floats truncation is caught.
        assert len(out_floats) == dim_h