Modify finalize of DomainEvaluationAccumulator.

Author: leo-starkware

crates/stwo/src/prover/air/accumulation.rs

@@ -86,42 +86,26 @@ impl<B: Backend> DomainEvaluationAccumulator<B> {
             class = "ConstraintInterpolation"
         )
         .entered();
-        let mut cur_poly: Option<SecureCirclePoly<B>> = None;
-        let twiddles = B::precompute_twiddles(
-            CanonicCoset::new(self.log_size())
-                .circle_domain()
-                .half_coset,
-        );
 
-        for (log_size, values) in self.sub_accumulations.into_iter().enumerate().skip(1) {
-            let Some(mut values) = values else {
-                continue;
-            };
-            if let Some(prev_poly) = cur_poly {
-                let eval = SecureColumnByCoords {
-                    columns: prev_poly.0.map(|c| {
-                        c.evaluate_with_twiddles(
-                            CanonicCoset::new(log_size as u32).circle_domain(),
-                            &twiddles,
-                        )
-                        .values
-                    }),
-                };
-                B::accumulate(&mut values, &eval);
-            }
-            cur_poly = Some(SecureCirclePoly(values.columns.map(|c| {
+        let sub_accumulations = self.sub_accumulations.into_iter().flatten().collect_vec();
+        let lifted_accumulation = B::lift_and_accumulate(sub_accumulations);
+
+        if let Some(eval) = lifted_accumulation {
+            let twiddles =
+                B::precompute_twiddles(CanonicCoset::new(log_size).circle_domain().half_coset);
+
+            SecureCirclePoly(eval.columns.map(|c| {
                 CircleEvaluation::<B, BaseField, BitReversedOrder>::new(
-                    CanonicCoset::new(log_size as u32).circle_domain(),
+                    CanonicCoset::new(log_size).circle_domain(),
                     c,
                 )
                 .interpolate_with_twiddles(&twiddles)
-            })));
-        }
-        cur_poly.unwrap_or_else(|| {
+            }))
+        } else {
             SecureCirclePoly(std::array::from_fn(|_| {
                 CirclePoly::new(Col::<B, BaseField>::zeros(1 << log_size))
             }))
-        })
+        }
     }
 }
 
@@ -167,17 +151,15 @@ mod tests {
 
     use super::*;
     use crate::core::circle::CirclePoint;
-    use crate::core::fields::m31::{M31, P};
+    use crate::core::fields::m31::M31;
     use crate::prover::backend::cpu::CpuCircleEvaluation;
     use crate::qm31;
 
     #[test]
-    fn test_domain_evaluation_accumulator() {
-        // Generate a vector of random sizes with a constant seed.
+    fn test_domain_evaluation_accumulator_lifted() {
         let mut rng = SmallRng::seed_from_u64(0);
         const LOG_SIZE_MIN: u32 = 4;
         const LOG_SIZE_BOUND: u32 = 10;
-        const MASK: u32 = P;
         let mut log_sizes = (0..100)
             .map(|_| rng.gen_range(LOG_SIZE_MIN..LOG_SIZE_BOUND))
             .collect::<Vec<_>>();
@@ -188,16 +170,15 @@ mod tests {
             .iter()
             .map(|log_size| {
                 (0..(1 << *log_size))
-                    .map(|_| M31::from_u32_unchecked(rng.gen::<u32>() & MASK))
+                    .map(|_| M31::from(rng.gen::<u32>()))
                     .collect::<Vec<_>>()
             })
             .collect::<Vec<_>>();
         let alpha = qm31!(2, 3, 4, 5);
 
-        // Use accumulator.
         let mut accumulator = DomainEvaluationAccumulator::<CpuBackend>::new(
             alpha,
-            LOG_SIZE_BOUND,
+            LOG_SIZE_BOUND - 1,
             evaluations.len(),
         );
         let n_cols_per_size: [(u32, usize); (LOG_SIZE_BOUND - LOG_SIZE_MIN) as usize] =
@@ -210,6 +191,7 @@ mod tests {
                     .count();
                 (current_log_size, n_cols)
             });
+
         let mut cols = accumulator.columns(n_cols_per_size);
         let mut eval_chunk_offset = 0;
         for (log_size, n_cols) in n_cols_per_size.iter() {
@@ -221,7 +203,6 @@ mod tests {
                     if *log_size != *col_log_size {
                         continue;
                     }
-
                     // The random coefficient powers chunk is in regular order.
                     let random_coeff_chunk =
                         &cols[(log_size - LOG_SIZE_MIN) as usize].random_coeff_powers;
@@ -239,13 +220,18 @@ mod tests {
         let point = CirclePoint::<SecureField>::get_point(98989892);
         let accumulator_res = accumulator_poly.eval_at_point(point);
 
-        // Use direct computation.
+        // Use direct computation: first interpolate each evaluation to obtain a polynomial,
+        // evaluate its lift at `point`, and accumulate over the evaluations.
         let mut res = SecureField::default();
         for (log_size, values) in log_sizes.into_iter().zip(evaluations) {
             res = res * alpha
-                + CpuCircleEvaluation::new(CanonicCoset::new(log_size).circle_domain(), values)
-                    .interpolate()
-                    .eval_at_point(point);
+                + CpuCircleEvaluation::<BaseField, BitReversedOrder>::new(
+                    CanonicCoset::new(log_size).circle_domain(),
+                    values,
+                )
+                .interpolate()
+                // The max log domain size is LOG_SIZE_BOUND - 1.
+                .eval_at_point(point.repeated_double(LOG_SIZE_BOUND - 1 - log_size));
         }
 
         assert_eq!(accumulator_res, res);