Adapt the queries for the preprocessed trees. (#1281)

Co-authored-by: Alon F <alonf@starkware.co>

Author: leo-starkware

crates/examples/src/poseidon/mod.rs

@@ -487,6 +487,7 @@ mod tests {
         );
     }
 
+    #[ignore = "AIRs with constraint degree >= 2 are not supported yet in the lifted protocol."]
     #[test_log::test]
     fn test_simd_poseidon_prove() {
         // Note: To see time measurement, run test with
@@ -530,6 +531,7 @@ mod tests {
         verify(&[&component], channel, commitment_scheme, proof).unwrap();
     }
 
+    #[ignore = "AIRs with constraint degree >= 2 are not supported yet in the lifted protocol."]
     #[cfg(feature = "tracing")]
     #[test]
     fn trace_simd_poseidon_prove() {

crates/stwo/src/core/pcs/mod.rs

@@ -8,7 +8,7 @@
 //! Note: Opened points cannot come from the commitment domain.
 
 pub mod quotients;
-mod utils;
+pub mod utils;
 mod verifier;
 
 use serde::{Deserialize, Serialize};

crates/stwo/src/core/pcs/utils.rs

@@ -174,3 +174,23 @@ impl<T> TreeVec<ColumnVec<Vec<T>>> {
         self.0.into_iter().flatten().flatten().collect()
     }
 }
+
+pub fn prepare_preprocessed_query_positions(
+    query_positions: &[usize],
+    max_log_size: u32,
+    pp_max_log_size: u32,
+) -> Vec<usize> {
+    if pp_max_log_size == 0 {
+        return vec![];
+    };
+    if max_log_size < pp_max_log_size {
+        return query_positions
+            .iter()
+            .map(|pos| (pos >> 1 << (pp_max_log_size - max_log_size + 1)) + (pos & 1))
+            .collect();
+    }
+    query_positions
+        .iter()
+        .map(|pos| (pos >> (max_log_size - pp_max_log_size + 1) << 1) + (pos & 1))
+        .collect()
+}

crates/stwo/src/core/pcs/verifier.rs

@@ -11,6 +11,7 @@ use super::utils::TreeVec;
 use super::PcsConfig;
 use crate::core::channel::{Channel, MerkleChannel};
 use crate::core::pcs::quotients::CommitmentSchemeProof;
+use crate::core::pcs::utils::prepare_preprocessed_query_positions;
 use crate::core::vcs_lifted::merkle_hasher::MerkleHasherLifted;
 use crate::core::vcs_lifted::verifier::MerkleVerifierLifted;
 use crate::core::verifier::VerificationError;
@@ -62,31 +63,69 @@ impl<MC: MerkleChannel> CommitmentSchemeVerifier<MC> {
     ) -> Result<(), VerificationError> {
         channel.mix_felts(&proof.sampled_values.clone().flatten_cols());
         let random_coeff = channel.draw_secure_felt();
-        let max_log_size = *self.column_log_sizes().iter().flatten().max().unwrap();
+        // The lifting log size is the length of the longest column which has at least one sample
+        // (i.e. a column which is actually used in the constraints). Usually, the only columns
+        // that have an empty vector of samples are among the preprocessed columns.
+        let lifting_log_size = self
+            .column_log_sizes()
+            .zip_cols(&sampled_points)
+            .flatten()
+            .iter()
+            .filter(|(_, sampled_points)| !sampled_points.is_empty())
+            .map(|(log_size, _)| *log_size)
+            .max()
+            .unwrap();
 
         let bound =
-            CirclePolyDegreeBound::new(max_log_size - self.config.fri_config.log_blowup_factor);
+            CirclePolyDegreeBound::new(lifting_log_size - self.config.fri_config.log_blowup_factor);
 
         // FRI commitment phase on OODS quotients.
         let mut fri_verifier =
             FriVerifier::<MC>::commit(channel, self.config.fri_config, proof.fri_proof, bound)?;
 
         // Verify proof of work.
-
         if !channel.verify_pow_nonce(self.config.pow_bits, proof.proof_of_work) {
             return Err(VerificationError::ProofOfWork);
         }
         channel.mix_u64(proof.proof_of_work);
         // Get FRI query positions.
         let query_positions = fri_verifier.sample_query_positions(channel);
-        // Verify merkle decommitments.
+        let preprocessed_query_positions = prepare_preprocessed_query_positions(
+            &query_positions,
+            lifting_log_size,
+            self.column_log_sizes()[0]
+                .iter()
+                .max()
+                .copied()
+                .unwrap_or_default(),
+        );
+
+        // Build the query positions tree: the preprocessed tree needs a different treatment than
+        // the other trees.
+        let query_positions_tree = TreeVec::new(
+            self.trees
+                .iter()
+                .enumerate()
+                .map(|(i, _)| {
+                    if i == 0 {
+                        preprocessed_query_positions.as_slice()
+                    } else {
+                        query_positions.as_slice()
+                    }
+                })
+                .collect::<Vec<_>>(),
+        );
+        // Verify decommitments.
         self.trees
             .as_ref()
             .zip_eq(proof.decommitments)
             .zip_eq(proof.queried_values.clone())
-            .map(|((tree, decommitment), queried_values)| {
-                tree.verify(&query_positions, queried_values, decommitment)
-            })
+            .zip_eq(query_positions_tree)
+            .map(
+                |(((tree, decommitment), queried_values), query_positions)| {
+                    tree.verify(query_positions, queried_values, decommitment)
+                },
+            )
             .0
             .into_iter()
             .collect::<Result<(), _>>()?;

crates/stwo/src/core/vcs_lifted/verifier.rs

@@ -103,6 +103,7 @@ impl<H: MerkleHasherLifted> MerkleVerifierLifted<H> {
         let mut prev_layer_hashes: Vec<(usize, H::Hash)> = query_positions
             .iter()
             .zip_eq(queried_values.chunks_exact(self.column_log_sizes.len()))
+            .dedup_by(|(idx, _), (idx2, _)| idx == idx2)
             .map(|(idx, column_values)| {
                 let mut hasher = H::default_with_initial_state();
                 hasher.update_leaf(column_values);

crates/stwo/src/prover/pcs/mod.rs

@@ -12,6 +12,7 @@ use crate::core::fields::qm31::SecureField;
 use crate::core::pcs::quotients::{
     CommitmentSchemeProof, CommitmentSchemeProofAux, ExtendedCommitmentSchemeProof, PointSample,
 };
+use crate::core::pcs::utils::prepare_preprocessed_query_positions;
 use crate::core::pcs::{PcsConfig, TreeSubspan, TreeVec};
 use crate::core::poly::circle::CanonicCoset;
 use crate::core::vcs_lifted::merkle_hasher::MerkleHasherLifted;
@@ -199,13 +200,30 @@ impl<'a, B: BackendForChannel<MC>, MC: MerkleChannel> CommitmentSchemeProver<'a,
             unsorted_query_locations,
         } = fri_prover.decommit(channel);
 
-        // Decommit the FRI queries on the merkle trees.
-        let decommitment_results = self
+        // Build the query position tree.
+        let preprocessed_query_positions = prepare_preprocessed_query_positions(
+            &query_positions,
+            max_log_size,
+            self.trees[0].commitment.layers.len() as u32 - 1,
+        );
+        let query_positions_tree = TreeVec::new(
+            self.trees
+                .iter()
+                .enumerate()
+                .map(|(i, _)| {
+                    if i == 0 {
+                        preprocessed_query_positions.as_slice()
+                    } else {
+                        query_positions.as_slice()
+                    }
+                })
+                .collect::<Vec<_>>(),
+        );
+        let (queried_values, decommitments, aux): (Vec<_>, Vec<_>, Vec<_>) = self
             .trees
             .as_ref()
-            .map(|tree| tree.decommit(&query_positions));
-
-        let (queried_values, decommitments, aux): (Vec<_>, Vec<_>, Vec<_>) = decommitment_results
+            .zip_eq(query_positions_tree)
+            .map(|(tree, query_positions)| tree.decommit(query_positions))
             .0
             .into_iter()
             .map(|(v, x)| (v, x.decommitment, x.aux))

crates/stwo/src/prover/vcs_lifted/prover.rs

@@ -89,6 +89,7 @@ impl<B: MerkleOpsLifted<H>, H: MerkleHasherLifted> MerkleProverLifted<B, H> {
         }
 
         let mut prev_layer_queries = queries_position.to_vec();
+        prev_layer_queries.dedup();
         // The largest log size of a layer is equal to `self.layers.len() - 1`. We start iterating
         // from the layer of log size `self.layers.len() - 2` so that we always have a previous
         // layer available for the computation.