aere-research/pq-stark/fri_verify_reference.py

#!/usr/bin/env python3
# Independent reference for the FRI VERIFIER (the low-degree test at the heart of the STARK),
# component (d) of the PQ STARK-verify precompile 0x0AE8 (direct SP1/Plonky3 inner FRI/STARK verify).
# See docs/AERE-STARK-VERIFIER-PORT-SPEC.md section 5 and spec-fri-babybear.md.
#
# ============================ HONEST SCOPE (read first) ============================
# This module implements the FRI query verification (verify_query / verify_challenges) that Plonky3
# p3-fri performs: for each query, walk the commit-phase folding, checking (1) the MMCS opening of the
# sibling pair at each layer against that layer's commitment (component (c), CONFIRMED), and (2) that
# the arity-2 folding relation holds layer to layer (interpolate the pair through beta at the coset
# point, in F_{p^4}, component (a)), down to a single constant that must equal the committed final
# polynomial. It is CONFIRMED-FROM-SOURCE against the pinned Plonky3 p3-fri 0.4.3-succinct, and a real
# known-answer test PASSES (see conformance() and test_fri_verify.py): the ground-truth extractor
# (pq-stark/fri-extractor) ran the pinned p3-fri PROVER to emit real FRI proofs over known low-degree
# polynomials and confirmed the pinned p3-fri VERIFIER accepts them; this reference reproduces the
# accept and rejects tampered proofs (corrupted opening, wrong fold, non-matching final poly).
#
# WHAT IS AND IS NOT PORTED (the honest boundary):
#   - PORTED + CONFIRMED here (component (d)): the FOLD + OPENING relations, i.e. verify_query. It takes
#     the transcript-derived betas and query indices (and the reduced openings) as INPUTS.
#   - NOT ported (component (f), the duplex-sponge Fiat-Shamir challenger): deriving those betas / query
#     indices / grinding IN-CIRCUIT from the transcript. In this KAT the betas and indices come from the
#     pinned p3-fri challenger (the ground-truth extractor), supplied as inputs; the transcript BINDING
#     that would make them trustless is component (f) and remains un-ported (Challenger.spongePorted =
#     false). So even with this passing FRI-relation KAT, the top-level 0x0AE8 stays FAIL-CLOSED.
#   - The proof-of-work / grinding check is part of the transcript phase (component (f)); it is not part
#     of the fold+opening relation this module verifies.
#
# Source (pinned, checksum-matched to the repo Cargo.lock):
#   p3-fri        0.4.3-succinct  5cbc4965ee488f3247867b7ec4bb005b8afa72cb0d461a4dcb1387ecab6426d5
#   p3-challenger 0.4.3-succinct  b6a908924d43e4cfb93fb41c8346cac211b70314385a9037e9241f5b7f3eaf77
#   p3-dft        0.4.3-succinct  be6408b10a2c27eb13a7d5580c546c2179a8dc7dbc10a990657311891f9b41c0
#   (commit-phase MMCS: p3-commit ExtensionMmcs over the CONFIRMED p3-merkle-tree ValMmcs, component (c);
#    leaf hash/compress: the CONFIRMED Poseidon2-BabyBear permutation, component (b)).

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      # component (a): F_p and F_{p^4} arithmetic (CONFIRMED)
import mmcs_babybear_reference as M       # component (c): FieldMerkleTreeMmcs verify_batch (CONFIRMED)

GROUND_TRUTH = os.path.join(HERE, "fri_ground_truth.json")

P = F.P


# ============================ index arithmetic ============================

def reverse_bits_len(x: int, bit_len: int) -> int:
    """Reverse the low `bit_len` bits of x. Mirrors p3-util reverse_bits_len (x.reverse_bits() >>
    (usize::BITS - bit_len)). Used to map a query index to its coset point exponent."""
    r = 0
    for i in range(bit_len):
        r = (r << 1) | ((x >> i) & 1)
    return r


# ============================ the FRI query verifier ============================
# Elements of F_{p^4} are 4-int lists [c0, c1, c2, c3] (component (a)'s convention, matching Plonky3
# BinomialExtensionField<BabyBear,4>::as_base_slice). A digest is 8 canonical BabyBear ints.

def verify_query(log_blowup, log_max_height, commits, betas, index, ro_full, layers):
    """Reproduce p3-fri verifier::verify_query for one query. Returns the folded evaluation in
    F_{p^4} (a 4-int list) on success, or None if any commit-phase MMCS opening fails.

    Arguments (all supplied as inputs; the transcript derivation of betas/index is component (f),
    un-ported):
      log_blowup      : FRI log blowup (num_fold_rounds = log_max_height - log_blowup).
      log_max_height  : log2 of the first-layer (LDE) domain.
      commits         : the commit-phase MMCS roots, one per fold layer (each 8 BabyBear ints).
      betas           : the folding challenges, one per layer (each F_{p^4}).
      index           : this query's index in [0, 2^log_max_height).
      ro_full         : the reduced openings array indexed by log-height (len >= log_max_height+1),
                        each an F_{p^4} element (mostly zero; ro_full[h] injected before the fold at
                        log_folded_height = h-1).
      layers          : per fold layer, {"sibling_value": F_{p^4}, "opening_proof": [digest,...]}.
    """
    folded = [0, 0, 0, 0]  # F_{p^4} zero
    g = F.two_adic_generator(log_max_height)                      # base-field two-adic generator
    x = F.ext_from_base(F.pow_(g, reverse_bits_len(index, log_max_height)))  # coset point, in F_{p^4}
    gen1 = F.ext_from_base(F.two_adic_generator(1))              # order-2 root = -1, embedded
    idx = index
    num_layers = log_max_height - log_blowup
    for layer in range(num_layers):
        log_folded_height = log_max_height - 1 - layer
        # inject the reduced opening at this height (single-codeword LDT: only ro_full[log_max_height])
        folded = F.ext_add(folded, ro_full[log_folded_height + 1])

        index_sibling = idx ^ 1
        index_pair = idx >> 1

        # evals = [folded, folded]; the sibling slot is replaced by the opened sibling value.
        evals = [list(folded), list(folded)]
        evals[index_sibling % 2] = list(layers[layer]["sibling_value"])

        # MMCS opening: the pair of F_{p^4} evals flattens to 8 BabyBear coords (ExtensionMmcs), a
        # single width-8 row of a height-2^log_folded_height tree, checked by the CONFIRMED base MMCS.
        row = evals[0] + evals[1]  # 8 base ints
        height = 1 << log_folded_height
        if not M.verify_batch(commits[layer], [(8, height)], index_pair, [row], layers[layer]["opening_proof"]):
            return None  # commit-phase MMCS opening failed -> reject

        # xs = [x, x]; the sibling's x is negated (x and -x are the arity-2 pair).
        xs = [x, F.ext_mul(x, gen1)] if index_sibling % 2 == 1 else [F.ext_mul(x, gen1), x]
        # interpolate the pair through beta and evaluate:
        #   folded = evals[0] + (beta - xs[0]) * (evals[1] - evals[0]) / (xs[1] - xs[0])
        beta = betas[layer]
        num = F.ext_mul(F.ext_sub(beta, xs[0]), F.ext_sub(evals[1], evals[0]))
        den = F.ext_sub(xs[1], xs[0])
        folded = F.ext_add(evals[0], F.ext_mul(num, F.ext_inv(den)))

        idx = index_pair
        x = F.ext_mul(x, x)  # x = x^2
    return folded


def build_ro_full(log_max_height, ro_top):
    """Construct the reduced-openings array for a single top-codeword LDT: zero everywhere except
    ro_full[log_max_height] = ro_top (the codeword value at the query index)."""
    ro = [[0, 0, 0, 0] for _ in range(log_max_height + 2)]
    ro[log_max_height] = list(ro_top)
    return ro


def verify_case(case, tamper=None):
    """Run verify_challenges over every query of a case: accept iff every query's MMCS openings pass
    AND every query's folded value equals the committed final polynomial. `tamper` injects a fault to
    demonstrate rejection: 'sibling' (corrupt an opened value), 'proof' (corrupt an MMCS sibling
    digest), 'beta' (wrong fold challenge), 'final' (non-matching final polynomial)."""
    log_blowup = case["log_blowup"]
    log_max_height = case["log_max_height"]
    commits = case["commit_phase_commits"]
    final_poly = list(case["final_poly"])
    betas = [list(b) for b in case["betas"]]

    if tamper == "final":
        final_poly[0] = (final_poly[0] + 1) % P
    if tamper == "beta":
        betas[0][0] = (betas[0][0] + 1) % P

    for q in case["queries"]:
        layers = [
            {"sibling_value": list(s["sibling_value"]),
             "opening_proof": [list(d) for d in s["opening_proof"]]}
            for s in q["layers"]
        ]
        if tamper == "sibling":
            layers[0]["sibling_value"][0] = (layers[0]["sibling_value"][0] + 1) % P
        if tamper == "proof":
            layers[0]["opening_proof"][0][0] = (layers[0]["opening_proof"][0][0] + 1) % P

        ro_full = build_ro_full(log_max_height, q["ro_top"])
        folded = verify_query(log_blowup, log_max_height, commits, betas, q["index"], ro_full, layers)
        if folded is None:
            return False   # a commit-phase MMCS opening failed
        if folded != final_poly:
            return False   # FinalPolyMismatch
    return True


# ============================ conformance KAT (real ground truth) ============================

def load_ground_truth():
    with open(GROUND_TRUTH) as f:
        return json.load(f)


def conformance_kats():
    """Return (passed, total) after: (0) a perm-sanity + two-adic-generator sanity check tying the
    ground truth to the CONFIRMED Poseidon2 permutation and component (a)'s generators; (1) accepting
    every genuine FRI proof; (2) rejecting each of four tamper variants per case."""
    gt = load_ground_truth()
    passed = 0
    total = 0

    # (0) sanity: the extractor's perm and generators match the confirmed sub-components.
    total += 1
    if M.permute([0] * 16) == gt["perm_zeros"]:
        passed += 1
    for b in range(28):
        total += 1
        if F.two_adic_generator(b) == gt["two_adic_generators"][b]:
            passed += 1

    # (1) accept genuine proofs; (2) reject tampered ones.
    for case in gt["cases"]:
        total += 1
        if verify_case(case):
            passed += 1
        for tamper in ("sibling", "proof", "beta", "final"):
            total += 1
            if not verify_case(case, tamper):
                passed += 1
    return passed, total


# ============================ shared cross-language vector set ============================
# All three languages verify the same ground truth and emit, per case: the accept flag, the four
# tamper-reject flags, and the per-query FOLDED evaluation (F_{p^4}). The harness asserts byte-identical
# folded values across Python/Node/Java (three independent implementations of the fold arithmetic) AND
# that every accept/reject flag holds.

def _case_folded(case):
    log_blowup = case["log_blowup"]
    log_max_height = case["log_max_height"]
    commits = case["commit_phase_commits"]
    betas = [list(b) for b in case["betas"]]
    folded_per_query = []
    for q in case["queries"]:
        ro_full = build_ro_full(log_max_height, q["ro_top"])
        folded = verify_query(log_blowup, log_max_height, commits, betas, q["index"], ro_full, q["layers"])
        folded_per_query.append(folded)
    return folded_per_query


def shared_vectors():
    gt = load_ground_truth()
    cases = []
    for case in gt["cases"]:
        cases.append({
            "name": case["name"],
            "logBlowup": case["log_blowup"],
            "logMaxHeight": case["log_max_height"],
            "numQueries": case["num_queries"],
            "folded": _case_folded(case),
            "finalPoly": list(case["final_poly"]),
            "accept": verify_case(case),
            "rejectSibling": not verify_case(case, "sibling"),
            "rejectProof": not verify_case(case, "proof"),
            "rejectBeta": not verify_case(case, "beta"),
            "rejectFinal": not verify_case(case, "final"),
        })
    return {
        "permZeros": M.permute([0] * 16),
        "twoAdicGenerators": [F.two_adic_generator(b) for b in range(28)],
        "cases": cases,
    }


if __name__ == "__main__":
    json.dump(shared_vectors(), sys.stdout)