mithril_stm/merkle_tree/
commitment.rs

1use std::marker::PhantomData;
2
3use blake2::digest::{Digest, FixedOutput};
4use serde::{Deserialize, Serialize};
5
6use crate::error::MerkleTreeError;
7use crate::merkle_tree::{MerkleBatchPath, MerklePath, MerkleTreeLeaf, parent, sibling};
8
9/// `MerkleTree` commitment.
10/// This structure differs from `MerkleTree` in that it does not contain all elements, which are not always necessary.
11/// Instead, it only contains the root of the tree.
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct MerkleTreeCommitment<D: Digest> {
14    /// Root of the merkle commitment.
15    pub root: Vec<u8>,
16    hasher: PhantomData<D>,
17}
18
19impl<D: Digest + FixedOutput> MerkleTreeCommitment<D> {
20    pub(crate) fn new(root: Vec<u8>) -> Self {
21        MerkleTreeCommitment {
22            root,
23            hasher: PhantomData,
24        }
25    }
26
27    /// Check an inclusion proof that `val` is part of the tree by traveling the whole path until the root.
28    /// # Error
29    /// If the merkle tree path is invalid, then the function fails.
30    pub(crate) fn verify_leaf_membership_from_path(
31        &self,
32        val: &MerkleTreeLeaf,
33        proof: &MerklePath<D>,
34    ) -> Result<(), MerkleTreeError<D>>
35    where
36        D: FixedOutput + Clone,
37    {
38        let mut idx = proof.index;
39
40        let mut h = D::digest(val.to_bytes()).to_vec();
41        for p in &proof.values {
42            if (idx & 0b1) == 0 {
43                h = D::new().chain_update(h).chain_update(p).finalize().to_vec();
44            } else {
45                h = D::new().chain_update(p).chain_update(h).finalize().to_vec();
46            }
47            idx >>= 1;
48        }
49
50        if h == self.root {
51            return Ok(());
52        }
53        Err(MerkleTreeError::PathInvalid(proof.clone()))
54    }
55
56    /// Check an inclusion proof that `val` is part of the tree by traveling the whole path until the root.
57    /// # Error
58    /// If the merkle tree path is invalid, then the function fails.
59    #[deprecated(
60        since = "0.5.0",
61        note = "Use `verify_leaf_membership_from_path` instead"
62    )]
63    pub fn check(
64        &self,
65        val: &MerkleTreeLeaf,
66        proof: &MerklePath<D>,
67    ) -> Result<(), MerkleTreeError<D>>
68    where
69        D: FixedOutput + Clone,
70    {
71        Self::verify_leaf_membership_from_path(self, val, proof)
72    }
73
74    /// Serializes the Merkle Tree commitment together with a message in a single vector of bytes.
75    /// Outputs `msg || self` as a vector of bytes.
76    fn concatenate_with_message(&self, msg: &[u8]) -> Vec<u8>
77    where
78        D: Digest,
79    {
80        let mut msgp = msg.to_vec();
81        let mut bytes = self.root.clone();
82        msgp.append(&mut bytes);
83
84        msgp
85    }
86
87    /// Serializes the Merkle Tree commitment together with a message in a single vector of bytes.
88    /// Outputs `msg || self` as a vector of bytes.
89    #[deprecated(since = "0.5.0", note = "Use `concatenate_with_message` instead")]
90    pub fn concat_with_msg(&self, msg: &[u8]) -> Vec<u8>
91    where
92        D: Digest,
93    {
94        Self::concatenate_with_message(self, msg)
95    }
96
97    /// Convert to bytes
98    /// # Layout
99    /// * Root of the Merkle commitment
100    pub fn to_bytes(&self) -> Vec<u8> {
101        let mut output = Vec::new();
102        output.extend_from_slice(&self.root);
103        output
104    }
105
106    /// Extract a `MerkleTreeCommitment` from a byte slice.
107    pub fn from_bytes(bytes: &[u8]) -> Result<MerkleTreeCommitment<D>, MerkleTreeError<D>> {
108        let root = bytes.to_vec();
109
110        Ok(Self {
111            root,
112            hasher: PhantomData,
113        })
114    }
115}
116
117/// Batch compatible `MerkleTree` commitment .
118/// This structure differs from `MerkleTreeCommitment` in that it stores the number of leaves in the tree
119/// as well as the root of the tree.
120/// Number of leaves is required by the batch path generation/verification.
121#[derive(Debug, Clone, Serialize, Deserialize)]
122pub struct MerkleTreeBatchCommitment<D: Digest> {
123    /// Root of the merkle commitment.
124    pub root: Vec<u8>,
125    nr_leaves: usize,
126    hasher: PhantomData<D>,
127}
128
129impl<D: Digest + FixedOutput> MerkleTreeBatchCommitment<D> {
130    pub(crate) fn new(root: Vec<u8>, nr_leaves: usize) -> Self {
131        Self {
132            root,
133            nr_leaves,
134            hasher: Default::default(),
135        }
136    }
137
138    #[cfg(test)]
139    /// Used in property test of `tree`: `test_bytes_tree_commitment_batch_compat`
140    pub(crate) fn get_number_of_leaves(&self) -> usize {
141        self.nr_leaves
142    }
143
144    /// Serializes the Merkle Tree commitment together with a message in a single vector of bytes.
145    /// Outputs `msg || self` as a vector of bytes.
146    // todo: Do we need to concat msg to whole commitment (nr_leaves and root) or just the root?
147    pub(crate) fn concatenate_with_message(&self, msg: &[u8]) -> Vec<u8> {
148        let mut msgp = msg.to_vec();
149        let mut bytes = self.root.clone();
150        msgp.append(&mut bytes);
151
152        msgp
153    }
154
155    /// Serializes the Merkle Tree commitment together with a message in a single vector of bytes.
156    /// Outputs `msg || self` as a vector of bytes.
157    // todo: Do we need to concat msg to whole commitment (nr_leaves and root) or just the root?
158    #[deprecated(since = "0.5.0", note = "Use `concatenate_with_message` instead")]
159    pub fn concat_with_msg(&self, msg: &[u8]) -> Vec<u8> {
160        Self::concatenate_with_message(self, msg)
161    }
162
163    /// Check a proof of a batched opening. The indices must be ordered.
164    ///
165    /// # Error
166    /// Returns an error if the proof is invalid.
167    // todo: Update doc.
168    // todo: Simplify the algorithm.
169    // todo: Maybe we want more granular errors, rather than only `BatchPathInvalid`
170    pub(crate) fn verify_leaves_membership_from_batch_path(
171        &self,
172        batch_val: &[MerkleTreeLeaf],
173        proof: &MerkleBatchPath<D>,
174    ) -> Result<(), MerkleTreeError<D>>
175    where
176        D: FixedOutput + Clone,
177    {
178        if batch_val.len() != proof.indices.len() {
179            return Err(MerkleTreeError::BatchPathInvalid(proof.clone()));
180        }
181        let mut ordered_indices: Vec<usize> = proof.indices.clone();
182        ordered_indices.sort_unstable();
183
184        if ordered_indices != proof.indices {
185            return Err(MerkleTreeError::BatchPathInvalid(proof.clone()));
186        }
187
188        let nr_nodes = self.nr_leaves + self.nr_leaves.next_power_of_two() - 1;
189
190        ordered_indices = ordered_indices
191            .into_iter()
192            .map(|i| i + self.nr_leaves.next_power_of_two() - 1)
193            .collect();
194
195        let mut idx = ordered_indices[0];
196        // First we need to hash the leave values
197        let mut leaves: Vec<Vec<u8>> = batch_val
198            .iter()
199            .map(|val| D::digest(val.to_bytes()).to_vec())
200            .collect();
201
202        let mut values = proof.values.clone();
203
204        while idx > 0 {
205            let mut new_hashes = Vec::with_capacity(ordered_indices.len());
206            let mut new_indices = Vec::with_capacity(ordered_indices.len());
207            let mut i = 0;
208            idx = parent(idx);
209            while i < ordered_indices.len() {
210                new_indices.push(parent(ordered_indices[i]));
211                if ordered_indices[i] & 1 == 0 {
212                    new_hashes.push(
213                        D::new()
214                            .chain(values.first().ok_or(MerkleTreeError::SerializationError)?)
215                            .chain(&leaves[i])
216                            .finalize()
217                            .to_vec(),
218                    );
219                    values.remove(0);
220                } else {
221                    let sibling = sibling(ordered_indices[i]);
222                    if i < ordered_indices.len() - 1 && ordered_indices[i + 1] == sibling {
223                        new_hashes.push(
224                            D::new().chain(&leaves[i]).chain(&leaves[i + 1]).finalize().to_vec(),
225                        );
226                        i += 1;
227                    } else if sibling < nr_nodes {
228                        new_hashes.push(
229                            D::new()
230                                .chain(&leaves[i])
231                                .chain(values.first().ok_or(MerkleTreeError::SerializationError)?)
232                                .finalize()
233                                .to_vec(),
234                        );
235                        values.remove(0);
236                    } else {
237                        new_hashes.push(
238                            D::new().chain(&leaves[i]).chain(D::digest([0u8])).finalize().to_vec(),
239                        );
240                    }
241                }
242                i += 1;
243            }
244            leaves.clone_from(&new_hashes);
245            ordered_indices.clone_from(&new_indices);
246        }
247
248        if leaves.len() == 1 && leaves[0] == self.root {
249            return Ok(());
250        }
251
252        Err(MerkleTreeError::BatchPathInvalid(proof.clone()))
253    }
254
255    /// Check a proof of a batched opening. The indices must be ordered.
256    ///
257    /// # Error
258    /// Returns an error if the proof is invalid.
259    // todo: Update doc.
260    // todo: Simplify the algorithm.
261    // todo: Maybe we want more granular errors, rather than only `BatchPathInvalid`
262    #[deprecated(
263        since = "0.5.0",
264        note = "Use `verify_leaves_membership_from_batch_path` instead"
265    )]
266    pub fn check(
267        &self,
268        batch_val: &[MerkleTreeLeaf],
269        proof: &MerkleBatchPath<D>,
270    ) -> Result<(), MerkleTreeError<D>>
271    where
272        D: FixedOutput + Clone,
273    {
274        Self::verify_leaves_membership_from_batch_path(self, batch_val, proof)
275    }
276
277    /// Convert to bytes
278    /// * Number of leaves as u64
279    /// * Root of the Merkle commitment as bytes
280    pub fn to_bytes(&self) -> Vec<u8> {
281        let mut output = Vec::new();
282        output.extend_from_slice(&u64::try_from(self.nr_leaves).unwrap().to_be_bytes());
283        output.extend_from_slice(&self.root);
284
285        output
286    }
287
288    /// Extract a `MerkleTreeBatchCommitment` from a byte slice.
289    pub fn from_bytes(bytes: &[u8]) -> Result<MerkleTreeBatchCommitment<D>, MerkleTreeError<D>> {
290        let mut u64_bytes = [0u8; 8];
291        u64_bytes.copy_from_slice(bytes.get(..8).ok_or(MerkleTreeError::SerializationError)?);
292        let nr_leaves = usize::try_from(u64::from_be_bytes(u64_bytes))
293            .map_err(|_| MerkleTreeError::SerializationError)?;
294        let mut root = Vec::new();
295        root.extend_from_slice(bytes.get(8..).ok_or(MerkleTreeError::SerializationError)?);
296
297        Ok(MerkleTreeBatchCommitment {
298            root,
299            nr_leaves,
300            hasher: PhantomData,
301        })
302    }
303}
304
305impl<D: Digest> PartialEq for MerkleTreeBatchCommitment<D> {
306    fn eq(&self, other: &Self) -> bool {
307        self.root == other.root && self.nr_leaves == other.nr_leaves
308    }
309}
310
311impl<D: Digest> Eq for MerkleTreeBatchCommitment<D> {}