mithril_stm/merkle_tree/
commitment.rs

1use crate::error::MerkleTreeError;
2use crate::merkle_tree::{parent, sibling};
3use crate::merkle_tree::{BatchPath, MTLeaf, Path};
4use blake2::digest::{Digest, FixedOutput};
5use serde::{Deserialize, Serialize};
6use std::marker::PhantomData;
7
8/// `MerkleTree` commitment.
9/// This structure differs from `MerkleTree` in that it does not contain all elements, which are not always necessary.
10/// Instead, it only contains the root of the tree.
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct MerkleTreeCommitment<D: Digest> {
13    /// Root of the merkle commitment.
14    pub root: Vec<u8>,
15    hasher: PhantomData<D>,
16}
17
18impl<D: Digest + FixedOutput> MerkleTreeCommitment<D> {
19    pub(crate) fn new(root: Vec<u8>) -> Self {
20        MerkleTreeCommitment {
21            root,
22            hasher: PhantomData,
23        }
24    }
25
26    /// Check an inclusion proof that `val` is part of the tree by traveling the whole path until the root.
27    /// # Error
28    /// If the merkle tree path is invalid, then the function fails.
29    pub fn check(&self, val: &MTLeaf, proof: &Path<D>) -> Result<(), MerkleTreeError<D>>
30    where
31        D: FixedOutput + Clone,
32    {
33        let mut idx = proof.index;
34
35        let mut h = D::digest(val.to_bytes()).to_vec();
36        for p in &proof.values {
37            if (idx & 0b1) == 0 {
38                h = D::new().chain_update(h).chain_update(p).finalize().to_vec();
39            } else {
40                h = D::new().chain_update(p).chain_update(h).finalize().to_vec();
41            }
42            idx >>= 1;
43        }
44
45        if h == self.root {
46            return Ok(());
47        }
48        Err(MerkleTreeError::PathInvalid(proof.clone()))
49    }
50
51    /// Serializes the Merkle Tree commitment together with a message in a single vector of bytes.
52    /// Outputs `msg || self` as a vector of bytes.
53    pub fn concat_with_msg(&self, msg: &[u8]) -> Vec<u8>
54    where
55        D: Digest,
56    {
57        let mut msgp = msg.to_vec();
58        let mut bytes = self.root.clone();
59        msgp.append(&mut bytes);
60
61        msgp
62    }
63}
64
65/// Batch compatible `MerkleTree` commitment .
66/// This structure differs from `MerkleTreeCommitment` in that it stores the number of leaves in the tree
67/// as well as the root of the tree.
68/// Number of leaves is required by the batch path generation/verification.
69#[derive(Debug, Clone, Serialize, Deserialize)]
70pub struct MerkleTreeCommitmentBatchCompat<D: Digest> {
71    /// Root of the merkle commitment.
72    pub root: Vec<u8>,
73    nr_leaves: usize,
74    hasher: PhantomData<D>,
75}
76
77impl<D: Digest> MerkleTreeCommitmentBatchCompat<D> {
78    pub(crate) fn new(root: Vec<u8>, nr_leaves: usize) -> Self {
79        Self {
80            root,
81            nr_leaves,
82            hasher: Default::default(),
83        }
84    }
85
86    #[cfg(test)]
87    /// Used in property test of `tree`: `test_bytes_tree_commitment_batch_compat`
88    pub(crate) fn get_nr_leaves(&self) -> usize {
89        self.nr_leaves
90    }
91
92    /// Serializes the Merkle Tree commitment together with a message in a single vector of bytes.
93    /// Outputs `msg || self` as a vector of bytes.
94    // todo: Do we need to concat msg to whole commitment (nr_leaves and root) or just the root?
95    pub fn concat_with_msg(&self, msg: &[u8]) -> Vec<u8> {
96        let mut msgp = msg.to_vec();
97        let mut bytes = self.root.clone();
98        msgp.append(&mut bytes);
99
100        msgp
101    }
102
103    /// Check a proof of a batched opening. The indices must be ordered.
104    ///
105    /// # Error
106    /// Returns an error if the proof is invalid.
107    // todo: Update doc.
108    // todo: Simplify the algorithm.
109    // todo: Maybe we want more granular errors, rather than only `BatchPathInvalid`
110    pub fn check(
111        &self,
112        batch_val: &[MTLeaf],
113        proof: &BatchPath<D>,
114    ) -> Result<(), MerkleTreeError<D>>
115    where
116        D: FixedOutput + Clone,
117    {
118        if batch_val.len() != proof.indices.len() {
119            return Err(MerkleTreeError::BatchPathInvalid(proof.clone()));
120        }
121        let mut ordered_indices: Vec<usize> = proof.indices.clone();
122        ordered_indices.sort_unstable();
123
124        if ordered_indices != proof.indices {
125            return Err(MerkleTreeError::BatchPathInvalid(proof.clone()));
126        }
127
128        let nr_nodes = self.nr_leaves + self.nr_leaves.next_power_of_two() - 1;
129
130        ordered_indices = ordered_indices
131            .into_iter()
132            .map(|i| i + self.nr_leaves.next_power_of_two() - 1)
133            .collect();
134
135        let mut idx = ordered_indices[0];
136        // First we need to hash the leave values
137        let mut leaves: Vec<Vec<u8>> = batch_val
138            .iter()
139            .map(|val| D::digest(val.to_bytes()).to_vec())
140            .collect();
141
142        let mut values = proof.values.clone();
143
144        while idx > 0 {
145            let mut new_hashes = Vec::with_capacity(ordered_indices.len());
146            let mut new_indices = Vec::with_capacity(ordered_indices.len());
147            let mut i = 0;
148            idx = parent(idx);
149            while i < ordered_indices.len() {
150                new_indices.push(parent(ordered_indices[i]));
151                if ordered_indices[i] & 1 == 0 {
152                    new_hashes.push(
153                        D::new()
154                            .chain(values.first().ok_or(MerkleTreeError::SerializationError)?)
155                            .chain(&leaves[i])
156                            .finalize()
157                            .to_vec(),
158                    );
159                    values.remove(0);
160                } else {
161                    let sibling = sibling(ordered_indices[i]);
162                    if i < ordered_indices.len() - 1 && ordered_indices[i + 1] == sibling {
163                        new_hashes.push(
164                            D::new()
165                                .chain(&leaves[i])
166                                .chain(&leaves[i + 1])
167                                .finalize()
168                                .to_vec(),
169                        );
170                        i += 1;
171                    } else if sibling < nr_nodes {
172                        new_hashes.push(
173                            D::new()
174                                .chain(&leaves[i])
175                                .chain(values.first().ok_or(MerkleTreeError::SerializationError)?)
176                                .finalize()
177                                .to_vec(),
178                        );
179                        values.remove(0);
180                    } else {
181                        new_hashes.push(
182                            D::new()
183                                .chain(&leaves[i])
184                                .chain(D::digest([0u8]))
185                                .finalize()
186                                .to_vec(),
187                        );
188                    }
189                }
190                i += 1;
191            }
192            leaves.clone_from(&new_hashes);
193            ordered_indices.clone_from(&new_indices);
194        }
195
196        if leaves.len() == 1 && leaves[0] == self.root {
197            return Ok(());
198        }
199
200        Err(MerkleTreeError::BatchPathInvalid(proof.clone()))
201    }
202}
203
204impl<D: Digest> PartialEq for MerkleTreeCommitmentBatchCompat<D> {
205    fn eq(&self, other: &Self) -> bool {
206        self.root == other.root && self.nr_leaves == other.nr_leaves
207    }
208}
209
210impl<D: Digest> Eq for MerkleTreeCommitmentBatchCompat<D> {}
mithril_stm/merkle_tree/commitment.rs

mithril_stm/merkle_tree/
commitment.rs
