aere-research/pq-stark/air_quotient_reference.py

#!/usr/bin/env python3
# Independent reference for the GENERIC AIR constraint / quotient-consistency check (component (e)) that
# a p3-uni-stark STARK verifier performs, for the PQ STARK-verify precompile 0x0AE8 (direct SP1/Plonky3
# inner FRI/STARK verify). See docs/AERE-STARK-VERIFIER-PORT-SPEC.md section 6 and spec-stark-air.md.
#
# ============================ HONEST SCOPE (read first) ============================
# This module implements the GENERIC quotient-consistency MECHANISM that p3-uni-stark's verify() runs
# AFTER the PCS opening argument (verifier.rs lines 90-141): given the out-of-domain trace openings at
# zeta (trace_local) and at g*zeta (trace_next), a random challenge alpha, the quotient-chunk openings,
# and the trace-domain degree, it (1) folds the AIR's constraints with alpha into a single value
# folded_constraints(zeta) (Horner: acc = acc*alpha + constraint), (2) reconstructs quotient(zeta) from
# the chunk openings and the split-domain normalization zps, (3) computes the vanishing polynomial
# Z_H(zeta) and the first/last/transition selectors, and (4) checks the identity
#     folded_constraints(zeta) == Z_H(zeta) * quotient(zeta)      (equivalently  folded * inv_zeroifier == quotient),
# all over F_{p^4}. It is CONFIRMED-FROM-SOURCE against the pinned Plonky3 p3-uni-stark 0.4.3-succinct
# verifier.rs / p3-commit domain.rs / p3-air folder.rs, and a real known-answer test PASSES (see
# conformance() and test_air_quotient.py): the ground-truth extractor (pq-stark/airquotient-extractor)
# ran the pinned p3-uni-stark PROVER + VERIFIER on two KNOWN example AIRs (the Fibonacci AIR from
# Plonky3's own tests/fib_air.rs, and a degree-3 multiply AIR matching tests/mul_air.rs), confirmed the
# library ACCEPTS, and emitted alpha/zeta/openings; this reference reproduces the accept and rejects
# tampered inputs (a corrupted trace opening, a wrong quotient chunk, a wrong alpha).
#
# WHAT IS AND IS NOT PORTED (the honest boundary):
#   - PORTED + CONFIRMED here (component (e) GENERIC mechanism): the constraint-fold + quotient
#     reconstruction + Z_H + selectors + the identity check, for a SUPPLIED AIR constraint evaluator.
#     The constraint evaluators for the two example AIRs are transcribed BY HAND from their eval().
#   - NOT ported (the SPECIFIC piece that keeps the precompile fail-closed): the SP1 RECURSION AIR (its
#     exact multi-thousand-constraint set, interactions/permutation argument, public-value layout) and
#     the verifying-key digest that commits to it. That is a large, program-specific, multi-week port
#     needing the SP1 toolchain, and is NOT here. So this passing GENERIC KAT does NOT let the precompile
#     verify a real SP1 proof: the top-level 0x0AE8 STAYS FAIL-CLOSED for real proofs (see the precompile
#     StarkConstraints.SP1_RECURSION_AIR_PORTED = false gate).
#
# Source (pinned, checksum-matched to the repo Cargo.lock):
#   p3-uni-stark  0.4.3-succinct  verifier.rs  (the quotient-consistency check)
#   p3-commit     0.4.3-succinct  domain.rs    (TwoAdicMultiplicativeCoset selectors / zp_at_point / split)
#   p3-air        0.4.3-succinct  folder.rs, air.rs (VerifierConstraintFolder Horner fold; when_* filters)

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)

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

P = F.P
ONE = F.ext_from_base(1)


# ============================ domain / selector math (p3-commit domain.rs) ============================
# Elements of F_{p^4} are 4-int lists [c0,c1,c2,c3] (matching BinomialExtensionField<BabyBear,4>).

def ext_exp_power_of_2(x, log_n):
    """x^(2^log_n) in F_{p^4} by repeated squaring (Field::exp_power_of_2)."""
    r = list(x)
    for _ in range(log_n):
        r = F.ext_mul(r, r)
    return r


def ext_div(a, b):
    return F.ext_mul(a, F.ext_inv(b))


def zp_at_point(log_n, shift_base, point_ext):
    """Z_D(point) for a TwoAdicMultiplicativeCoset D of size 2^log_n and (base-field) shift:
    (point * shift^-1)^(2^log_n) - 1. Mirrors domain.rs zp_at_point."""
    shift_inv = F.ext_from_base(F.inv(shift_base))
    return F.ext_sub(ext_exp_power_of_2(F.ext_mul(point_ext, shift_inv), log_n), ONE)


def selectors_at_point(degree_bits, zeta):
    """domain.rs selectors_at_point for the trace domain (shift = 1, log_n = degree_bits).
    Returns (is_first_row, is_last_row, is_transition, inv_zeroifier), all F_{p^4}."""
    g = F.two_adic_generator(degree_bits)
    g_inv = F.ext_from_base(F.inv(g))
    unshifted = list(zeta)                       # shift = 1
    z_h = F.ext_sub(ext_exp_power_of_2(unshifted, degree_bits), ONE)
    is_first = ext_div(z_h, F.ext_sub(unshifted, ONE))
    is_last = ext_div(z_h, F.ext_sub(unshifted, g_inv))
    is_transition = F.ext_sub(unshifted, g_inv)
    inv_zeroifier = F.ext_inv(z_h)
    return is_first, is_last, is_transition, inv_zeroifier


def monomial(e_i):
    """The F_{p^4} basis element x^{e_i} (Challenge::monomial): [0,..,1 at e_i,..,0]."""
    m = [0, 0, 0, 0]
    m[e_i] = 1
    return m


def reconstruct_quotient(degree_bits, quotient_chunks, zeta):
    """Reconstruct quotient(zeta) from the chunk openings (verifier.rs lines 90-116).

    quotient_domain: log_n = degree_bits + log_quotient_degree, shift = GENERATOR (base).
    chunk i domain: log_n = degree_bits, shift = GENERATOR * gen_q^i, gen_q = two_adic_generator(log_n_q).
    zps[i] = prod_{j != i} zp_j(zeta) * zp_j(first_point_i)^-1.
    quotient = sum_i zps[i] * sum_{e} monomial(e) * chunk[i][e]."""
    quotient_degree = len(quotient_chunks)
    log_quotient_degree = quotient_degree.bit_length() - 1
    assert (1 << log_quotient_degree) == quotient_degree, "quotient_degree must be a power of two"
    log_n_q = degree_bits + log_quotient_degree
    gen_q = F.two_adic_generator(log_n_q)
    # chunk-domain base-field shifts
    shifts = [F.mul(F.GENERATOR, F.pow_(gen_q, i)) for i in range(quotient_degree)]

    zps = []
    for i in range(quotient_degree):
        acc = list(ONE)
        first_point_i = F.ext_from_base(shifts[i])          # domain_i.first_point() = shift_i
        for j in range(quotient_degree):
            if j == i:
                continue
            num = zp_at_point(degree_bits, shifts[j], zeta)
            den = zp_at_point(degree_bits, shifts[j], first_point_i)
            acc = F.ext_mul(acc, F.ext_mul(num, F.ext_inv(den)))
        zps.append(acc)

    quotient = [0, 0, 0, 0]
    for ch_i, ch in enumerate(quotient_chunks):
        for e_i, c in enumerate(ch):
            term = F.ext_mul(F.ext_mul(zps[ch_i], monomial(e_i)), c)
            quotient = F.ext_add(quotient, term)
    return quotient


# ============================ example AIR constraint evaluators (transcribed by hand) ============
# Each returns folded_constraints(zeta): the AIR's constraints folded with alpha via Horner, exactly as
# VerifierConstraintFolder.assert_zero does (acc = acc*alpha + constraint), in the eval() emission order.

def _horner(constraints, alpha):
    acc = [0, 0, 0, 0]
    for c in constraints:
        acc = F.ext_add(F.ext_mul(acc, alpha), c)
    return acc


def fold_fibonacci(tl, tn, pis, sel_first, sel_last, sel_trans, alpha):
    """tests/fib_air.rs eval(): columns [left,right], public values [a,b,x]. Constraint order:
    is_first*(left-a), is_first*(right-b), is_trans*(right-next.left),
    is_trans*(left+right-next.right), is_last*(right-x)."""
    a = F.ext_from_base(pis[0])
    b = F.ext_from_base(pis[1])
    x = F.ext_from_base(pis[2])
    left, right = tl[0], tl[1]
    nleft, nright = tn[0], tn[1]
    c1 = F.ext_mul(sel_first, F.ext_sub(left, a))
    c2 = F.ext_mul(sel_first, F.ext_sub(right, b))
    c3 = F.ext_mul(sel_trans, F.ext_sub(right, nleft))
    c4 = F.ext_mul(sel_trans, F.ext_sub(F.ext_add(left, right), nright))
    c5 = F.ext_mul(sel_last, F.ext_sub(right, x))
    return _horner([c1, c2, c3, c4, c5], alpha)


def fold_mul_deg3(tl, tn, pis, sel_first, sel_last, sel_trans, alpha):
    """tests/mul_air.rs eval() (REPETITIONS=1, degree=3): columns [a,b,c]. Constraint order:
    a^2*b - c (all rows), is_first*((a*a+1)-b), is_trans*((a+1)-next_a)."""
    a, b, c = tl[0], tl[1], tl[2]
    next_a = tn[0]
    c1 = F.ext_sub(F.ext_mul(F.ext_mul(a, a), b), c)                     # a^2*b - c
    c2 = F.ext_mul(sel_first, F.ext_sub(F.ext_add(F.ext_mul(a, a), ONE), b))   # is_first*(a*a+1 - b)
    c3 = F.ext_mul(sel_trans, F.ext_sub(F.ext_add(a, ONE), next_a))      # is_trans*(a+1 - next_a)
    return _horner([c1, c2, c3], alpha)


AIR_FOLDERS = {
    "fibonacci": fold_fibonacci,
    "mul_deg3": fold_mul_deg3,
}


# ============================ the generic quotient-consistency check ============================

def check_identity(case):
    """Reproduce p3-uni-stark verify()'s quotient-consistency check for one case. Returns True iff
    folded_constraints(zeta) * inv_zeroifier == quotient(zeta) (i.e. == Z_H(zeta)*quotient(zeta))."""
    degree_bits = case["degree_bits"]
    alpha = list(case["alpha"])
    zeta = list(case["zeta"])
    pis = list(case["public_values"])
    tl = [list(v) for v in case["trace_local"]]
    tn = [list(v) for v in case["trace_next"]]
    chunks = [[list(v) for v in ch] for ch in case["quotient_chunks"]]

    is_first, is_last, is_transition, inv_zeroifier = selectors_at_point(degree_bits, zeta)
    quotient = reconstruct_quotient(degree_bits, chunks, zeta)
    folder = AIR_FOLDERS[case["air"]]
    folded = folder(tl, tn, pis, is_first, is_last, is_transition, alpha)
    lhs = F.ext_mul(folded, inv_zeroifier)
    return F.ext_eq(lhs, quotient)


def check_case(case, tamper=None):
    """Run the identity check, optionally injecting a fault to demonstrate rejection: 'trace'
    (corrupt a trace opening), 'quotient' (corrupt a quotient chunk opening), 'alpha' (wrong alpha)."""
    c = json.loads(json.dumps(case))  # deep copy
    if tamper == "trace":
        c["trace_local"][0][0] = (c["trace_local"][0][0] + 1) % P
    elif tamper == "quotient":
        c["quotient_chunks"][0][0][0] = (c["quotient_chunks"][0][0][0] + 1) % P
    elif tamper == "alpha":
        c["alpha"][0] = (c["alpha"][0] + 1) % P
    return check_identity(c)


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

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


def conformance_kats():
    """(passed, total) after: (0) perm/generator sanity tying the ground truth to the CONFIRMED
    Poseidon2 permutation + component (a); (1) accepting every genuine proof's identity; (2) rejecting
    each of three tamper variants per case."""
    gt = load_ground_truth()
    passed = 0
    total = 0

    # (0) sanity
    import mmcs_babybear_reference as M
    total += 1
    if M.permute([0] * 16) == gt["perm_zeros"]:
        passed += 1
    total += 1
    if gt["val_generator"] == F.GENERATOR:
        passed += 1

    # (1) accept genuine identity; (2) reject tampered.
    for case in gt["cases"]:
        total += 1
        if case["library_accept"] and check_case(case):
            passed += 1
        for tamper in ("trace", "quotient", "alpha"):
            total += 1
            if not check_case(case, tamper):
                passed += 1
    return passed, total


# ============================ shared cross-language vector set ============================
# All three languages emit, per case: the accept flag, the three tamper-reject flags, and the recomputed
# quotient(zeta) and folded_constraints(zeta) (F_{p^4}). The harness asserts byte-identical values across
# Python/Node/Java (three independent implementations) AND that every accept/reject flag holds.

def _case_derived(case):
    degree_bits = case["degree_bits"]
    zeta = list(case["zeta"])
    chunks = [[list(v) for v in ch] for ch in case["quotient_chunks"]]
    is_first, is_last, is_transition, inv_zeroifier = selectors_at_point(degree_bits, zeta)
    quotient = reconstruct_quotient(degree_bits, chunks, zeta)
    folder = AIR_FOLDERS[case["air"]]
    folded = folder([list(v) for v in case["trace_local"]],
                    [list(v) for v in case["trace_next"]],
                    list(case["public_values"]),
                    is_first, is_last, is_transition, list(case["alpha"]))
    return {"quotient": quotient, "folded": folded, "invZeroifier": inv_zeroifier}


def shared_vectors():
    gt = load_ground_truth()
    cases = []
    for case in gt["cases"]:
        d = _case_derived(case)
        cases.append({
            "name": case["name"],
            "air": case["air"],
            "degreeBits": case["degree_bits"],
            "quotientDegree": case["quotient_degree"],
            "quotient": d["quotient"],
            "folded": d["folded"],
            "invZeroifier": d["invZeroifier"],
            "accept": check_case(case),
            "rejectTrace": not check_case(case, "trace"),
            "rejectQuotient": not check_case(case, "quotient"),
            "rejectAlpha": not check_case(case, "alpha"),
        })
    return {"valGenerator": F.GENERATOR, "cases": cases}


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