Use partial product chain

src/plonk/prover.rs

@@ -268,8 +268,7 @@ fn wires_permutation_partial_products<F: RichField + Extendable<D>, const D: usi
             let quotient_partials = partial_products(&quotient_values, degree);
 
             // This is the final product for the quotient.
-            let quotient = quotient_partials
+            let quotient = quotient_partials[common_data.num_partial_products.1..]
-                [common_data.num_partial_products.0 - common_data.num_partial_products.1..]
                 .iter()
                 .copied()
                 .product();

src/util/partial_products.rs

@@ -1,5 +1,5 @@
 use std::iter::Product;
-use std::ops::Sub;
+use std::ops::{MulAssign, Sub};
 
 use crate::field::extension_field::target::ExtensionTarget;
 use crate::field::extension_field::Extendable;
@@ -9,17 +9,18 @@ use crate::util::ceil_div_usize;
 
 /// Compute partial products of the original vector `v` such that all products consist of `max_degree`
 /// or less elements. This is done until we've computed the product `P` of all elements in the vector.
-pub fn partial_products<T: Product + Copy>(v: &[T], max_degree: usize) -> Vec<T> {
+pub fn partial_products<T: MulAssign + Product + Copy>(v: &[T], max_degree: usize) -> Vec<T> {
+    debug_assert!(max_degree > 1);
     let mut res = Vec::new();
-    let mut remainder = v.to_vec();
+    let mut acc = v[0];
-    while remainder.len() > max_degree {
+    let chunk_size = max_degree - 1;
-        let new_partials = remainder
+    let num_chunks = ceil_div_usize(v.len() - 1, chunk_size) - 1;
-            .chunks(max_degree)
+    for i in 0..num_chunks {
-            // TODO: can filter out chunks of length 1.
+        acc *= v[1 + i * chunk_size..1 + (i + 1) * chunk_size]
-            .map(|chunk| chunk.iter().copied().product())
+            .iter()
-            .collect::<Vec<_>>();
+            .copied()
-        res.extend_from_slice(&new_partials);
+            .product();
-        remainder = new_partials;
+        res.push(acc);
     }
 
     res
@@ -29,34 +30,33 @@ pub fn partial_products<T: Product + Copy>(v: &[T], max_degree: usize) -> Vec<T>
 /// vector of length `n`, and `b` is the number of elements needed to compute the final product.
 pub fn num_partial_products(n: usize, max_degree: usize) -> (usize, usize) {
     debug_assert!(max_degree > 1);
-    let mut res = 0;
+    let chunk_size = max_degree - 1;
-    let mut remainder = n;
+    let num_chunks = ceil_div_usize(n - 1, chunk_size) - 1;
-    while remainder > max_degree {
-        let new_partials_len = ceil_div_usize(remainder, max_degree);
-        res += new_partials_len;
-        remainder = new_partials_len;
-    }
 
-    (res, remainder)
+    (num_chunks, 1 + num_chunks * chunk_size)
 }
 
 /// Checks that the partial products of `v` are coherent with those in `partials` by only computing
 /// products of size `max_degree` or less.
-pub fn check_partial_products<T: Product + Copy + Sub<Output = T>>(
+pub fn check_partial_products<T: MulAssign + Product + Copy + Sub<Output = T>>(
     v: &[T],
     mut partials: &[T],
     max_degree: usize,
 ) -> Vec<T> {
+    debug_assert!(max_degree > 1);
+    let mut partials = partials.iter();
     let mut res = Vec::new();
-    let mut remainder = v;
+    let mut acc = v[0];
-    while remainder.len() > max_degree {
+    let chunk_size = max_degree - 1;
-        let products = remainder
+    let num_chunks = ceil_div_usize(v.len() - 1, chunk_size) - 1;
-            .chunks(max_degree)
+    for i in 0..num_chunks {
-            .map(|chunk| chunk.iter().copied().product::<T>());
+        acc *= v[1 + i * chunk_size..1 + (i + 1) * chunk_size]
-        let products_len = products.len();
+            .iter()
-        res.extend(products.zip(partials).map(|(a, &b)| a - b));
+            .copied()
-        (remainder, partials) = partials.split_at(products_len);
+            .product();
+        res.push(acc - *partials.next().unwrap());
     }
+    debug_assert!(partials.next().is_none());
 
     res
 }
@@ -67,22 +67,20 @@ pub fn check_partial_products_recursively<F: RichField + Extendable<D>, const D:
     partials: &[ExtensionTarget<D>],
     max_degree: usize,
 ) -> Vec<ExtensionTarget<D>> {
+    debug_assert!(max_degree > 1);
+    let mut partials = partials.iter();
     let mut res = Vec::new();
-    let mut remainder = v.to_vec();
+    let mut acc = v[0];
-    let mut partials = partials.to_vec();
+    let chunk_size = max_degree - 1;
-    while remainder.len() > max_degree {
+    let num_chunks = ceil_div_usize(v.len() - 1, chunk_size) - 1;
-        let products = remainder
+    for i in 0..num_chunks {
-            .chunks(max_degree)
+        let mut chunk = v[1 + i * chunk_size..1 + (i + 1) * chunk_size].to_vec();
-            .map(|chunk| builder.mul_many_extension(chunk))
+        chunk.push(acc);
-            .collect::<Vec<_>>();
+        acc = builder.mul_many_extension(&chunk);
-        res.extend(
+
-            products
+        res.push(builder.sub_extension(acc, *partials.next().unwrap()));
-                .iter()
-                .zip(&partials)
-                .map(|(&a, &b)| builder.sub_extension(a, b)),
-        );
-        remainder = partials.drain(..products.len()).collect();
     }
+    debug_assert!(partials.next().is_none());
 
     res
 }
@@ -97,15 +95,15 @@ mod tests {
     fn test_partial_products() {
         let v = vec![1, 2, 3, 4, 5, 6];
         let p = partial_products(&v, 2);
-        assert_eq!(p, vec![2, 12, 30, 24, 30]);
+        assert_eq!(p, vec![2, 6, 24, 120]);
         let nums = num_partial_products(v.len(), 2);
         assert_eq!(p.len(), nums.0);
         assert!(check_partial_products(&v, &p, 2)
             .iter()
             .all(|x| x.is_zero()));
         assert_eq!(
+            *p.last().unwrap() * v[nums.1..].iter().copied().product::<i32>(),
             v.into_iter().product::<i32>(),
-            p[p.len() - nums.1..].iter().copied().product(),
         );
 
         let v = vec![1, 2, 3, 4, 5, 6];
@@ -117,8 +115,8 @@ mod tests {
             .iter()
             .all(|x| x.is_zero()));
         assert_eq!(
+            *p.last().unwrap() * v[nums.1..].iter().copied().product::<i32>(),
             v.into_iter().product::<i32>(),
-            p[p.len() - nums.1..].iter().copied().product(),
         );
     }
 }