soft3/zheng/rs/src/recursion/whir/prove.rs

//! The batched field-native WHIR prover: the messages
//! [`super::verify`] reads, in its transcript order.

use lens::rspcs::field::{
    add_scaled_eq, eq_table, fold_coeffs_ext, fold_evals_ext, pow_point, root_of_unity,
};
use lens::rspcs::whir::RoundSpec;
use nebu::Fp3;

use super::{BatchProof, Config, Proof, Round, Tree};
use crate::accumulate::sumcheck::{self, Weight};
use crate::recursion::sponge::ProverTranscript;
use crate::recursion::word::{LeafOpening, Word};

/// `(h(0), h(2))` of `Σ_b f(b)·w(b)` in its first (low) variable.
fn round_poly(f: &[Fp3], w: &[Fp3]) -> (Fp3, Fp3) {
    let pairs = f.len() / 2;
    let part = |lo: usize, hi: usize| {
        let (mut a, mut c) = (Fp3::ZERO, Fp3::ZERO);
        for i in lo..hi {
            let (f0, f1, w0, w1) = (f[2 * i], f[2 * i + 1], w[2 * i], w[2 * i + 1]);
            a += f0 * w0;
            c += (f1 + f1 - f0) * (w1 + w1 - w0);
        }
        (a, c)
    };
    if pairs < 1 << 14 {
        return part(0, pairs);
    }
    let threads = std::thread::available_parallelism()
        .map_or(1, |n| n.get())
        .min(16);
    let size = pairs.div_ceil(threads);
    std::thread::scope(|s| {
        let hs: Vec<_> = (0..threads)
            .map(|k| s.spawn(move || part(k * size, ((k + 1) * size).min(pairs))))
            .collect();
        hs.into_iter()
            .map(|h| h.join().expect("round"))
            .fold((Fp3::ZERO, Fp3::ZERO), |a, x| (a.0 + x.0, a.1 + x.1))
    })
}

/// `dst += c · src`, in threads for long tables (exact field arithmetic).
fn axpy(dst: &mut [Fp3], c: Fp3, src: &[Fp3]) {
    let threads = std::thread::available_parallelism()
        .map_or(1, |n| n.get())
        .min(16);
    if dst.len() < 1 << 14 || threads == 1 {
        for (x, &y) in dst.iter_mut().zip(src) {
            *x += c * y;
        }
        return;
    }
    let per = dst.len().div_ceil(threads);
    std::thread::scope(|s| {
        for (d, sr) in dst.chunks_mut(per).zip(src.chunks(per)) {
            s.spawn(move || {
                for (x, &y) in d.iter_mut().zip(sr) {
                    *x += c * y;
                }
            });
        }
    });
}

/// The folding sumcheck: `rounds` rounds over `f·w`, folding `f`, `w`
/// and the coefficients.
fn fold_rounds(
    t: &mut ProverTranscript,
    f: &mut Vec<Fp3>,
    w: &mut Vec<Fp3>,
    coeffs: &mut Vec<Fp3>,
    rounds: usize,
    pow: u32,
) -> (Vec<Fp3>, Vec<u64>) {
    let mut msgs = Vec::with_capacity(2 * rounds);
    let mut nonces = Vec::new();
    for _ in 0..rounds {
        let (h0, h2) = round_poly(f, w);
        t.absorb_ext(h0);
        t.absorb_ext(h2);
        msgs.push(h0);
        msgs.push(h2);
        if pow > 0 {
            nonces.push(t.grind(pow));
        }
        let a = t.squeeze_ext();
        fold_evals_ext(f, a);
        fold_evals_ext(w, a);
        fold_coeffs_ext(coeffs, a);
    }
    (msgs, nonces)
}

fn queries(t: &mut ProverTranscript, s: &RoundSpec) -> Result<Vec<usize>, String> {
    if s.opens_all() {
        return Ok((0..1usize << s.log_leaves()).collect());
    }
    t.indices(s.queries, s.log_leaves() as usize)
}

/// Prove the batched opening of the words of `trees` (each in round 0's
/// layout, `cfg.groups` words a tree) with `claims[i]` on word `i`
/// (`(point, value)`, a univariate claim as its `pow` point).
pub fn prove(
    cfg: &Config,
    t: &mut ProverTranscript,
    trees: &[&dyn Tree],
    claims: &[Vec<(Vec<Fp3>, Fp3)>],
) -> Result<Proof, String> {
    let wc = &cfg.wc;
    let ell = wc.num_vars;
    let words: Vec<&Word> = trees.iter().flat_map(|tr| tr.members()).collect();
    let m = words.len();
    let shape: Vec<usize> = trees.iter().map(|tr| tr.members().len()).collect();
    if cfg.direct()
        || m != cfg.inputs
        || shape != cfg.groups
        || trees.iter().any(|tr| tr.arity() != cfg.arity)
        || claims.len() != m
        || words
            .iter()
            .any(|w| w.num_vars != ell || w.layout != cfg.layout(0))
    {
        return Err("whir: input words".into());
    }
    // batch
    let gamma = t.squeeze_ext();
    let mut g = Fp3::ONE;
    let weights: Vec<Weight> = claims
        .iter()
        .map(|cs| Weight {
            terms: cs
                .iter()
                .map(|(pt, _)| {
                    let term = (pt.clone(), g);
                    g *= gamma;
                    term
                })
                .collect(),
        })
        .collect();
    let lap = crate::recursion::ivc::timer_pub("      whir prove ");
    let tables: Vec<Vec<Fp3>> = words.iter().map(|w| w.table()).collect();
    let (msgs, rho, evals) = sumcheck::prove(t, tables.clone(), &weights, ell);
    lap("batch sumcheck");
    for &e in &evals {
        t.absorb_ext(e);
    }
    let comb_nonce = t.grind(cfg.comb_pow);
    lap(&format!("combination grinding ({} bits)", cfg.comb_pow));
    let r = t.squeeze_ext();
    let mut coef = Vec::with_capacity(m);
    let mut c = Fp3::ONE;
    for _ in 0..m {
        coef.push(c);
        c *= r;
    }
    let mut f = vec![Fp3::ZERO; 1 << ell];
    let mut coeffs = vec![Fp3::ZERO; 1 << ell];
    for ((tb, w), &c) in tables.iter().zip(words).zip(&coef) {
        axpy(&mut f, c, tb);
        axpy(&mut coeffs, c, &w.coeffs());
    }
    drop(tables);
    let w = eq_table(&rho);
    lap("combination");
    core(
        cfg,
        t,
        trees,
        f,
        coeffs,
        w,
        BatchProof {
            sumcheck: msgs,
            evals,
            comb_nonce,
        },
    )
}

/// Prove an opening of one word with claims given by their weights'
/// tables (`Σ_x w(x)·f(x) = v`), in [`super::verify_direct`]'s order.
pub fn prove_direct(
    cfg: &Config,
    t: &mut ProverTranscript,
    word: &Word,
    claims: &[(Vec<Fp3>, Fp3)],
) -> Result<Proof, String> {
    let ell = cfg.wc.num_vars;
    if !cfg.direct()
        || word.arity != cfg.arity
        || word.num_vars != ell
        || word.layout != cfg.layout(0)
        || claims.len() > cfg.claims
    {
        return Err("whir: a direct opening".into());
    }
    let gamma = t.squeeze_ext();
    let mut w = vec![Fp3::ZERO; 1 << ell];
    let mut g = Fp3::ONE;
    for (tb, _) in claims {
        axpy(&mut w, g, tb);
        g *= gamma;
    }
    core(
        cfg,
        t,
        &[word as &dyn Tree],
        word.table(),
        word.coeffs(),
        w,
        BatchProof {
            sumcheck: vec![],
            evals: vec![],
            comb_nonce: 0,
        },
    )
}

/// WHIR on `f_0` (table `f`, coefficients) from its initial weight table.
fn core(
    cfg: &Config,
    t: &mut ProverTranscript,
    trees: &[&dyn Tree],
    mut f: Vec<Fp3>,
    mut coeffs: Vec<Fp3>,
    mut w: Vec<Fp3>,
    batch: BatchProof,
) -> Result<Proof, String> {
    let wc = &cfg.wc;
    let ell = wc.num_vars;
    let lap = crate::recursion::ivc::timer_pub("      whir core ");
    // round 0
    let s0 = wc.rounds[0];
    let mut ood0 = Vec::with_capacity(s0.ood);
    let mut zs = Vec::with_capacity(s0.ood);
    for _ in 0..s0.ood {
        let z = t.squeeze_ext();
        let y = lens::rspcs::field::univariate_ext(&coeffs, z);
        t.absorb_ext(y);
        ood0.push(y);
        zs.push(z);
    }
    let gamma = t.squeeze_ext();
    let mut g = gamma;
    for &z in &zs {
        add_scaled_eq(&mut w, &pow_point(z, ell), g);
        g *= gamma;
    }
    lap("ood 0");
    let (sumcheck0, fold_nonces0) =
        fold_rounds(t, &mut f, &mut w, &mut coeffs, s0.fold, s0.fold_pow);
    lap(&format!("folds 0 (grinding {} bits)", s0.fold_pow));
    let mut prev_words: Option<Word> = None;
    let mut prev = s0;
    let mut rounds = Vec::with_capacity(wc.rounds.len() - 1);
    // openings in parallel (a streamed word recomputes each leaf)
    let open = |prev_words: &Option<Word>, idx: &[usize]| -> Vec<Vec<LeafOpening>> {
        let one = |j: usize| match prev_words {
            Some(wd) => vec![wd.open(j)],
            None => trees.iter().map(|tr| tr.open(j)).collect::<Vec<_>>(),
        };
        let threads = std::thread::available_parallelism()
            .map_or(1, |n| n.get())
            .min(16)
            .min(idx.len().max(1));
        let size = idx.len().div_ceil(threads).max(1);
        std::thread::scope(|s| {
            let hs: Vec<_> = idx
                .chunks(size)
                .map(|c| s.spawn(move || c.iter().map(|&j| one(j)).collect::<Vec<_>>()))
                .collect();
            hs.into_iter()
                .flat_map(|h| h.join().expect("opening"))
                .collect()
        })
    };
    for (i, s) in wc.rounds.iter().enumerate().skip(1) {
        let wd = Word::commit_coeffs(cfg.layout(i), coeffs.clone(), cfg.arity);
        lap(&format!("round {i} commit"));
        t.absorb_all(&wd.root());
        let mut ood = Vec::with_capacity(s.ood);
        let mut zs = Vec::with_capacity(s.ood);
        for _ in 0..s.ood {
            let z = t.squeeze_ext();
            let y = wd.univariate(z);
            t.absorb_ext(y);
            ood.push(y);
            zs.push(z);
        }
        let query_nonce = t.grind(prev.query_pow);
        lap(&format!(
            "round {i} query grinding ({} bits)",
            prev.query_pow
        ));
        let idx = queries(t, &prev)?;
        let gamma = t.squeeze_ext();
        let opening = open(&prev_words, &idx);
        lap(&format!("round {i} openings"));
        let mut g = gamma;
        for &z in &zs {
            add_scaled_eq(&mut w, &pow_point(z, s.num_vars), g);
            g *= gamma;
        }
        let omega_l = root_of_unity(prev.log_leaves());
        for &j in &idx {
            let x = Fp3::from_base(omega_l.exp(j as u64));
            add_scaled_eq(&mut w, &pow_point(x, s.num_vars), g);
            g *= gamma;
        }
        let (sumcheck, fold_nonces) =
            fold_rounds(t, &mut f, &mut w, &mut coeffs, s.fold, s.fold_pow);
        lap(&format!("round {i} folds (grinding {} bits)", s.fold_pow));
        rounds.push(Round {
            root: wd.root(),
            ood,
            query_nonce,
            open: opening,
            sumcheck,
            fold_nonces,
        });
        prev_words = Some(wd);
        prev = *s;
    }
    if cfg.direct() {
        t.absorb_all(&crate::recursion::msg::digest_native(&coeffs));
    } else {
        for &c in &coeffs {
            t.absorb_ext(c);
        }
    }
    let final_nonce = t.grind(prev.query_pow);
    let idx = queries(t, &prev)?;
    let final_open = open(&prev_words, &idx);
    lap("final");
    if t.sp.has_pending() {
        t.sp.flush(&mut t.o);
    }
    Ok(Proof {
        batch,
        ood0,
        sumcheck0,
        fold_nonces0,
        rounds,
        final_poly: coeffs,
        final_nonce,
        final_open,
    })
}