use nebu::{Fp3, Goldilocks};
use super::coins::{Coins, Label};
use crate::types::SumcheckPoly;
#[derive(Clone)]
pub(crate) struct Mask {
coeffs: Vec<Vec<Fp3>>,
}
fn two_pow(e: usize) -> Fp3 {
Fp3::from_base(Goldilocks::new(2).exp(e as u64))
}
fn eval(c: &[Fp3], x: Fp3) -> Fp3 {
c.iter().rev().fold(Fp3::ZERO, |acc, &k| acc * x + k)
}
pub(crate) fn limb_basis(l: usize) -> Fp3 {
let t = Fp3::new(Goldilocks::ZERO, Goldilocks::ONE, Goldilocks::ZERO);
(0..l).fold(Fp3::ONE, |acc, _| acc * t)
}
impl Mask {
pub fn entries(rounds: usize, degree: usize) -> usize {
rounds * (degree + 1) * 3
}
pub fn sample(rounds: usize, degree: usize, coins: &Coins, label: Label) -> Self {
let all = coins.exts(label, rounds * (degree + 1));
Self {
coeffs: all.chunks(degree + 1).map(<[Fp3]>::to_vec).collect(),
}
}
pub fn limbs(&self) -> Vec<Goldilocks> {
self.coeffs
.iter()
.flatten()
.flat_map(|c| [c.c0, c.c1, c.c2])
.collect()
}
pub fn sum(&self) -> Fp3 {
let r = self.coeffs.len();
if r == 0 {
return Fp3::ZERO;
}
let ends = self.coeffs.iter().fold(Fp3::ZERO, |acc, c| {
acc + eval(c, Fp3::ZERO) + eval(c, Fp3::ONE)
});
two_pow(r - 1) * ends
}
pub fn round_poly(&self, prefix: &[Fp3]) -> Vec<Fp3> {
let r = self.coeffs.len();
let j = prefix.len();
let bound = prefix
.iter()
.zip(&self.coeffs)
.fold(Fp3::ZERO, |acc, (&x, c)| acc + eval(c, x));
let scale = two_pow(r - 1 - j);
let mut out: Vec<Fp3> = self.coeffs[j].iter().map(|&c| scale * c).collect();
out[0] += scale * bound;
if j + 1 < r {
let rest = self.coeffs[j + 1..].iter().fold(Fp3::ZERO, |acc, c| {
acc + eval(c, Fp3::ZERO) + eval(c, Fp3::ONE)
});
out[0] += two_pow(r - 2 - j) * rest;
}
out
}
pub fn eval_at(&self, point: &[Fp3]) -> Fp3 {
point
.iter()
.zip(&self.coeffs)
.fold(Fp3::ZERO, |acc, (&x, c)| acc + eval(c, x))
}
}
pub(crate) fn weights(point: &[Fp3], degree: usize, scale: Fp3) -> Vec<Fp3> {
let basis = [limb_basis(0), limb_basis(1), limb_basis(2)];
let mut out = Vec::with_capacity(Mask::entries(point.len(), degree));
for &x in point {
let mut power = scale;
for _ in 0..=degree {
for b in basis {
out.push(power * b);
}
power *= x;
}
}
out
}
pub(crate) fn masked(f: SumcheckPoly<Fp3>, rho: Fp3, g: &[Fp3]) -> SumcheckPoly<Fp3> {
debug_assert_eq!(f.coeffs.len(), g.len());
SumcheckPoly {
degree: f.degree,
coeffs: f.coeffs.iter().zip(g).map(|(&a, &b)| a + rho * b).collect(),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rounds_sums_and_weights_agree_with_brute_force() {
let coins = Coins::seeded([3; 32]);
let mut rng = coins.stream(crate::execution::veil::coins::label::TEST);
for (rounds, degree) in [(1, 2), (3, 2), (4, 8)] {
let g = Mask::sample(rounds, degree, &coins, Label(1 << 33 | rounds as u64));
let brute = (0..1usize << rounds).fold(Fp3::ZERO, |acc, x| {
let p: Vec<Fp3> = (0..rounds)
.map(|i| Fp3::from_base(Goldilocks::new(((x >> (rounds - 1 - i)) & 1) as u64)))
.collect();
acc + g.eval_at(&p)
});
assert_eq!(g.sum(), brute);
let mut claim = g.sum();
let mut prefix = vec![];
for j in 0..rounds {
let s = g.round_poly(&prefix);
assert_eq!(eval(&s, Fp3::ZERO) + eval(&s, Fp3::ONE), claim, "round {j}");
let r = rng.ext();
claim = eval(&s, r);
prefix.push(r);
}
assert_eq!(claim, g.eval_at(&prefix));
let scale = rng.ext();
let w = weights(&prefix, degree, scale);
let dot = w
.iter()
.zip(g.limbs())
.fold(Fp3::ZERO, |acc, (&w, v)| acc + w * Fp3::from_base(v));
assert_eq!(dot, scale * g.eval_at(&prefix));
}
}
}