use iterators instead of slices for looping over polys in merge

a16z · Jan 25, 2024 · 2307aa6 · 2307aa6
1 parent 570831e
commit 2307aa6
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 8 deletions.
diff --git a/jolt-core/src/jolt/vm/instruction_lookups.rs b/jolt-core/src/jolt/vm/instruction_lookups.rs
@@ -122,7 +122,7 @@ where
     fn batch(&self) -> Self::BatchedPolynomials {
         use rayon::prelude::*;
         let (batched_dim_read, (batched_final, batched_E, batched_flag)) = rayon::join(
-            || DensePolynomial::merge(&[self.dim.as_slice(), self.read_cts.as_slice()].concat()),
+            || DensePolynomial::merge(self.dim.iter().chain(&self.read_cts)),
             || {
                 let batched_final = DensePolynomial::merge(&self.final_cts);
                 let (batched_E, batched_flag) = rayon::join(

diff --git a/jolt-core/src/lasso/surge.rs b/jolt-core/src/lasso/surge.rs
@@ -65,7 +65,7 @@ where
     #[tracing::instrument(skip_all, name = "SurgePolys::batch")]
     fn batch(&self) -> Self::BatchedPolynomials {
         let (batched_dim_read, (batched_final, batched_E)) = rayon::join(
-            || DensePolynomial::merge(&[self.dim.as_slice(), self.read_cts.as_slice()].concat()),
+            || DensePolynomial::merge(self.dim.iter().chain(&self.read_cts)),
             || {
                 rayon::join(
                     || DensePolynomial::merge(&self.final_cts),

diff --git a/jolt-core/src/poly/dense_mlpoly.rs b/jolt-core/src/poly/dense_mlpoly.rs
@@ -398,13 +398,14 @@ impl<F: PrimeField> DensePolynomial<F> {
     }
 
     #[tracing::instrument(skip_all, name = "DensePoly.merge")]
-    pub fn merge<T>(polys: &[T]) -> DensePolynomial<F>
-    where
-        T: AsRef<DensePolynomial<F>>,
-    {
-        let total_len: usize = polys.iter().map(|poly| poly.as_ref().vec().len()).sum();
+    pub fn merge(polys: impl IntoIterator<Item = impl AsRef<Self>> + Clone) -> DensePolynomial<F> {
+        let polys_iter_cloned = polys.clone().into_iter();
+        let total_len: usize = polys
+            .into_iter()
+            .map(|poly| poly.as_ref().vec().len())
+            .sum();
         let mut Z: Vec<F> = Vec::with_capacity(total_len.next_power_of_two());
-        for poly in polys {
+        for poly in polys_iter_cloned {
             Z.extend_from_slice(poly.as_ref().vec());
         }