sumcheck.rs

//! The claim-batching sumcheck of an accumulation step (and of the
//! decider): for words `f_1..f_m` (multilinear in `ℓ` variables, tables in
//! lens's order — bit `k` of an index is variable `k`) and weights
//! `w_i = Σ_j c_ij · eq(z_ij, ·)`, prove
//!
//! ```text
//! Σ_{b ∈ {0,1}^ℓ} Σ_i w_i(b) · f_i(b) = σ
//! ```
//!
//! Variable `k` is bound in round `k` (low bit first), so the final point
//! `ρ` is in lens order. Round `k` sends `h(0), h(2)` of the degree-2 round
//! polynomial `h`; `h(1) = claim − h(0)`. The verifier ends with the claim
//! `Σ_i w_i(ρ) · f_i(ρ)` and the prover sends every `f_i(ρ)`.

use crate::fs::FiatShamir;
use lens::rspcs::field::{eq_eval, eq_table, quadratic_at};
use nebu::Fp3;

/// Weighted sum of evaluation claims on one word: `Σ_j coef_j · eq(point_j, ·)`.
#[derive(Clone, Debug, Default)]
pub struct Weight {
    pub terms: Vec<(Vec<Fp3>, Fp3)>,
}

impl Weight {
    /// `w(x) = Σ_j coef_j · eq(point_j, x)`.
    pub fn eval(&self, x: &[Fp3]) -> Fp3 {
        self.terms
            .iter()
            .fold(Fp3::ZERO, |acc, (p, c)| acc + *c * eq_eval(p, x))
    }
    /// The table of `w` over `{0,1}^ℓ`.
    pub fn table(&self, vars: usize) -> Vec<Fp3> {
        let mut out = vec![Fp3::ZERO; 1 << vars];
        add_eq_terms(&mut out, &self.terms);
        out
    }
}

/// `out += Σ_j c_j·eq(p_j, ·)` over `{0,1}^ℓ` (`out.len() = 2^ℓ`), in
/// parallel: an index splits as `hi·2^k + lo` (`p[..k]` the low
/// variables), so `eq(p, x) = eq(p_lo, lo)·eq(p_hi, hi)` and every block
/// of `2^k` entries adds `Σ_j (c_j·eq(p_hi, hi))·eq_table(p_lo)`.
pub fn add_eq_terms(out: &mut [Fp3], terms: &[(Vec<Fp3>, Fp3)]) {
    let vars = out.len().trailing_zeros() as usize;
    if terms.is_empty() {
        return;
    }
    let k = vars.min(12);
    // a point whose low coordinates are base (a spot check's ω^s powers)
    // has a base low table: its blocks add in mixed arithmetic
    let is_base = |p: &[Fp3]| {
        p.iter()
            .all(|x| x.c1 == nebu::Goldilocks::ZERO && x.c2 == nebu::Goldilocks::ZERO)
    };
    let lo: Vec<Vec<Fp3>> = terms.iter().map(|(p, _)| eq_table(&p[..k])).collect();
    let lo_base: Vec<Option<Vec<nebu::Goldilocks>>> = terms
        .iter()
        .zip(&lo)
        .map(|((p, _), t)| is_base(&p[..k]).then(|| t.iter().map(|x| x.c0).collect()))
        .collect();
    let hi: Vec<Vec<Fp3>> = terms
        .iter()
        .map(|(p, c)| {
            let mut t = eq_table(&p[k..]);
            for x in &mut t {
                *x *= *c;
            }
            t
        })
        .collect();
    let block = 1usize << k;
    let threads = if vars < 16 {
        1
    } else {
        std::thread::available_parallelism()
            .map_or(1, |n| n.get())
            .min(16)
    };
    let blocks = out.len() / block;
    let per = blocks.div_ceil(threads).max(1);
    std::thread::scope(|s| {
        for (ci, chunk) in out.chunks_mut(per * block).enumerate() {
            let (lo, hi, lo_base) = (&lo, &hi, &lo_base);
            s.spawn(move || {
                for (bi, b) in chunk.chunks_mut(block).enumerate() {
                    let h = ci * per + bi;
                    for ((l, lb), hv) in lo.iter().zip(lo_base).zip(hi) {
                        let sc = hv[h];
                        match lb {
                            Some(lb) => {
                                for (o, &e) in b.iter_mut().zip(lb) {
                                    *o += lens::rspcs::field::mul_base(sc, e);
                                }
                            }
                            None => {
                                for (o, &e) in b.iter_mut().zip(l) {
                                    *o += sc * e;
                                }
                            }
                        }
                    }
                }
            });
        }
    });
}

#[cfg(test)]
mod eq_terms_tests {
    use super::*;
    use nebu::Goldilocks;

    #[test]
    fn split_eq_terms_equal_the_full_tables() {
        for vars in [3usize, 12, 13, 17] {
            let pt = |s: u64| {
                (0..vars as u64)
                    .map(|i| {
                        Fp3::new(
                            Goldilocks::new(s * 31 + i),
                            Goldilocks::new(i * i + s),
                            Goldilocks::new(7),
                        )
                    })
                    .collect::<Vec<_>>()
            };
            let base_pt: Vec<Fp3> = (0..vars as u64)
                .map(|i| Fp3::from_base(Goldilocks::new(i * 7 + 2)))
                .collect();
            let terms = vec![
                (pt(1), Fp3::from_base(Goldilocks::new(5))),
                (
                    pt(2),
                    Fp3::new(Goldilocks::new(2), Goldilocks::new(9), Goldilocks::ZERO),
                ),
                (
                    base_pt,
                    Fp3::new(Goldilocks::new(4), Goldilocks::new(1), Goldilocks::new(8)),
                ),
            ];
            let mut want = vec![Fp3::ZERO; 1 << vars];
            for (p, c) in &terms {
                for (o, e) in want.iter_mut().zip(eq_table(p)) {
                    *o += *c * e;
                }
            }
            let mut got = vec![Fp3::ZERO; 1 << vars];
            add_eq_terms(&mut got, &terms);
            assert_eq!(got, want, "{vars}");
        }
    }
}

/// A word's table as the prover holds it: base-field entries (a fresh
/// trace word) or Fp3.
#[derive(Clone, Copy)]
pub enum Table<'a> {
    Base(&'a [nebu::Goldilocks]),
    Ext(&'a [Fp3]),
}

impl Table<'_> {
    fn len(&self) -> usize {
        match self {
            Table::Base(t) => t.len(),
            Table::Ext(t) => t.len(),
        }
    }
}

/// [`prove`] over borrowed tables, base-field ones summed and folded in
/// mixed arithmetic in the first round (the same messages: `f(2) = 2f(1)
/// − f(0)` and the fold are exact in either field). The weight tables
/// are built, and every first round summed and folded, in one pass of
/// threads per word.
pub fn prove_tables(
    t: &mut impl FiatShamir,
    f: &[Table<'_>],
    weights: &[Weight],
    vars: usize,
) -> (Vec<Fp3>, Vec<Fp3>, Vec<Fp3>) {
    use lens::rspcs::field::mul_base;
    assert!(vars >= 1 && f.iter().all(|x| x.len() == 1 << vars));
    let mut w: Vec<Vec<Fp3>> = weights.iter().map(|w| w.table(vars)).collect();
    let (mut h0, mut h2) = (Fp3::ZERO, Fp3::ZERO);
    for (fi, wi) in f.iter().zip(&w) {
        let (a, b) = match fi {
            Table::Ext(e) => round_sums(e, wi),
            Table::Base(bt) => par_sums(bt.len() / 2, |i| {
                let (f0, f1) = (bt[2 * i], bt[2 * i + 1]);
                let (w0, w1) = (wi[2 * i], wi[2 * i + 1]);
                (mul_base(w0, f0), mul_base(w1 + w1 - w0, f1 + f1 - f0))
            }),
        };
        h0 += a;
        h2 += b;
    }
    t.absorb_fp3(h0);
    t.absorb_fp3(h2);
    let mut msgs = vec![h0, h2];
    let alpha = t.squeeze_fp3();
    let mut point = vec![alpha];
    let mut fo: Vec<Vec<Fp3>> = f
        .iter()
        .map(|fi| match fi {
            Table::Ext(e) => par_fold(e.len() / 2, |i| {
                e[2 * i] + alpha * (e[2 * i + 1] - e[2 * i])
            }),
            Table::Base(bt) => par_fold(bt.len() / 2, |i| {
                mul_base(alpha, bt[2 * i + 1] - bt[2 * i]) + Fp3::from_base(bt[2 * i])
            }),
        })
        .collect();
    for wi in &mut w {
        fold(wi, alpha);
    }
    for _ in 1..vars {
        let (mut h0, mut h2) = (Fp3::ZERO, Fp3::ZERO);
        for (fi, wi) in fo.iter().zip(&w) {
            let (a, b) = round_sums(fi, wi);
            h0 += a;
            h2 += b;
        }
        t.absorb_fp3(h0);
        t.absorb_fp3(h2);
        msgs.push(h0);
        msgs.push(h2);
        let alpha = t.squeeze_fp3();
        for fi in &mut fo {
            fold(fi, alpha);
        }
        for wi in &mut w {
            fold(wi, alpha);
        }
        point.push(alpha);
    }
    let evals = fo.iter().map(|fi| fi[0]).collect();
    (msgs, point, evals)
}

/// `Σ_i term(i)` over `pairs` indices, componentwise, in threads.
fn par_sums(pairs: usize, term: impl Fn(usize) -> (Fp3, Fp3) + Sync) -> (Fp3, Fp3) {
    let t = threads(pairs);
    let size = pairs.div_ceil(t).max(1);
    std::thread::scope(|s| {
        let term = &term;
        let hs: Vec<_> = (0..pairs)
            .step_by(size)
            .map(|lo| {
                s.spawn(move || {
                    (lo..(lo + size).min(pairs)).fold((Fp3::ZERO, Fp3::ZERO), |a, i| {
                        let (x, y) = term(i);
                        (a.0 + x, a.1 + y)
                    })
                })
            })
            .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))
    })
}

/// `(0..half).map(at)` in threads.
fn par_fold(half: usize, at: impl Fn(usize) -> Fp3 + Sync) -> Vec<Fp3> {
    let mut out = vec![Fp3::ZERO; half];
    let size = half.div_ceil(threads(half)).max(1);
    std::thread::scope(|s| {
        for (k, chunk) in out.chunks_mut(size).enumerate() {
            let at = &at;
            s.spawn(move || {
                for (i, o) in chunk.iter_mut().enumerate() {
                    *o = at(k * size + i);
                }
            });
        }
    });
    out
}

/// Run the prover over tables `f_i` and their weights; returns the round
/// messages `(h(0), h(2))…`, the point `ρ` and every `f_i(ρ)`.
pub fn prove(
    t: &mut impl FiatShamir,
    mut f: Vec<Vec<Fp3>>,
    weights: &[Weight],
    vars: usize,
) -> (Vec<Fp3>, Vec<Fp3>, Vec<Fp3>) {
    let mut w: Vec<Vec<Fp3>> = weights.iter().map(|w| w.table(vars)).collect();
    let mut msgs = Vec::with_capacity(2 * vars);
    let mut point = Vec::with_capacity(vars);
    for _ in 0..vars {
        let (mut h0, mut h2) = (Fp3::ZERO, Fp3::ZERO);
        for (fi, wi) in f.iter().zip(&w) {
            let (a, b) = round_sums(fi, wi);
            h0 += a;
            h2 += b;
        }
        t.absorb_fp3(h0);
        t.absorb_fp3(h2);
        msgs.push(h0);
        msgs.push(h2);
        let alpha = t.squeeze_fp3();
        for fi in &mut f {
            fold(fi, alpha);
        }
        for wi in &mut w {
            fold(wi, alpha);
        }
        point.push(alpha);
    }
    let evals = f.iter().map(|fi| fi[0]).collect();
    (msgs, point, evals)
}

/// Threads for a pass over `pairs` pairs (one below `2^14`).
fn threads(pairs: usize) -> usize {
    if pairs < 1 << 14 {
        1
    } else {
        std::thread::available_parallelism()
            .map_or(1, |n| n.get())
            .min(16)
    }
}

/// `(Σ_b' f(0,b')w(0,b'), Σ_b' f(2,b')w(2,b'))` with variable 0 the low bit.
fn round_sums(f: &[Fp3], w: &[Fp3]) -> (Fp3, Fp3) {
    let part = |f: &[Fp3], w: &[Fp3]| {
        let (mut h0, mut h2) = (Fp3::ZERO, Fp3::ZERO);
        for (fp, wp) in f.as_chunks::<2>().0.iter().zip(w.as_chunks::<2>().0.iter()) {
            h0 += fp[0] * wp[0];
            let f2 = fp[1] + fp[1] - fp[0];
            let w2 = wp[1] + wp[1] - wp[0];
            h2 += f2 * w2;
        }
        (h0, h2)
    };
    let t = threads(f.len() / 2);
    if t == 1 {
        return part(f, w);
    }
    let size = (f.len() / 2).div_ceil(t) * 2;
    std::thread::scope(|s| {
        let hs: Vec<_> = f
            .chunks(size)
            .zip(w.chunks(size))
            .map(|(a, b)| s.spawn(move || part(a, b)))
            .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))
    })
}

fn fold(v: &mut Vec<Fp3>, alpha: Fp3) {
    let half = v.len() / 2;
    let t = threads(half);
    if t == 1 {
        for i in 0..half {
            let a = v[2 * i];
            v[i] = a + alpha * (v[2 * i + 1] - a);
        }
    } else {
        let mut out = vec![Fp3::ZERO; half];
        let size = half.div_ceil(t);
        let src = &v[..];
        std::thread::scope(|s| {
            for (k, chunk) in out.chunks_mut(size).enumerate() {
                s.spawn(move || {
                    for (i, o) in chunk.iter_mut().enumerate() {
                        let j = k * size + i;
                        let a = src[2 * j];
                        *o = a + alpha * (src[2 * j + 1] - a);
                    }
                });
            }
        });
        *v = out;
        return;
    }
    v.truncate(half);
}

/// Verifier: replay the rounds on claim `sigma`; returns `ρ` and the final
/// claim, or `None` on a malformed message list.
pub fn verify(
    t: &mut impl FiatShamir,
    sigma: Fp3,
    msgs: &[Fp3],
    vars: usize,
) -> Option<(Vec<Fp3>, Fp3)> {
    if msgs.len() != 2 * vars {
        return None;
    }
    let mut claim = sigma;
    let mut point = Vec::with_capacity(vars);
    for pair in msgs.as_chunks::<2>().0 {
        let (h0, h2) = (pair[0], pair[1]);
        let h1 = claim - h0;
        t.absorb_fp3(h0);
        t.absorb_fp3(h2);
        let alpha = t.squeeze_fp3();
        claim = quadratic_at(h0, h1, h2, alpha);
        point.push(alpha);
    }
    Some((point, claim))
}

#[cfg(test)]
mod tests {
    use super::*;
    use lens::Transcript;
    use lens::rspcs::field::ml_eval_ext;
    use nebu::Goldilocks;

    fn e(i: u64) -> Fp3 {
        Fp3::new(
            Goldilocks::new(i),
            Goldilocks::new(3 * i + 1),
            Goldilocks::new(i * i + 7),
        )
    }

    #[test]
    fn batched_sumcheck_reduces_to_the_point_evaluations() {
        let vars = 5;
        let f: Vec<Vec<Fp3>> = (0..3u64)
            .map(|i| (0..32u64).map(|b| e(100 * i + b)).collect())
            .collect();
        let weights: Vec<Weight> = (0..3u64)
            .map(|i| Weight {
                terms: (0..2u64)
                    .map(|j| {
                        (
                            (0..vars as u64).map(|k| e(7 * i + j + k)).collect(),
                            e(i + j + 2),
                        )
                    })
                    .collect(),
            })
            .collect();
        let sigma = f
            .iter()
            .zip(&weights)
            .map(|(fi, wi)| {
                wi.terms
                    .iter()
                    .fold(Fp3::ZERO, |a, (p, c)| a + *c * ml_eval_ext(fi, p))
            })
            .fold(Fp3::ZERO, |a, b| a + b);
        let mut tp = Transcript::new(b"t");
        let (msgs, point, evals) = prove(&mut tp, f.clone(), &weights, vars);
        let mut tv = Transcript::new(b"t");
        let (vp, claim) = verify(&mut tv, sigma, &msgs, vars).unwrap();
        assert_eq!(vp, point);
        for (fi, &v) in f.iter().zip(&evals) {
            assert_eq!(ml_eval_ext(fi, &point), v);
        }
        let expect = weights
            .iter()
            .zip(&evals)
            .fold(Fp3::ZERO, |a, (w, &v)| a + w.eval(&point) * v);
        assert_eq!(claim, expect);
        let mut tb = Transcript::new(b"t");
        let (_, bad) = verify(&mut tb, sigma + Fp3::ONE, &msgs, vars).unwrap();
        assert_ne!(bad, expect);
        // borrowed and base-field tables: the same messages, point, values
        let base: Vec<Goldilocks> = (0..32u64).map(|b| Goldilocks::new(b * b + 3)).collect();
        let lifted: Vec<Fp3> = base.iter().map(|&x| Fp3::from_base(x)).collect();
        let mut f2 = f.clone();
        f2[1] = lifted;
        let mut ta = Transcript::new(b"u");
        let want = prove(&mut ta, f2.clone(), &weights, vars);
        let tables = [Table::Ext(&f2[0]), Table::Base(&base), Table::Ext(&f2[2])];
        let mut tc = Transcript::new(b"u");
        assert_eq!(prove_tables(&mut tc, &tables, &weights, vars), want);
    }
}
soft3/zheng/rs/src/accumulate/sumcheck.rs

Graph
