Return opening positions and query positions per column size in FRI

crates/prover/src/core/fri.rs

@@ -265,15 +265,19 @@ impl<B: FriOps + MerkleOps<MC::H>, MC: MerkleChannel> FriProver<B, MC> {
     /// Generates a FRI proof and returns it with the opening positions for the committed columns.
     ///
     /// Returned column opening positions are mapped by their log size.
-    pub fn decommit(
-        self,
-        channel: &mut MC::C,
-    ) -> (FriProof<MC::H>, BTreeMap<u32, SparseSubCircleDomain>) {
+    pub fn decommit(self, channel: &mut MC::C) -> (FriProof<MC::H>, ColumnQueries) {
         let max_column_log_size = self.column_log_sizes[0];
         let queries = Queries::generate(channel, max_column_log_size, self.config.n_queries);
-        let positions = get_opening_positions(&queries, &self.column_log_sizes);
+        let opening_positions = get_opening_positions(&queries, &self.column_log_sizes);
+        let query_positions = get_query_positions(&queries, &self.column_log_sizes);
         let proof = self.decommit_on_queries(&queries);
-        (proof, positions)
+        (
+            proof,
+            ColumnQueries {
+                query_positions,
+                opening_positions,
+            },
+        )
     }
 
     /// # Panics
@@ -501,20 +505,21 @@ impl<MC: MerkleChannel> FriVerifier<MC> {
     /// Samples queries and returns the opening positions for each unique column size.
     ///
     /// The order of the opening positions corresponds to the order of the column commitment.
-    pub fn column_query_positions(
-        &mut self,
-        channel: &mut MC::C,
-    ) -> BTreeMap<u32, SparseSubCircleDomain> {
+    pub fn column_queries(&mut self, channel: &mut MC::C) -> ColumnQueries {
         let column_log_sizes = self
             .column_bounds
             .iter()
             .dedup()
             .map(|b| b.log_degree_bound + self.config.log_blowup_factor)
             .collect_vec();
         let queries = Queries::generate(channel, column_log_sizes[0], self.config.n_queries);
-        let positions = get_opening_positions(&queries, &column_log_sizes);
+        let query_positions = get_query_positions(&queries, &column_log_sizes);
+        let opening_positions = get_opening_positions(&queries, &column_log_sizes);
         self.queries = Some(queries);
-        positions
+        ColumnQueries {
+            query_positions,
+            opening_positions,
+        }
     }
 }
 
@@ -541,6 +546,26 @@ fn get_opening_positions(
     positions
 }
 
+/// Returns the column query positions needed for verification.
+///
+/// The column log sizes must be unique and in descending order.
+/// Returned column query positions are mapped by their log size.
+fn get_query_positions(queries: &Queries, column_log_sizes: &[u32]) -> BTreeMap<u32, Vec<usize>> {
+    let mut prev_log_size = column_log_sizes[0];
+    assert!(prev_log_size == queries.log_domain_size);
+    let mut prev_queries = queries.clone();
+    let mut positions = BTreeMap::new();
+    positions.insert(prev_log_size, prev_queries.positions.clone());
+    for &log_size in column_log_sizes.iter().skip(1) {
+        let n_folds = prev_log_size - log_size;
+        let queries = prev_queries.fold(n_folds);
+        positions.insert(log_size, queries.positions.clone());
+        prev_log_size = log_size;
+        prev_queries = queries;
+    }
+    positions
+}
+
 #[derive(Clone, Copy, Debug, Error)]
 pub enum FriVerificationError {
     #[error("proof contains an invalid number of FRI layers")]
@@ -613,6 +638,15 @@ pub struct FriProof<H: MerkleHasher> {
     pub last_layer_poly: LinePoly,
 }
 
+/// Values of interest obtained from the execution of the FRI protocol.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ColumnQueries {
+    /// Query positions mapped by log size.
+    pub query_positions: BTreeMap<u32, Vec<usize>>,
+    /// Opening positions mapped by log size.
+    pub opening_positions: BTreeMap<u32, SparseSubCircleDomain>,
+}
+
 /// Number of folds for univariate polynomials.
 // TODO(andrew): Support different step sizes.
 pub const FOLD_STEP: u32 = 1;
@@ -1160,16 +1194,19 @@ mod tests {
             &evaluations,
             &CpuBackend::precompute_twiddles(evaluations[0].domain.half_coset),
         );
-        let (proof, prover_opening_positions) = prover.decommit(&mut test_channel());
-        let decommitment_values = zip(&evaluations, prover_opening_positions.values().rev())
-            .map(|(poly, positions)| open_polynomial(poly, positions))
-            .collect();
+        let (proof, prover_col_queries) = prover.decommit(&mut test_channel());
+        let decommitment_values = zip(
+            &evaluations,
+            prover_col_queries.opening_positions.values().rev(),
+        )
+        .map(|(poly, positions)| open_polynomial(poly, positions))
+        .collect();
         let bounds = LOG_DEGREES.map(CirclePolyDegreeBound::new).to_vec();
 
         let mut verifier = FriVerifier::commit(&mut test_channel(), config, proof, bounds).unwrap();
-        let verifier_opening_positions = verifier.column_query_positions(&mut test_channel());
+        let verifier_col_queries = verifier.column_queries(&mut test_channel());
 
-        assert_eq!(prover_opening_positions, verifier_opening_positions);
+        assert_eq!(prover_col_queries, verifier_col_queries);
         verifier.decommit(decommitment_values)
     }
 

crates/prover/src/core/pcs/prover.rs

@@ -126,16 +126,18 @@ impl<'a, B: BackendForChannel<MC>, MC: MerkleChannel> CommitmentSchemeProver<'a,
         channel.mix_u64(proof_of_work);
 
         // FRI decommitment phase.
-        let (fri_proof, fri_query_domains) = fri_prover.decommit(channel);
+        let (fri_proof, col_queries) = fri_prover.decommit(channel);
 
         // Decommit the FRI queries on the merkle trees.
-        let decommitment_results = self.trees.as_ref().map(|tree| {
-            let queries = fri_query_domains
-                .iter()
-                .map(|(&log_size, domain)| (log_size, domain.flatten()))
-                .collect();
-            tree.decommit(&queries)
-        });
+        let opening_positions = col_queries
+            .opening_positions
+            .iter()
+            .map(|(&log_size, domain)| (log_size, domain.flatten()))
+            .collect();
+        let decommitment_results = self
+            .trees
+            .as_ref()
+            .map(|tree| tree.decommit(&opening_positions));
 
         let queried_values = decommitment_results.as_ref().map(|(v, _)| v.clone());
         let decommitments = decommitment_results.map(|(_, d)| d);

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

@@ -88,15 +88,16 @@ impl<MC: MerkleChannel> CommitmentSchemeVerifier<MC> {
         }
 
         // Get FRI query domains.
-        let fri_query_domains = fri_verifier.column_query_positions(channel);
+        let column_queries = fri_verifier.column_queries(channel);
 
         // Verify merkle decommitments.
         self.trees
             .as_ref()
             .zip_eq(proof.decommitments)
             .zip_eq(proof.queried_values.clone())
             .map(|((tree, decommitment), queried_values)| {
-                let queries = fri_query_domains
+                let queries = column_queries
+                    .opening_positions
                     .iter()
                     .map(|(&log_size, domain)| (log_size, domain.flatten()))
                     .collect();
@@ -120,7 +121,7 @@ impl<MC: MerkleChannel> CommitmentSchemeVerifier<MC> {
             self.column_log_sizes().flatten().into_iter().collect(),
             &samples,
             random_coeff,
-            fri_query_domains,
+            column_queries.opening_positions,
             &proof.queried_values.flatten(),
         )?;
 

crates/prover/src/core/queries.rs

@@ -85,7 +85,7 @@ impl Deref for Queries {
     }
 }
 
-#[derive(Debug, Eq, PartialEq)]
+#[derive(Debug, Clone, Eq, PartialEq)]
 pub struct SparseSubCircleDomain {
     pub domains: Vec<SubCircleDomain>,
     pub large_domain_log_size: u32,
@@ -109,7 +109,7 @@ impl Deref for SparseSubCircleDomain {
 
 /// Represents a circle domain relative to a larger circle domain. The `initial_index` is the bit
 /// reversed query index in the larger domain.
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
 pub struct SubCircleDomain {
     pub coset_index: usize,
     pub log_size: u32,