aere-research/pq-stark/test_fri_query_index.py

#!/usr/bin/env python3
# Test harness for the FRI query-index derivation sub-component (port spec section 8). Runs:
#   1. hand-computed golden vectors (worked out below and asserted),
#   2. structural invariants over many random cases,
#   3. an alternate-algorithm cross-check (a deliberately different implementation of the walk),
#   4. cross-language agreement: Python vs Node vs Java produce byte-identical results on a shared
#      vector set.
# Writes ../results/kat-results-fri-query-index.json and exits non-zero on any failure.
#
# HONEST SCOPE. This validates the constant-free index LOGIC against the public Plonky3 p3-fri spec
# and across three independent implementations. It does NOT verify any real STARK proof; the
# top-level 0x0AE8 precompile stays fail-closed. Conformance to a real SP1 v6.1.0 trace is [MEASURE]
# (see docs/AERE-STARK-VERIFIER-PORT-SPEC.md section 8 for the exact command).

import json
import os
import random
import subprocess
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

import fri_query_index_reference as ref  # noqa: E402

RESULTS = os.path.normpath(os.path.join(HERE, "..", "results", "kat-results-fri-query-index.json"))

_pass = 0
_fail = 0
_notes = []


def check(name, cond):
    global _pass, _fail
    if cond:
        _pass += 1
    else:
        _fail += 1
        print(f"FAIL {name}")


# ---- 1. hand-computed golden vectors --------------------------------------------------------------
def golden_vectors():
    # index=11 (0b1011), log_max_height=4, final poly len 0 -> 4 folding rounds.
    #   round0: 11 -> sib 10, pair 5, par 1
    #   round1:  5 -> sib  4, pair 2, par 1
    #   round2:  2 -> sib  3, pair 1, par 0
    #   round3:  1 -> sib  0, pair 0, par 1 ; final index 11>>4 = 0
    w = ref.walk(11, 4, 0)
    expect = [
        (0, 11, 10, 5, 1),
        (1, 5, 4, 2, 1),
        (2, 2, 3, 1, 0),
        (3, 1, 0, 0, 1),
    ]
    check("golden.walk.len", len(w) == 4)
    for i, e in enumerate(expect):
        s = w[i]
        check(f"golden.walk.row{i}", (s.layer, s.index, s.sibling, s.index_pair, s.parity) == e)
    check("golden.walk.final", ref.final_index(11, 4, 0) == 0)

    # index=22 (0b10110), log_max_height=5, log_final_poly_len=1 -> 4 rounds, final index 22>>4 = 1.
    w2 = ref.walk(22, 5, 1)
    expect2 = [
        (0, 22, 23, 11, 0),
        (1, 11, 10, 5, 1),
        (2, 5, 4, 2, 1),
        (3, 2, 3, 1, 0),
    ]
    check("golden.walk2.len", len(w2) == 4)
    for i, e in enumerate(expect2):
        s = w2[i]
        check(f"golden.walk2.row{i}", (s.layer, s.index, s.sibling, s.index_pair, s.parity) == e)
    check("golden.walk2.final", ref.final_index(22, 5, 1) == 1)

    # sample_bits hand goldens: 20 = 0b10100.
    check("golden.samplebits.20_2", ref.sample_bits(20, 2) == 0)
    check("golden.samplebits.20_3", ref.sample_bits(20, 3) == 4)
    check("golden.samplebits.20_4", ref.sample_bits(20, 4) == 4)
    check("golden.samplebits.20_5", ref.sample_bits(20, 5) == 20)
    check("golden.samplebits.mixed", ref.sample_bits(0x6AE81234, 5) == 0x14)


# ---- 2. structural invariants + 3. alternate-algorithm cross-check --------------------------------
def alt_walk(index, log_max_height, log_final_poly_len):
    # Deliberately different algorithm: precompute the bit decomposition of `index` and rebuild each
    # layer's value from the surviving high bits, instead of iteratively shifting. If this agrees
    # with the reference over many cases, the shifting logic is corroborated by an independent path.
    rounds = log_max_height - log_final_poly_len
    bits = [(index >> b) & 1 for b in range(log_max_height)]
    out = []
    for layer in range(rounds):
        cur = 0
        for b in range(layer, log_max_height):
            cur |= bits[b] << (b - layer)
        out.append((layer, cur, cur ^ 1, cur >> 1, cur & 1))
    return out


def invariants_and_alt():
    rnd = random.Random(0xAE8)
    for _ in range(20000):
        lmh = rnd.randint(1, 27)              # single-sample safe
        lfp = rnd.randint(0, lmh)
        idx = rnd.randint(0, (1 << lmh) - 1)
        w = ref.walk(idx, lmh, lfp)
        check("inv.rounds", len(w) == lmh - lfp)
        cur = idx
        ok = True
        for s in w:
            ok &= s.index == cur
            ok &= s.sibling == (cur ^ 1)
            ok &= bin(s.index ^ s.sibling).count("1") == 1   # differ in exactly one bit
            ok &= s.index_pair == (cur >> 1)
            ok &= s.parity == (cur & 1)
            cur >>= 1
        ok &= ref.final_index(idx, lmh, lfp) == cur
        ok &= cur < (1 << lfp)                                # collapses into final-poly domain
        check("inv.walk", ok)

        # alternate-algorithm agreement
        alt = alt_walk(idx, lmh, lfp)
        native = [(s.layer, s.index, s.sibling, s.index_pair, s.parity) for s in w]
        check("alt.walk.agree", alt == native)

        bits = rnd.randint(0, 29)
        v = rnd.randint(0, ref.BABYBEAR_P - 1)
        sb = ref.sample_bits(v, bits)
        check("inv.samplebits.range", 0 <= sb < (1 << bits))
        check("inv.samplebits.lowbits", sb == (v & ((1 << bits) - 1)))
    check("inv.samplebits.identity", ref.sample_bits(0x0ABCDEF, 28) == 0x0ABCDEF)


# ---- documented extra: the Poseidon2 S-box exponent (a field fact, not the FRI component) ----------
def sbox_exponent_note():
    # Not part of the FRI component, but a cheap, honest, offline-checkable fact the port spec (3)
    # relies on: x^7 is the Poseidon2 S-box for BabyBear because 7 is the smallest d>1 coprime to
    # p-1 = 2^27 * 3 * 5. Verify it here so the claim is grounded, not asserted.
    pm1 = ref.BABYBEAR_P - 1
    import math
    smallest = next(d for d in range(2, 64) if math.gcd(d, pm1) == 1)
    check("note.sbox_exponent_is_7", smallest == 7)
    _notes.append(f"Poseidon2 S-box exponent d = smallest d>1 with gcd(d, p-1)=1 = {smallest} (x^7)")


# ---- 4. cross-language agreement (Python vs Node vs Java) ------------------------------------------
def canonicalize(vec):
    # Normalize each language's JSON to a comparable structure (lists of tuples), ignoring key order.
    walks = []
    for wc in vec["walks"]:
        steps = [
            (s["layer"], s["index"], s["sibling"], s["index_pair"], s["parity"])
            for s in wc["steps"]
        ]
        walks.append((wc["index"], wc["logMaxHeight"], wc["logFinalPolyLen"], wc["final_index"], steps))
    samples = [(s["v"], s["bits"], s["out"]) for s in vec["samples"]]
    return {"walks": walks, "samples": samples}


def python_shared_vectors():
    walk_cases = [(11, 4, 0), (22, 5, 1), (0, 6, 0), (63, 6, 0), (12345, 20, 3), (1, 1, 0)]
    sample_cases = [(20, 2), (20, 3), (20, 5), (0x6AE81234, 5), (0x6AE81234, 20), (ref.BABYBEAR_P - 1, 10)]
    walks = []
    for idx, lmh, lfp in walk_cases:
        walks.append({
            "index": idx, "logMaxHeight": lmh, "logFinalPolyLen": lfp,
            "final_index": ref.final_index(idx, lmh, lfp),
            "steps": [
                {"layer": s.layer, "index": s.index, "sibling": s.sibling,
                 "index_pair": s.index_pair, "parity": s.parity}
                for s in ref.walk(idx, lmh, lfp)
            ],
        })
    samples = [{"v": v, "bits": b, "out": ref.sample_bits(v, b)} for v, b in sample_cases]
    return {"walks": walks, "samples": samples}


def cross_language():
    py = canonicalize(python_shared_vectors())
    langs = {"python": py}

    # Node
    node = _which("node")
    if node:
        try:
            out = subprocess.run(
                [node, os.path.join(HERE, "fri_query_index_reference.mjs")],
                capture_output=True, text=True, timeout=60, check=True,
            ).stdout.strip()
            langs["node"] = canonicalize(json.loads(out))
        except Exception as e:  # noqa: BLE001
            _notes.append(f"node cross-check unavailable: {e}")
    else:
        _notes.append("node not found on PATH; node cross-check skipped")

    # Java (compile the standalone self-test, then run it with --emit)
    javac = _which("javac")
    java = _which("java")
    if javac and java:
        try:
            subprocess.run(
                [javac, "-d", os.path.join(HERE, "out"),
                 os.path.join(HERE, "FriQueryIndexSelfTest.java")],
                capture_output=True, text=True, timeout=120, check=True,
            )
            out = subprocess.run(
                [java, "-cp", os.path.join(HERE, "out"), "FriQueryIndexSelfTest", "--emit"],
                capture_output=True, text=True, timeout=60, check=True,
            ).stdout.strip()
            langs["java"] = canonicalize(json.loads(out))
        except Exception as e:  # noqa: BLE001
            _notes.append(f"java cross-check unavailable: {e}")
    else:
        _notes.append("javac/java not found on PATH; java cross-check skipped")

    # Every available language must agree with Python exactly.
    for name, data in langs.items():
        if name == "python":
            continue
        check(f"crosslang.{name}.agree", data == py)
    _notes.append("cross-language implementations compared: " + ", ".join(sorted(langs)))
    return sorted(langs)


def _which(exe):
    from shutil import which
    return which(exe)


def main():
    golden_vectors()
    invariants_and_alt()
    sbox_exponent_note()
    langs = cross_language()

    report = {
        "component": "FRI query-index derivation (sample_bits + folding-index walk)",
        "precompile": "0x0000000000000000000000000000000000000ae8",
        "scope": "port spec section 8; ONE constant-free slice of FRI component (d)",
        "validates": "index LOGIC vs public Plonky3 p3-fri spec, across independent implementations",
        "does_not_validate": "conformance to a real SP1 v6.1.0 trace ([MEASURE]; needs Poseidon2 "
                             "challenger + exported proof)",
        "top_level_verifier": "unchanged, still fail-closed (returns EMPTY for every input)",
        "cross_language": langs,
        "flags": [
            "[VERIFY] sample_bits masks LOW bits of as_canonical_u32 (LSB) vs p3-challenger",
            "[VERIFY] arity-2 fold: sibling = index^1, index_pair = index>>1 vs p3-fri verifier.rs",
            "[VERIFY] log_max_height = log_max_degree + log_blowup; num_fold_rounds = "
            "log_max_height - log_final_poly_len for the pinned config",
            "[MEASURE] bit-reversal index-to-domain-point mapping vs a real Plonky3 trace",
            "[MEASURE] end-to-end derive_query_indices vs a real SP1 v6.1.0 inner proof",
        ],
        "notes": _notes,
        "passed": _pass,
        "failed": _fail,
        "total": _pass + _fail,
    }
    os.makedirs(os.path.dirname(RESULTS), exist_ok=True)
    with open(RESULTS, "w") as f:
        json.dump(report, f, indent=2)

    print(f"\n=== FRI query-index derivation KAT ===")
    print(f"cross-language: {', '.join(langs)}")
    for n in _notes:
        print(f"note: {n}")
    print(f"PASS={_pass} FAIL={_fail} TOTAL={_pass + _fail}")
    print(f"results -> {RESULTS}")
    if _fail != 0:
        print("FRI_QUERY_INDEX_KAT_FAILED")
        sys.exit(1)
    print("FRI_QUERY_INDEX_KAT_OK")


if __name__ == "__main__":
    main()