commit before fri fold in each basefold_fri_round

Author: kunxian-xia

crates/mpcs/src/basefold.rs

@@ -406,13 +406,13 @@ where
         assert_eq!(commits.len(), num_rounds);
         let sumcheck_messages = proof.sumcheck_proof.as_ref().unwrap();
         for i in 0..num_rounds {
-            write_digest_to_transcript(&commits[i], transcript);
             transcript.append_field_element_exts(sumcheck_messages[i].evaluations.as_slice());
             fold_challenges.push(
                 transcript
                     .sample_and_append_challenge(b"commit round")
                     .elements,
             );
+            write_digest_to_transcript(&commits[i], transcript);
         }
         #[cfg(debug_assertions)]
         {

crates/mpcs/src/basefold/commit_phase.rs

@@ -150,30 +150,6 @@ where
             .collect_vec(),
     );
     exit_span!(batch_codeword_span);
-
-    // commit to codewords with max height using mmcs_ext
-    let max_height = batched_codewords
-        .front()
-        .expect("empty batched_codewords")
-        .height();
-    let mut highest_codeword = batched_codewords.pop_front().unwrap();
-    while let Some(new_codeword) = batched_codewords.front() {
-        if new_codeword.height() == max_height {
-            let new_codeword = batched_codewords.pop_front().unwrap();
-            // sum up the rows in each codeword
-            highest_codeword
-                .par_rows_mut()
-                .zip(new_codeword.par_rows())
-                .for_each(|(row_acc, row)| {
-                    row_acc.iter_mut().zip(row).for_each(|(acc, v)| *acc += v);
-                });
-        }
-    }
-    let (commit, mmcs) = mmcs_ext.commit_matrix(highest_codeword);
-    write_digest_to_transcript(&commit, transcript);
-    trees.push(mmcs);
-    commits.push(commit);
-
     exit_span!(prepare_span);
 
     // eq is the evaluation representation of the eq(X,r) polynomial over the hypercube
@@ -219,6 +195,7 @@ where
     let mut sumcheck_messages = Vec::with_capacity(num_rounds);
 
     let mut challenge = None;
+    let mut running_codeword_opt: Option<RowMajorMatrix<E>> = None;
     let sumcheck_phase1 = entered_span!("sumcheck_phase1");
     let phase1_rounds = num_rounds.min(max_num_vars - log2_num_threads);
 
@@ -228,6 +205,7 @@ where
             &mut prover_states,
             challenge,
             &mut sumcheck_messages,
+            &mut running_codeword_opt,
             &mut batched_codewords,
             transcript,
             &mut trees,
@@ -264,6 +242,7 @@ where
             &mut prover_states,
             challenge,
             &mut sumcheck_messages,
+            &mut running_codeword_opt,
             &mut batched_codewords,
             transcript,
             &mut trees,
@@ -326,6 +305,7 @@ where
 #[allow(clippy::too_many_arguments)]
 pub(crate) fn basefold_fri_round<E: ExtensionField, Spec: BasefoldSpec<E>>(
     pp: &<Spec::EncodingScheme as EncodingScheme<E>>::ProverParameters,
+    running_codeword_opt: &mut Option<RowMajorMatrix<E>>,
     codewords: &mut VecDeque<RowMajorMatrix<E>>,
     trees: &mut Vec<MerkleTreeExt<E>>,
     commits: &mut Vec<<Poseidon2ExtMerkleMmcs<E> as Mmcs<E>>::Commitment>,
@@ -342,11 +322,8 @@ pub(crate) fn basefold_fri_round<E: ExtensionField, Spec: BasefoldSpec<E>>(
     <Poseidon2ExtMerkleMmcs<E> as Mmcs<E>>::Commitment:
         IntoIterator<Item = E::BaseField> + PartialEq,
 {
-    let running_codeword_opt = trees
-        .last()
-        .and_then(|mktree| mmcs_ext.get_matrices(mktree).pop())
-        .map(|m| m.as_view());
     let target_len = running_codeword_opt
+        .as_ref()
         .map(|running_codeword| running_codeword.values.len())
         .unwrap_or_else(|| {
             codewords
@@ -355,7 +332,6 @@ pub(crate) fn basefold_fri_round<E: ExtensionField, Spec: BasefoldSpec<E>>(
                 .max()
                 .expect("empty codeword")
         });
-    let next_level_target_len = target_len >> 1;
     let level = log2_strict_usize(target_len) - 1;
     let folding_coeffs =
         <Spec::EncodingScheme as EncodingScheme<E>>::prover_folding_coeffs_level(pp, level);
@@ -365,76 +341,39 @@ pub(crate) fn basefold_fri_round<E: ExtensionField, Spec: BasefoldSpec<E>>(
     // take codewords match with target length then fold
     let codewords_matched =
         pop_front_while(codewords, |codeword| codeword.values.len() == target_len);
-    // take codewords match next target length in preparation of being committed together
-    let codewords_next_level_matched = pop_front_while(codewords, |codeword| {
-        codeword.values.len() == next_level_target_len
-    });
 
-    // optimize for single codeword match
-    let folded_codeword = if (usize::from(running_codeword_opt.is_some()) + codewords_matched.len())
-        == 1
-        && codewords_next_level_matched.is_empty()
-    {
-        RowMajorMatrix::new(
+    // aggregate codeword with same length
+    let codeword_to_fold = (0..target_len)
+        .into_par_iter()
+        .map(|index| {
             running_codeword_opt
-                .or_else(|| codewords_matched.first().map(|m| m.as_view()))
-                .unwrap()
-                .values
-                .par_chunks_exact(2)
-                .zip(folding_coeffs)
-                .map(|(ys, coeff)| codeword_fold_with_challenge(ys, challenge, *coeff, inv_2))
-                .collect::<Vec<_>>(),
-            2,
-        )
-    } else {
-        // aggregate codeword with same length
-        let codeword_to_fold = (0..target_len)
-            .into_par_iter()
-            .map(|index| {
-                running_codeword_opt
-                    .into_iter()
-                    .chain(codewords_matched.iter().map(|m| m.as_view()))
-                    .map(|codeword| codeword.values[index])
-                    .sum::<E>()
-            })
-            .collect::<Vec<E>>();
-
-        RowMajorMatrix::new(
-            (0..target_len)
-                .into_par_iter()
-                .step_by(2)
-                .map(|index| {
-                    let coeff = &folding_coeffs[index >> 1];
-
-                    // 1st part folded with challenge then sum
-                    let cur_same_pos_sum = codeword_fold_with_challenge(
-                        &codeword_to_fold[index..index + 2],
-                        challenge,
-                        *coeff,
-                        inv_2,
-                    );
-                    // 2nd part: retrieve respective index then sum
-                    let next_same_pos_sum = codewords_next_level_matched
-                        .iter()
-                        .map(|codeword| codeword.values[index >> 1])
-                        .sum::<E>();
-                    cur_same_pos_sum + next_same_pos_sum
-                })
-                .collect::<Vec<_>>(),
-            2,
-        )
-    };
+                .iter()
+                .chain(codewords_matched.iter())
+                .map(|codeword| codeword.values[index])
+                .sum::<E>()
+        })
+        .collect::<Vec<E>>();
+
+    // commit
+    let codeword_as_matrix = RowMajorMatrix::new(codeword_to_fold.clone(), 2);
+    let (commitment, merkle_tree) = mmcs_ext.commit_matrix(codeword_as_matrix);
+    write_digest_to_transcript(&commitment, transcript);
+    commits.push(commitment);
+    trees.push(merkle_tree);
+
+    // fri fold and add codewords with next target length
+    *running_codeword_opt = Some(RowMajorMatrix::new(
+        codeword_to_fold
+            .par_chunks_exact(2)
+            .zip(folding_coeffs)
+            .map(|(ys, coeff)| codeword_fold_with_challenge(ys, challenge, *coeff, inv_2))
+            .collect::<Vec<_>>(),
+        2,
+    ));
 
     if cfg!(feature = "sanity-check") && is_last_round {
-        let (commitment, merkle_tree) = mmcs_ext.commit_matrix(folded_codeword.clone());
-        commits.push(commitment);
-        trees.push(merkle_tree);
-    }
-
-    // skip last round commitment as verifer need to derive encode(final_message) = final_codeword itself
-    if !is_last_round {
-        let (commitment, merkle_tree) = mmcs_ext.commit_matrix(folded_codeword);
-        write_digest_to_transcript(&commitment, transcript);
+        let (commitment, merkle_tree) =
+            mmcs_ext.commit_matrix(running_codeword_opt.as_ref().map(|c| c.clone()).unwrap());
         commits.push(commitment);
         trees.push(merkle_tree);
     }
@@ -447,6 +386,7 @@ fn basefold_one_round<E: ExtensionField, Spec: BasefoldSpec<E>>(
     prover_states: &mut Vec<IOPProverState<'_, E>>,
     challenge: Option<Challenge<E>>,
     sumcheck_messages: &mut Vec<IOPProverMessage<E>>,
+    running_codeword_opt: &mut Option<RowMajorMatrix<E>>,
     codewords: &mut VecDeque<RowMajorMatrix<E>>,
     transcript: &mut impl Transcript<E>,
     trees: &mut Vec<MerkleTreeExt<E>>,
@@ -490,6 +430,7 @@ where
     let fri_round_span = entered_span!("basefold::fri_one_round");
     basefold_fri_round::<E, Spec>(
         pp,
+        running_codeword_opt,
         codewords,
         trees,
         commits,

crates/mpcs/src/basefold/encoding/rs.rs

@@ -334,6 +334,7 @@ mod tests {
             izip!(&codeword.values, &codeword_ext.values).all(|(base, ext)| E::from(*base) == *ext)
         );
 
+        let mut running_codeword_opt = None;
         let mut codeword_ext = VecDeque::from(vec![codeword_ext]);
         let mut transcript = BasicTranscript::new(b"test");
 
@@ -342,6 +343,7 @@ mod tests {
         let r = E::from_canonical_u64(97);
         basefold_fri_round::<E, BasefoldRSParams>(
             &pp,
+            &mut running_codeword_opt,
             &mut codeword_ext,
             &mut prove_data,
             &mut vec![],

crates/mpcs/src/basefold/query_phase.rs

@@ -49,7 +49,6 @@ where
         log2_max_codeword_size,
     );
 
-    println!("log2_max_codeword_size: {}", log2_max_codeword_size);
     queries
         .iter()
         .map(|idx| {