use alloc::vec;
use alloc::vec::Vec;
use p3_field::{ExtensionField, Field};
use super::grid::evals_01inf_grid_into;
use crate::layout::{EqPartials, NextPartials, ProverMultiClaim};
use crate::strategy::VariableOrder;
pub(crate) type SvoAccumulators<EF> = Vec<[Vec<EF>; 2]>;
pub(crate) fn calculate_accumulators_batch<F: Field, EF: ExtensionField<F>>(
claim: &ProverMultiClaim<F, EF>,
alphas: &[EF],
) -> SvoAccumulators<EF> {
assert_eq!(claim.len(), alphas.len());
let k = claim.point().num_variables_svo();
(0..k)
.map(|round_idx| {
let l = round_idx + 1;
let mut eq_acc = EqPartials::zero(l);
let mut next_acc = NextPartials::zero(l);
for (opening, &alpha) in claim.current_openings().iter().zip(alphas.iter()) {
eq_acc.accumulate(&opening.data().rounds()[round_idx], alpha);
}
for (opening, &alpha) in claim
.next_openings()
.iter()
.zip(alphas.iter().skip(claim.current_openings().len()))
{
next_acc.accumulate(&opening.data().rounds()[round_idx], alpha);
}
let stride = 3usize.pow(round_idx as u32);
let mut acc0 = EF::zero_vec(stride);
let mut acc_inf = EF::zero_vec(stride);
match claim.point().var_order() {
VariableOrder::Prefix => {
let (svo_active, _) = claim.point().z_svo().split_at(l);
if !claim.current_openings().is_empty() {
eq_acc.accumulate_prefix(svo_active.as_slice(), &mut acc0, &mut acc_inf);
}
if !claim.next_openings().is_empty() {
next_acc.accumulate_prefix(svo_active.as_slice(), &mut acc0, &mut acc_inf);
}
}
VariableOrder::Suffix => {
let (_, svo_active) = claim.point().z_svo().split_at(k - l);
if !claim.current_openings().is_empty() {
eq_acc.accumulate_suffix(svo_active.as_slice(), &mut acc0, &mut acc_inf);
}
if !claim.next_openings().is_empty() {
next_acc.accumulate_suffix(svo_active.as_slice(), &mut acc0, &mut acc_inf);
}
}
}
[acc0, acc_inf]
})
.collect()
}
pub(crate) fn calculate_product_accumulator<F: Field>(
l: usize,
left: &[F],
right: &[F],
) -> [Vec<F>; 2] {
assert_eq!(left.len(), 1 << l);
assert_eq!(right.len(), 1 << l);
match l {
1 => calculate_accumulator_1(left.try_into().unwrap(), right.try_into().unwrap()),
2 => calculate_accumulator_2(left.try_into().unwrap(), right.try_into().unwrap()),
3 => calculate_accumulator_3(left.try_into().unwrap(), right.try_into().unwrap()),
_ => calculate_accumulator_general(l, left, right),
}
}
fn calculate_accumulator_1<EF: Field>(eq0: &[EF; 2], reduced: &[EF; 2]) -> [Vec<EF>; 2] {
let &[e0, e1] = eq0;
let &[r0, r1] = reduced;
let e_inf = e1 - e0;
let r_inf = r1 - r0;
[vec![e0 * r0], vec![e_inf * r_inf]]
}
fn calculate_accumulator_2<EF: Field>(eq0: &[EF; 4], reduced: &[EF; 4]) -> [Vec<EF>; 2] {
let &[e00, e10, e01, e11] = eq0;
let &[r00, r10, r01, r11] = reduced;
let e20 = e10 - e00;
let e21 = e11 - e01;
let r20 = r10 - r00;
let r21 = r11 - r01;
let e02 = e01 - e00;
let r02 = r01 - r00;
[
vec![e00 * r00, e01 * r01, e02 * r02],
vec![e20 * r20, e21 * r21, (e21 - e20) * (r21 - r20)],
]
}
fn calculate_accumulator_3<EF: Field>(eq0: &[EF; 8], reduced: &[EF; 8]) -> [Vec<EF>; 2] {
let &[e_000, e_100, e_010, e_110, e_001, e_101, e_011, e_111] = eq0;
let &[r_000, r_100, r_010, r_110, r_001, r_101, r_011, r_111] = reduced;
let e_200 = e_100 - e_000;
let e_210 = e_110 - e_010;
let e_201 = e_101 - e_001;
let e_211 = e_111 - e_011;
let r_200 = r_100 - r_000;
let r_210 = r_110 - r_010;
let r_201 = r_101 - r_001;
let r_211 = r_111 - r_011;
let e_020 = e_010 - e_000;
let e_220 = e_210 - e_200;
let e_021 = e_011 - e_001;
let e_221 = e_211 - e_201;
let r_020 = r_010 - r_000;
let r_220 = r_210 - r_200;
let r_021 = r_011 - r_001;
let r_221 = r_211 - r_201;
let e_002 = e_001 - e_000;
let e_012 = e_011 - e_010;
let e_022 = e_021 - e_020;
let e_202 = e_201 - e_200;
let e_212 = e_211 - e_210;
let e_222 = e_221 - e_220;
let r_002 = r_001 - r_000;
let r_012 = r_011 - r_010;
let r_022 = r_021 - r_020;
let r_202 = r_201 - r_200;
let r_212 = r_211 - r_210;
let r_222 = r_221 - r_220;
[
vec![
e_000 * r_000,
e_001 * r_001,
e_002 * r_002,
e_010 * r_010,
e_011 * r_011,
e_012 * r_012,
e_020 * r_020,
e_021 * r_021,
e_022 * r_022,
],
vec![
e_200 * r_200,
e_201 * r_201,
e_202 * r_202,
e_210 * r_210,
e_211 * r_211,
e_212 * r_212,
e_220 * r_220,
e_221 * r_221,
e_222 * r_222,
],
]
}
fn calculate_accumulator_general<F: Field, EF: ExtensionField<F>>(
l: usize,
eq0: &[EF],
reduced_evals: &[EF],
) -> [Vec<EF>; 2] {
let grid_len = 3usize.pow(l as u32);
let mut eq0_grid = EF::zero_vec(grid_len);
let mut reduced_grid = EF::zero_vec(grid_len);
let mut scratch = EF::zero_vec(grid_len);
evals_01inf_grid_into(eq0, &mut eq0_grid, &mut scratch);
evals_01inf_grid_into(reduced_evals, &mut reduced_grid, &mut scratch);
let stride = 3usize.pow((l - 1) as u32);
debug_assert_eq!(eq0_grid.len(), 3 * stride);
debug_assert_eq!(reduced_grid.len(), 3 * stride);
let acc0 = eq0_grid[..stride]
.iter()
.copied()
.zip(reduced_grid[..stride].iter().copied())
.map(|(eq, eval)| eq * eval)
.collect();
let acc_inf = eq0_grid[2 * stride..]
.iter()
.copied()
.zip(reduced_grid[2 * stride..].iter().copied())
.map(|(eq, eval)| eq * eval)
.collect();
[acc0, acc_inf]
}
#[cfg(test)]
mod test {
use alloc::vec::Vec;
use p3_field::extension::BinomialExtensionField;
use p3_field::{ExtensionField, Field, PrimeCharacteristicRing, dot_product};
use p3_koala_bear::KoalaBear;
use p3_multilinear_util::point::Point;
use p3_multilinear_util::poly::Poly;
use p3_multilinear_util::split_eq::SplitEq;
use p3_util::log2_strict_usize;
use proptest::prelude::*;
use rand::rngs::SmallRng;
use rand::{RngExt, SeedableRng};
use super::*;
use crate::lagrange::lagrange_weights_01inf_multi;
use crate::layout::Opening;
use crate::strategy::VariableOrder;
use crate::svo::SvoPoint;
type F = KoalaBear;
type EF = BinomialExtensionField<F, 4>;
#[test]
fn test_batch_svo_accumulators() {
let k = 12;
let n_polys = 3;
let mut rng = SmallRng::seed_from_u64(0);
let polys = (0..n_polys)
.map(|_| Poly::<F>::rand(&mut rng, k))
.collect::<Vec<_>>();
let alphas = (0..polys.len())
.map(|_| rng.random::<EF>())
.collect::<Vec<_>>();
let point = Point::<EF>::rand(&mut rng, k);
for l0 in 0..=k / 2 {
let svo_point = SvoPoint::<F, EF>::new_unpacked(l0, &point, VariableOrder::Suffix);
let openings = polys
.iter()
.map(|poly| {
let (eval, partial_evals) = svo_point.eval(poly.as_view());
let opening = Opening {
poly_idx: None,
eval,
data: partial_evals,
};
assert_eq!(opening.eval(), poly.eval_base(&point));
opening
})
.collect::<Vec<_>>();
let claim = ProverMultiClaim::new(svo_point, openings, Vec::new());
let accumulators = calculate_accumulators_batch(&claim, &alphas);
if l0 == 0 {
assert!(accumulators.is_empty());
continue;
}
let mut poly = Poly::<EF>::zero(l0);
claim
.current_openings()
.iter()
.zip(alphas.iter())
.for_each(|(opening, &alpha)| {
let full_svo_poly = opening
.data()
.rounds()
.last()
.expect("l0 > 0 guarantees one SVO partial polynomial");
poly.as_mut_slice()
.iter_mut()
.zip(full_svo_poly.poly().iter())
.for_each(|(out, &value)| *out += alpha * value);
});
let mut eq = Poly::new_from_point(claim.point().z_svo().as_slice(), EF::ONE);
let mut rs = Vec::with_capacity(l0);
for [acc0, acc_inf] in accumulators.iter() {
let weights = lagrange_weights_01inf_multi(rs.as_slice());
let c0 = dot_product::<EF, _, _>(acc0.iter().copied(), weights.iter().copied());
let cinf =
dot_product::<EF, _, _>(acc_inf.iter().copied(), weights.iter().copied());
let msg =
VariableOrder::Suffix.sumcheck_coefficients(poly.as_slice(), eq.as_slice());
assert_eq!(c0, msg.c_a);
assert_eq!(cinf, msg.c_inf);
let r: EF = rng.random();
poly.fix_suffix_var_mut(r);
eq.fix_suffix_var_mut(r);
rs.push(r);
}
}
for l0 in 0..=k / 2 {
let svo_point = SvoPoint::<F, EF>::new_unpacked(l0, &point, VariableOrder::Prefix);
let openings = polys
.iter()
.map(|poly| {
let (eval, partial_evals) = svo_point.eval(poly.as_view());
let opening = Opening {
poly_idx: None,
eval,
data: partial_evals,
};
assert_eq!(opening.eval(), poly.eval_base(&point));
opening
})
.collect::<Vec<_>>();
let claim = ProverMultiClaim::new(svo_point, openings, Vec::new());
let accumulators = calculate_accumulators_batch(&claim, &alphas);
if l0 == 0 {
assert!(accumulators.is_empty());
continue;
}
let mut poly = Poly::<EF>::zero(l0);
claim
.current_openings()
.iter()
.zip(alphas.iter())
.for_each(|(opening, &alpha)| {
let full_svo_poly = opening
.data()
.rounds()
.last()
.expect("l0 > 0 guarantees one SVO partial polynomial");
poly.as_mut_slice()
.iter_mut()
.zip(full_svo_poly.poly().iter())
.for_each(|(out, &value)| *out += alpha * value);
});
let mut eq = Poly::new_from_point(claim.point().z_svo().as_slice(), EF::ONE);
let mut rs = Vec::with_capacity(l0);
for [acc0, acc_inf] in accumulators.iter() {
let weights = lagrange_weights_01inf_multi(rs.as_slice());
let c0 = dot_product::<EF, _, _>(acc0.iter().copied(), weights.iter().copied());
let cinf =
dot_product::<EF, _, _>(acc_inf.iter().copied(), weights.iter().copied());
let msg =
VariableOrder::Prefix.sumcheck_coefficients(poly.as_slice(), eq.as_slice());
assert_eq!(c0, msg.c_a);
assert_eq!(cinf, msg.c_inf);
let r: EF = rng.random();
poly.fix_prefix_var_mut(r);
eq.fix_prefix_var_mut(r);
rs.push(r);
}
}
}
fn calculate_accumulator<F: Field, EF: ExtensionField<F>>(
l: usize,
partial_evals: &[EF],
point: &[EF],
) -> [Vec<EF>; 2] {
let total_vars = log2_strict_usize(partial_evals.len());
let offset = total_vars - l;
let (z0, z1) = point.split_at(point.len() - offset);
let eq0 = Poly::new_from_point(z0, EF::ONE);
let eq1 = Poly::new_from_point(z1, EF::ONE);
let reduced_evals = partial_evals
.chunks(eq1.num_evals())
.map(|chunk| dot_product::<EF, _, _>(eq1.iter().copied(), chunk.iter().copied()))
.collect::<Vec<_>>();
calculate_product_accumulator(l, eq0.as_slice(), &reduced_evals)
}
proptest! {
#[test]
fn prop_accumulators_specialization_matches_general(k in 10usize..=14) {
let mut rng = SmallRng::seed_from_u64(k as u64);
let poly = Poly::new((0..1 << k).map(|_| rng.random()).collect::<Vec<_>>());
let point = Point::<EF>::rand(&mut rng, k);
let (z_svo, z_rest) = point.split_at(k / 2);
let split_eq = SplitEq::<F, EF>::new_packed(&z_rest, EF::ONE);
let partial_evals = split_eq.compress_suffix(poly.as_view());
for l in 1..k / 2 {
let dispatched =
calculate_accumulator::<F, EF>(l, partial_evals.as_slice(), z_svo.as_slice());
let eq0 = Poly::new_from_point(&z_svo.as_slice()[..l], EF::ONE);
let eq1 = Poly::new_from_point(&z_svo.as_slice()[l..], EF::ONE);
let reduced: Vec<EF> = partial_evals
.as_slice()
.chunks(eq1.num_evals())
.map(|chunk| dot_product::<EF, _, _>(eq1.iter().copied(), chunk.iter().copied()))
.collect();
let general = calculate_accumulator_general::<F, EF>(l, eq0.as_slice(), &reduced);
prop_assert_eq!(dispatched, general);
}
}
#[test]
fn prop_accumulators_1_matches_general(seed in 0u64..1000) {
let mut rng = SmallRng::seed_from_u64(seed);
let eq0: [EF; 2] = rng.random();
let reduced: [EF; 2] = rng.random();
let fast = calculate_accumulator_1(&eq0, &reduced);
let general = calculate_accumulator_general::<F, EF>(1, &eq0, &reduced);
prop_assert_eq!(fast, general);
}
#[test]
fn prop_accumulators_2_matches_general(seed in 0u64..1000) {
let mut rng = SmallRng::seed_from_u64(seed);
let eq0: [EF; 4] = rng.random();
let reduced: [EF; 4] = rng.random();
let fast = calculate_accumulator_2(&eq0, &reduced);
let general = calculate_accumulator_general::<F, EF>(2, &eq0, &reduced);
prop_assert_eq!(fast, general);
}
#[test]
fn prop_accumulators_3_matches_general(seed in 0u64..1000) {
let mut rng = SmallRng::seed_from_u64(seed);
let eq0: [EF; 8] = rng.random();
let reduced: [EF; 8] = rng.random();
let fast = calculate_accumulator_3(&eq0, &reduced);
let general = calculate_accumulator_general::<F, EF>(3, &eq0, &reduced);
prop_assert_eq!(fast, general);
}
}
}