Logup generator returns multiple prefix sums

crates/prover/src/constraint_framework/logup.rs

@@ -18,6 +18,7 @@ use crate::core::fields::FieldExpOps;
 use crate::core::lookups::utils::Fraction;
 use crate::core::poly::circle::{CanonicCoset, CircleEvaluation};
 use crate::core::poly::BitReversedOrder;
+use crate::core::utils::{bit_reverse_index, coset_index_to_circle_domain_index};
 use crate::core::ColumnVec;
 
 /// Evaluates constraints for batched logups.
@@ -164,34 +165,53 @@ impl LogupTraceGenerator {
         }
     }
 
-    /// Finalize the trace. Returns the trace and the claimed sum of the last column.
-    pub fn finalize(
-        mut self,
+    /// Finalize the trace. Returns the trace and the total sum of the last column.
+    pub fn finalize_last(
+        self,
     ) -> (
         ColumnVec<CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>>,
         SecureField,
+    ) {
+        let log_size = self.log_size;
+        let (trace, [total_sum]) = self.finalize_at([(1 << log_size) - 1]);
+        (trace, total_sum)
+    }
+
+    /// Finalize the trace. Returns the trace and the prefix sum of the last column at
+    /// the corresponding `indices`.
+    pub fn finalize_at<const N: usize>(
+        mut self,
+        indices: [usize; N],
+    ) -> (
+        ColumnVec<CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>>,
+        [SecureField; N],
     ) {
         // Prefix sum the last column.
         let last_col_coords = self.trace.pop().unwrap().columns;
         let coord_prefix_sum = last_col_coords.map(inclusive_prefix_sum);
         let secure_prefix_sum = SecureColumnByCoords {
             columns: coord_prefix_sum,
         };
-
-        // The last element in bit-reversed circle domain order is at index 1.
-        let total_sum = secure_prefix_sum.at(1);
+        let returned_prefix_sums = indices.map(|idx| {
+            // Prefix sum column is in bit-reversed circle domain order.
+            let fixed_index = bit_reverse_index(
+                coset_index_to_circle_domain_index(idx, self.log_size),
+                self.log_size,
+            );
+            secure_prefix_sum.at(fixed_index)
+        });
         self.trace.push(secure_prefix_sum);
 
         let trace = self
             .trace
             .into_iter()
             .flat_map(|eval| {
-                eval.columns.map(|c| {
-                    CircleEvaluation::new(CanonicCoset::new(self.log_size).circle_domain(), c)
+                eval.columns.map(|col| {
+                    CircleEvaluation::new(CanonicCoset::new(self.log_size).circle_domain(), col)
                 })
             })
             .collect_vec();
-        (trace, total_sum)
+        (trace, returned_prefix_sums)
     }
 }
 

crates/prover/src/examples/blake/round/gen.rs

@@ -277,5 +277,5 @@ pub fn generate_interaction_trace(
     }
     col_gen.finalize_col();
 
-    logup_gen.finalize()
+    logup_gen.finalize_last()
 }

crates/prover/src/examples/blake/scheduler/gen.rs

@@ -167,5 +167,5 @@ pub fn gen_interaction_trace(
     }
     col_gen.finalize_col();
 
-    logup_gen.finalize()
+    logup_gen.finalize_last()
 }

crates/prover/src/examples/blake/xor_table/gen.rs

@@ -132,7 +132,7 @@ pub fn generate_interaction_trace<const ELEM_BITS: u32, const EXPAND_BITS: u32>(
         }
     }
 
-    logup_gen.finalize()
+    logup_gen.finalize_last()
 }
 
 /// Generates the constant trace for the xor table.

crates/prover/src/examples/plonk/mod.rs

@@ -141,7 +141,7 @@ pub fn gen_interaction_trace(
     }
     col_gen.finalize_col();
 
-    logup_gen.finalize()
+    logup_gen.finalize_last()
 }
 
 #[allow(unused)]

crates/prover/src/examples/poseidon/mod.rs

@@ -326,7 +326,7 @@ pub fn gen_interaction_trace(
         col_gen.finalize_col();
     }
 
-    logup_gen.finalize()
+    logup_gen.finalize_last()
 }
 
 pub fn prove_poseidon(