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}%)"
    )


# ----------------------------------------------------------- validation


def test_phi_prg_rejects_short_hash():
    with pytest.raises(ValueError, match="32 bytes"):
        phi_prg(b"too short", dim_h=8)


def test_phi_prg_rejects_long_hash():
    with pytest.raises(ValueError, match="32 bytes"):
        phi_prg(b"\x00" * 33, dim_h=8)


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_zero_dim_h():
    h = hashlib.sha256(b"x").digest()
    with pytest.raises(ValueError, match="positive"):
        phi_prg(h, dim_h=0)


def test_phi_prg_rejects_negative_dim_h():
    h = hashlib.sha256(b"x").digest()
    with pytest.raises(ValueError, match="positive"):
        phi_prg(h, dim_h=-1)


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