stwo_gpu/core/vcs/
verifier.rs

1use hashbrown::HashMap;
2use itertools::Itertools;
3use serde::{Deserialize, Serialize};
4use std_shims::{vec, BTreeMap, Vec};
5use thiserror::Error;
6
7use super::utils::{next_decommitment_node, option_flatten_peekable};
8use crate::core::fields::m31::BaseField;
9use crate::core::utils::PeekableExt;
10use crate::core::vcs::MerkleHasher;
11
12#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, PartialOrd)]
13pub struct MerkleDecommitment<H: MerkleHasher> {
14    /// Hash values that the verifier needs but cannot deduce from previous computations, in the
15    /// order they are needed.
16    pub hash_witness: Vec<H::Hash>,
17    /// Column values that the verifier needs but cannot deduce from previous computations, in the
18    /// order they are needed.
19    /// This complements the column values that were queried. These must be supplied directly to
20    /// the verifier.
21    pub column_witness: Vec<BaseField>,
22}
23impl<H: MerkleHasher> MerkleDecommitment<H> {
24    pub const fn empty() -> Self {
25        Self {
26            hash_witness: Vec::new(),
27            column_witness: Vec::new(),
28        }
29    }
30}
31
32/// Auxiliary data for Merkle decommitment.
33#[derive(Clone, Debug)]
34pub struct MerkleDecommitmentAux<H: MerkleHasher> {
35    /// For each layer, a map from node index to its hash value.
36    pub all_node_values: Vec<HashMap<usize, H::Hash>>,
37}
38
39pub struct ExtendedMerkleDecommitment<H: MerkleHasher> {
40    pub decommitment: MerkleDecommitment<H>,
41    pub aux: MerkleDecommitmentAux<H>,
42}
43
44pub struct MerkleVerifier<H: MerkleHasher> {
45    pub root: H::Hash,
46    pub column_log_sizes: Vec<u32>,
47    pub n_columns_per_log_size: BTreeMap<u32, usize>,
48}
49impl<H: MerkleHasher> MerkleVerifier<H> {
50    pub fn new(root: H::Hash, column_log_sizes: Vec<u32>) -> Self {
51        let mut n_columns_per_log_size = BTreeMap::new();
52        for log_size in &column_log_sizes {
53            *n_columns_per_log_size.entry(*log_size).or_insert(0) += 1;
54        }
55
56        Self {
57            root,
58            column_log_sizes,
59            n_columns_per_log_size,
60        }
61    }
62    /// Verifies the decommitment of the columns.
63    ///
64    /// Returns `Ok(())` if the decommitment is successfully verified.
65    ///
66    /// # Arguments
67    ///
68    /// * `queries_per_log_size` - A map from log_size to a vector of queries for columns of that
69    ///   log_size.
70    /// * `queried_values` - A vector of queried values according to the order in
71    ///   [`MerkleProver::decommit()`].
72    /// * `decommitment` - The decommitment object containing the witness and column values.
73    ///
74    /// # Errors
75    ///
76    /// Returns an error if any of the following conditions are met:
77    ///
78    /// * The witness is too long (not fully consumed).
79    /// * The witness is too short (missing values).
80    /// * Too many queried values (not fully consumed).
81    /// * Too few queried values (missing values).
82    /// * The computed root does not match the expected root.
83    ///
84    /// [`MerkleProver::decommit()`]: crate::core::...::MerkleProver::decommit
85    pub fn verify(
86        &self,
87        queries_per_log_size: &BTreeMap<u32, Vec<usize>>,
88        queried_values: Vec<BaseField>,
89        decommitment: MerkleDecommitment<H>,
90    ) -> Result<(), MerkleVerificationError> {
91        let Some(max_log_size) = self.column_log_sizes.iter().max() else {
92            return Ok(());
93        };
94
95        let mut queried_values = queried_values.into_iter();
96
97        // Prepare read buffers.
98
99        let mut hash_witness = decommitment.hash_witness.into_iter();
100        let mut column_witness = decommitment.column_witness.into_iter();
101
102        let mut last_layer_hashes: Option<Vec<(usize, H::Hash)>> = None;
103        for layer_log_size in (0..=*max_log_size).rev() {
104            let n_columns_in_layer = *self
105                .n_columns_per_log_size
106                .get(&layer_log_size)
107                .unwrap_or(&0);
108
109            // Prepare write buffer for queries to the current layer. This will propagate to the
110            // next layer.
111            let mut layer_total_queries = vec![];
112
113            // Queries to this layer come from queried node in the previous layer and queried
114            // columns in this one.
115            let mut prev_layer_queries = last_layer_hashes
116                .iter()
117                .flatten()
118                .map(|(q, _)| *q)
119                .collect_vec()
120                .into_iter()
121                .peekable();
122            let mut prev_layer_hashes = last_layer_hashes.as_ref().map(|x| x.iter().peekable());
123            let mut layer_column_queries =
124                option_flatten_peekable(queries_per_log_size.get(&layer_log_size));
125
126            // Merge previous layer queries and column queries.
127            while let Some(node_index) =
128                next_decommitment_node(&mut prev_layer_queries, &mut layer_column_queries)
129            {
130                prev_layer_queries
131                    .peek_take_while(|q| q / 2 == node_index)
132                    .for_each(drop);
133
134                let node_hashes = prev_layer_hashes
135                    .as_mut()
136                    .map(|prev_layer_hashes| {
137                        {
138                            // If the left child was not computed, read it from the witness.
139                            let left_hash = prev_layer_hashes
140                                .next_if(|(index, _)| *index == 2 * node_index)
141                                .map(|(_, hash)| Ok(*hash))
142                                .unwrap_or_else(|| {
143                                    hash_witness
144                                        .next()
145                                        .ok_or(MerkleVerificationError::WitnessTooShort)
146                                })?;
147
148                            // If the right child was not computed, read it to from the witness.
149                            let right_hash = prev_layer_hashes
150                                .next_if(|(index, _)| *index == 2 * node_index + 1)
151                                .map(|(_, hash)| Ok(*hash))
152                                .unwrap_or_else(|| {
153                                    hash_witness
154                                        .next()
155                                        .ok_or(MerkleVerificationError::WitnessTooShort)
156                                })?;
157                            Ok((left_hash, right_hash))
158                        }
159                    })
160                    .transpose()?;
161
162                // If the column values were queried, read them from `queried_value`.
163                let (err, node_values_iter) = match layer_column_queries.next_if_eq(&node_index) {
164                    Some(_) => (
165                        MerkleVerificationError::TooFewQueriedValues,
166                        &mut queried_values,
167                    ),
168                    // Otherwise, read them from the witness.
169                    None => (
170                        MerkleVerificationError::WitnessTooShort,
171                        &mut column_witness,
172                    ),
173                };
174
175                let node_values = node_values_iter.take(n_columns_in_layer).collect_vec();
176                if node_values.len() != n_columns_in_layer {
177                    return Err(err);
178                }
179
180                layer_total_queries.push((node_index, H::hash_node(node_hashes, &node_values)));
181            }
182
183            last_layer_hashes = Some(layer_total_queries);
184        }
185
186        // Check that all witnesses and values have been consumed.
187        if hash_witness.next().is_some() {
188            return Err(MerkleVerificationError::WitnessTooLong);
189        }
190        if queried_values.next().is_some() {
191            return Err(MerkleVerificationError::TooManyQueriedValues);
192        }
193        if column_witness.next().is_some() {
194            return Err(MerkleVerificationError::WitnessTooLong);
195        }
196
197        let [(_, computed_root)] = last_layer_hashes.unwrap().try_into().unwrap();
198        if computed_root != self.root {
199            return Err(MerkleVerificationError::RootMismatch);
200        }
201
202        Ok(())
203    }
204}
205
206// TODO(ilya): Make error messages consistent.
207#[derive(Clone, Copy, Debug, Error, PartialEq, Eq)]
208pub enum MerkleVerificationError {
209    #[error("Witness is too short")]
210    WitnessTooShort,
211    #[error("Witness is too long.")]
212    WitnessTooLong,
213    #[error("too many Queried values")]
214    TooManyQueriedValues,
215    #[error("too few queried values")]
216    TooFewQueriedValues,
217    #[error("Root mismatch.")]
218    RootMismatch,
219}
stwo_gpu/core/vcs/verifier.rs

stwo_gpu/core/vcs/
verifier.rs
