diff --git a/crates/prover/src/constraint_framework/logup.rs b/crates/prover/src/constraint_framework/logup.rs index e7287bac0..21ce14b4a 100644 --- a/crates/prover/src/constraint_framework/logup.rs +++ b/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,12 +165,26 @@ 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>, 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( + mut self, + indices: [usize; N], + ) -> ( + ColumnVec>, + [SecureField; N], ) { // Prefix sum the last column. let last_col_coords = self.trace.pop().unwrap().columns; @@ -177,21 +192,26 @@ impl LogupTraceGenerator { 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) } } diff --git a/crates/prover/src/examples/blake/round/gen.rs b/crates/prover/src/examples/blake/round/gen.rs index 7adbe6fd5..a4f29de44 100644 --- a/crates/prover/src/examples/blake/round/gen.rs +++ b/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() } diff --git a/crates/prover/src/examples/blake/scheduler/gen.rs b/crates/prover/src/examples/blake/scheduler/gen.rs index 4f9e6b1e3..e5dfc2549 100644 --- a/crates/prover/src/examples/blake/scheduler/gen.rs +++ b/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() } diff --git a/crates/prover/src/examples/blake/xor_table/gen.rs b/crates/prover/src/examples/blake/xor_table/gen.rs index 59d3fd92d..ed07d48bf 100644 --- a/crates/prover/src/examples/blake/xor_table/gen.rs +++ b/crates/prover/src/examples/blake/xor_table/gen.rs @@ -132,7 +132,7 @@ pub fn generate_interaction_trace( } } - logup_gen.finalize() + logup_gen.finalize_last() } /// Generates the constant trace for the xor table. diff --git a/crates/prover/src/examples/plonk/mod.rs b/crates/prover/src/examples/plonk/mod.rs index 1cbe3a82b..8117d7f3f 100644 --- a/crates/prover/src/examples/plonk/mod.rs +++ b/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)] diff --git a/crates/prover/src/examples/poseidon/mod.rs b/crates/prover/src/examples/poseidon/mod.rs index d89235f22..125a94fb6 100644 --- a/crates/prover/src/examples/poseidon/mod.rs +++ b/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(