aere-research/pq-stark/stark_e2e_reference.py

#!/usr/bin/env python3
# ASSEMBLED END-TO-END BabyBear + FRI STARK verifier for the PQ STARK-verify precompile 0x0AE8 reference.
#
# ============================ WHAT THIS IS (read first) ============================
# The six generic components of a BabyBear+FRI STARK verifier were each conformance-confirmed INDIVIDUALLY
# against the pinned Plonky3 0.4.3-succinct crates (field, Poseidon2, MMCS, FRI verify_query, duplex
# challenger, generic AIR quotient check). THIS module wires all six together into ONE full p3-uni-stark
# verify pipeline and runs it, end to end, against a REAL proof emitted by the pinned p3-uni-stark PROVER
# (e2e-extractor -> e2e_ground_truth.json), which the pinned p3-uni-stark VERIFIER accepts (the ground
# truth). It asserts the assembled pipeline ACCEPTS the genuine proof and REJECTS tampered ones (a trace
# opening, a FRI layer, an input-batch opening, the public values, the final poly).
#
# The pipeline follows p3-uni-stark verifier.rs + p3-fri two_adic_pcs.rs::verify + p3-fri verifier.rs
# VERBATIM (pinned 0.4.3-succinct):
#   1. observe(trace commit); alpha_c = sample_ext           (constraint challenge)
#   2. observe(quotient commit); zeta = sample_ext; zeta_next = zeta * g_trace
#   3. PCS.verify:
#        alpha_fri = sample_ext                              (FRI batch-combining challenge)
#        verify_shape_and_sample_challenges: per commit-phase commit observe+sample beta; observe final
#          poly; check grinding PoW; sample num_queries query indices
#        per query: reduced-opening combination sum_i alpha_fri^i (p(x)-p(z))/(x-z) over the input-batch
#          MMCS openings (input-batch root recomputed by the CONFIRMED MMCS verify_batch), then
#          verify_query (the CONFIRMED FRI fold + commit-phase MMCS opening) and folded == final_poly
#   4. AIR quotient consistency: folded_constraints(zeta) * inv_zeroifier == quotient(zeta)
#
# HONEST SCOPE. The end-to-end proof verified here is a p3-uni-stark EXAMPLE AIR (Fibonacci / degree-3
# multiply) under an EXAMPLE FRI config (log_blowup=2, num_queries=28, pow_bits=8), NOT an Aere production
# circuit and NOT an SP1 6.1.0 proof. Aere's own zk-circuits are SP1 guest programs, and SP1 6.1.0 is
# Hypercube (KoalaBear multilinear: BaseFold/Jagged + sumcheck-zerocheck + LogUp-GKR), which this
# BabyBear+FRI pipeline does not target. This module does NOT run in the precompile and does NOT make
# 0x0AE8 accept anything: the top level stays FAIL-CLOSED (SP1_RECURSION_AIR_PORTED=false). See the port
# spec and README for exactly what an Aere-own / SP1 end-to-end KAT still needs.

import json
import os
import sys

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

import babybear_field_reference as F        # (a) F_p and F_{p^4}                         CONFIRMED
import poseidon2_babybear_reference as P2    # (b) Poseidon2 permutation                    CONFIRMED
import mmcs_babybear_reference as M          # (c) FieldMerkleTreeMmcs verify_batch         CONFIRMED
import challenger_reference as CH            # (f) DuplexChallenger + grinding              CONFIRMED
import fri_verify_reference as FRI           # (d) FRI verify_query (fold + opening)        CONFIRMED
import air_quotient_reference as AQ          # (e) generic AIR quotient-consistency         CONFIRMED

GROUND_TRUTH = os.path.join(HERE, "e2e_ground_truth.json")
P = F.P


def _log2_exact(n):
    b = n.bit_length() - 1
    assert (1 << b) == n, f"{n} is not a power of two"
    return b


def verify_stark(case, tamper=None):
    """Full assembled BabyBear+FRI STARK verify of a real p3-uni-stark proof. Returns True iff every
    stage passes. `tamper` in {None, 'trace_open', 'fri_layer', 'input_batch', 'public_values',
    'final_poly'} injects a single fault to demonstrate rejection."""
    c = json.loads(json.dumps(case))  # deep copy so tampers do not mutate the shared ground truth

    log_blowup = c["log_blowup"]
    num_queries = c["num_queries"]
    pow_bits = c["pow_bits"]
    degree_bits = c["degree_bits"]
    air = c["air"]
    pis = list(c["public_values"])

    trace_commit = list(c["commitments"]["trace"])
    quotient_commit = list(c["commitments"]["quotient_chunks"])
    trace_local = [list(v) for v in c["opened_values"]["trace_local"]]
    trace_next = [list(v) for v in c["opened_values"]["trace_next"]]
    quotient_chunks = [[list(v) for v in ch] for ch in c["opened_values"]["quotient_chunks"]]

    fp = c["opening_proof"]["fri_proof"]
    commit_phase_commits = [list(x) for x in fp["commit_phase_commits"]]
    final_poly = list(fp["final_poly"])
    pow_witness = fp["pow_witness"]
    query_proofs = fp["query_proofs"]
    query_openings = c["opening_proof"]["query_openings"]

    # ---- tampers that live in the claimed openings / public inputs ----
    if tamper == "public_values":
        pis[-1] = (pis[-1] + 1) % P
    if tamper == "trace_open":
        trace_local[0][0] = (trace_local[0][0] + 1) % P
    if tamper == "final_poly":
        final_poly[0] = (final_poly[0] + 1) % P

    quotient_degree = len(quotient_chunks)
    log_max_height = len(commit_phase_commits) + log_blowup  # = log_global_max_height

    # ---- 1) transcript: observe trace, sample constraint alpha ----
    ch = CH.DuplexChallenger()
    ch.observe_slice(trace_commit)
    alpha_c = ch.sample_ext()

    # ---- 2) observe quotient, sample DEEP point zeta ----
    ch.observe_slice(quotient_commit)
    zeta = ch.sample_ext()
    g_trace = F.two_adic_generator(degree_bits)
    zeta_next = F.ext_mul(zeta, F.ext_from_base(g_trace))

    # ---- 3) PCS.verify ----
    alpha_fri = ch.sample_ext()  # batch-combining challenge (distinct from the constraint alpha)

    # verify_shape_and_sample_challenges (p3-fri verifier.rs)
    betas = []
    for comm in commit_phase_commits:
        ch.observe_slice(comm)
        betas.append(ch.sample_ext())
    ch.observe_ext(final_poly)
    if len(query_proofs) != num_queries:
        return False
    if not ch.check_witness(pow_bits, pow_witness):
        return False
    query_indices = [ch.sample_bits(log_max_height) for _ in range(num_queries)]

    # rounds structure (p3-uni-stark verifier.rs pcs.verify call):
    #   batch 0 = trace   : 1 matrix of size 2^degree_bits, opened at [(zeta, trace_local), (zeta_next, trace_next)]
    #   batch 1 = quotient : quotient_degree matrices of size 2^degree_bits, opened at [(zeta, chunk_i)]
    trace_size = 1 << degree_bits
    rounds = [
        (trace_commit, [(trace_size, [(zeta, trace_local), (zeta_next, trace_next)])]),
        (quotient_commit, [(trace_size, [(zeta, quotient_chunks[i])]) for i in range(quotient_degree)]),
    ]

    # verify_challenges: for each query, build the reduced openings, then FRI verify_query.
    for qi in range(num_queries):
        index = query_indices[qi]
        per_query_batches = query_openings[qi]
        if len(per_query_batches) != len(rounds):
            return False

        # ro / alpha_pow indexed by log_height (p3-fri two_adic_pcs.rs::verify)
        ro = [[0, 0, 0, 0] for _ in range(log_max_height + 2)]
        alpha_pow = [F.ext_from_base(1) for _ in range(log_max_height + 2)]

        for bi, (batch_commit, mats) in enumerate(rounds):
            batch = per_query_batches[bi]
            opened_values = [list(r) for r in batch["opened_values"]]     # per matrix: base LDE row
            opening_proof = [list(d) for d in batch["opening_proof"]]     # sibling digest path

            if tamper == "input_batch" and qi == 0 and bi == 0:
                opened_values[0][0] = (opened_values[0][0] + 1) % P

            batch_heights = [ms << log_blowup for (ms, _) in mats]
            batch_dims = [(0, h) for h in batch_heights]  # MMCS ignores width; heights drive the tree
            log_batch_max_height = _log2_exact(max(batch_heights))
            bits_reduced = log_max_height - log_batch_max_height
            reduced_index = index >> bits_reduced

            # input-batch MMCS opening (component (c)): recompute the batch root, compare to batch_commit.
            if not M.verify_batch(batch_commit, batch_dims, reduced_index, opened_values, opening_proof):
                return False

            for mat_opening, (ms, points) in zip(opened_values, mats):
                log_height = _log2_exact(ms) + log_blowup
                br2 = log_max_height - log_height
                rev = FRI.reverse_bits_len(index >> br2, log_height)
                x = F.mul(F.GENERATOR, F.pow_(F.two_adic_generator(log_height), rev))  # LDE coset point (base)
                x_ext = F.ext_from_base(x)
                for z, ps_at_z in points:
                    if len(mat_opening) != len(ps_at_z):
                        return False
                    for p_at_x, p_at_z in zip(mat_opening, ps_at_z):
                        # quotient = (p_at_x - p_at_z) / (x - z), in F_{p^4}
                        num = F.ext_sub(F.ext_from_base(p_at_x), list(p_at_z))
                        den = F.ext_sub(x_ext, z)
                        q = F.ext_mul(num, F.ext_inv(den))
                        ro[log_height] = F.ext_add(ro[log_height], F.ext_mul(alpha_pow[log_height], q))
                        alpha_pow[log_height] = F.ext_mul(alpha_pow[log_height], alpha_fri)

        # FRI verify_query over the CONFIRMED fold + commit-phase MMCS opening (components (d),(c)).
        layers = []
        for st in query_proofs[qi]["commit_phase_openings"]:
            layers.append({
                "sibling_value": list(st["sibling_value"]),
                "opening_proof": [list(d) for d in st["opening_proof"]],
            })
        if tamper == "fri_layer" and qi == 0:
            layers[0]["sibling_value"][0] = (layers[0]["sibling_value"][0] + 1) % P

        folded = FRI.verify_query(log_blowup, log_max_height, commit_phase_commits, betas,
                                  index, ro, layers)
        if folded is None:
            return False
        if folded != final_poly:
            return False

    # ---- 4) AIR quotient consistency (component (e)) with the constraint alpha + zeta ----
    is_first, is_last, is_transition, inv_zeroifier = AQ.selectors_at_point(degree_bits, zeta)
    quotient = AQ.reconstruct_quotient(degree_bits, quotient_chunks, zeta)
    folder = AQ.AIR_FOLDERS[air]
    folded_constraints = folder(trace_local, trace_next, pis,
                                is_first, is_last, is_transition, alpha_c)
    if not F.ext_eq(F.ext_mul(folded_constraints, inv_zeroifier), quotient):
        return False

    return True


TAMPERS = ("trace_open", "fri_layer", "input_batch", "public_values", "final_poly")


def conformance_kats():
    """(passed, total, details). (0) tie ground truth to the CONFIRMED Poseidon2 permutation + Val
    generator; (1) ACCEPT each genuine library-accepted proof end to end; (2) REJECT each tamper."""
    with open(GROUND_TRUTH) as fp:
        gt = json.load(fp)

    passed = 0
    total = 0
    details = []

    def check(name, cond):
        nonlocal passed, total
        total += 1
        if cond:
            passed += 1
        details.append({"check": name, "pass": bool(cond)})

    # (0) sanity: same permutation + generator the CONFIRMED components use.
    check("perm_zeros_matches_poseidon2", P2.permute([0] * 16) == gt["perm_zeros"])
    check("val_generator_is_31", gt["val_generator"] == F.GENERATOR == 31)

    for case in gt["cases"]:
        name = case["name"]
        check(f"{name}.library_accept", case["library_accept"] is True)
        # (1) genuine proof accepts end to end.
        check(f"{name}.e2e_accept", verify_stark(case) is True)
        # (2) each tamper is rejected.
        for t in TAMPERS:
            check(f"{name}.reject[{t}]", verify_stark(case, t) is False)

    return passed, total, details


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "kat":
        p, t, det = conformance_kats()
        for d in det:
            print(("PASS" if d["pass"] else "FAIL"), d["check"])
        print(f"conformance: PASS={p} FAIL={t - p} TOTAL={t}")
        sys.exit(0 if p == t else 1)
    # default: emit the accept/reject matrix as JSON
    with open(GROUND_TRUTH) as fp:
        gt = json.load(fp)
    out = {"cases": []}
    for case in gt["cases"]:
        out["cases"].append({
            "name": case["name"],
            "accept": verify_stark(case),
            "rejects": {t: (not verify_stark(case, t)) for t in TAMPERS},
        })
    json.dump(out, sys.stdout, indent=2)