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.
@@ -156,20 +157,29 @@ impl LogupTraceGenerator {
         }
     }
 
-    /// Finalize the trace. Returns the trace and the claimed sum of the last column.
-    pub fn finalize(
+    /// Finalize the trace. Returns the trace and the prefix sum of the last column at
+    /// the corresponding `indices`.
+    pub fn finalize<const N: usize>(
         mut self,
+        indices: [usize; N],
     ) -> (
         ColumnVec<CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>>,
-        SecureField,
+        [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,
         };
-        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
@@ -181,7 +191,7 @@ impl LogupTraceGenerator {
                 })
             })
             .collect_vec();
-        (trace, total_sum)
+        (trace, returned_prefix_sums)
     }
 }
 

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

@@ -277,5 +277,6 @@ pub fn generate_interaction_trace(
     }
     col_gen.finalize_col();
 
-    logup_gen.finalize()
+    let (trace, [total_sum]) = logup_gen.finalize([(1 << log_size) - 1]);
+    (trace, total_sum)
 }

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

@@ -167,5 +167,6 @@ pub fn gen_interaction_trace(
     }
     col_gen.finalize_col();
 
-    logup_gen.finalize()
+    let (trace, [total_sum]) = logup_gen.finalize([(1 << log_size) - 1]);
+    (trace, total_sum)
 }

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

@@ -132,7 +132,9 @@ pub fn generate_interaction_trace<const ELEM_BITS: u32, const EXPAND_BITS: u32>(
         }
     }
 
-    logup_gen.finalize()
+    let (trace, [total_sum]) =
+        logup_gen.finalize([(1 << column_bits::<ELEM_BITS, EXPAND_BITS>()) - 1]);
+    (trace, total_sum)
 }
 
 /// Generates the constant trace for the xor table.

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

@@ -141,7 +141,8 @@ pub fn gen_interaction_trace(
     }
     col_gen.finalize_col();
 
-    logup_gen.finalize()
+    let (trace, [total_sum]) = logup_gen.finalize([(1 << log_size) - 1]);
+    (trace, total_sum)
 }
 
 #[allow(unused)]

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

@@ -325,7 +325,8 @@ pub fn gen_interaction_trace(
         col_gen.finalize_col();
     }
 
-    logup_gen.finalize()
+    let (trace, [total_sum]) = logup_gen.finalize([(1 << log_size) - 1]);
+    (trace, total_sum)
 }
 
 pub fn prove_poseidon(