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 fn check(
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    /// Serializes the Merkle Tree commitment together with a message in a single vector of bytes.
57    /// Outputs `msg || self` as a vector of bytes.
58    pub fn concat_with_msg(&self, msg: &[u8]) -> Vec<u8>
59    where
60        D: Digest,
61    {
62        let mut msgp = msg.to_vec();
63        let mut bytes = self.root.clone();
64        msgp.append(&mut bytes);
65
66        msgp
67    }
68}
69
70/// Batch compatible `MerkleTree` commitment .
71/// This structure differs from `MerkleTreeCommitment` in that it stores the number of leaves in the tree
72/// as well as the root of the tree.
73/// Number of leaves is required by the batch path generation/verification.
74#[derive(Debug, Clone, Serialize, Deserialize)]
75pub struct MerkleTreeBatchCommitment<D: Digest> {
76    /// Root of the merkle commitment.
77    pub root: Vec<u8>,
78    nr_leaves: usize,
79    hasher: PhantomData<D>,
80}
81
82impl<D: Digest> MerkleTreeBatchCommitment<D> {
83    pub(crate) fn new(root: Vec<u8>, nr_leaves: usize) -> Self {
84        Self {
85            root,
86            nr_leaves,
87            hasher: Default::default(),
88        }
89    }
90
91    #[cfg(test)]
92    /// Used in property test of `tree`: `test_bytes_tree_commitment_batch_compat`
93    pub(crate) fn get_nr_leaves(&self) -> usize {
94        self.nr_leaves
95    }
96
97    /// Serializes the Merkle Tree commitment together with a message in a single vector of bytes.
98    /// Outputs `msg || self` as a vector of bytes.
99    // todo: Do we need to concat msg to whole commitment (nr_leaves and root) or just the root?
100    pub fn concat_with_msg(&self, msg: &[u8]) -> Vec<u8> {
101        let mut msgp = msg.to_vec();
102        let mut bytes = self.root.clone();
103        msgp.append(&mut bytes);
104
105        msgp
106    }
107
108    /// Check a proof of a batched opening. The indices must be ordered.
109    ///
110    /// # Error
111    /// Returns an error if the proof is invalid.
112    // todo: Update doc.
113    // todo: Simplify the algorithm.
114    // todo: Maybe we want more granular errors, rather than only `BatchPathInvalid`
115    pub fn check(
116        &self,
117        batch_val: &[MerkleTreeLeaf],
118        proof: &MerkleBatchPath<D>,
119    ) -> Result<(), MerkleTreeError<D>>
120    where
121        D: FixedOutput + Clone,
122    {
123        if batch_val.len() != proof.indices.len() {
124            return Err(MerkleTreeError::BatchPathInvalid(proof.clone()));
125        }
126        let mut ordered_indices: Vec<usize> = proof.indices.clone();
127        ordered_indices.sort_unstable();
128
129        if ordered_indices != proof.indices {
130            return Err(MerkleTreeError::BatchPathInvalid(proof.clone()));
131        }
132
133        let nr_nodes = self.nr_leaves + self.nr_leaves.next_power_of_two() - 1;
134
135        ordered_indices = ordered_indices
136            .into_iter()
137            .map(|i| i + self.nr_leaves.next_power_of_two() - 1)
138            .collect();
139
140        let mut idx = ordered_indices[0];
141        // First we need to hash the leave values
142        let mut leaves: Vec<Vec<u8>> = batch_val
143            .iter()
144            .map(|val| D::digest(val.to_bytes()).to_vec())
145            .collect();
146
147        let mut values = proof.values.clone();
148
149        while idx > 0 {
150            let mut new_hashes = Vec::with_capacity(ordered_indices.len());
151            let mut new_indices = Vec::with_capacity(ordered_indices.len());
152            let mut i = 0;
153            idx = parent(idx);
154            while i < ordered_indices.len() {
155                new_indices.push(parent(ordered_indices[i]));
156                if ordered_indices[i] & 1 == 0 {
157                    new_hashes.push(
158                        D::new()
159                            .chain(values.first().ok_or(MerkleTreeError::SerializationError)?)
160                            .chain(&leaves[i])
161                            .finalize()
162                            .to_vec(),
163                    );
164                    values.remove(0);
165                } else {
166                    let sibling = sibling(ordered_indices[i]);
167                    if i < ordered_indices.len() - 1 && ordered_indices[i + 1] == sibling {
168                        new_hashes.push(
169                            D::new().chain(&leaves[i]).chain(&leaves[i + 1]).finalize().to_vec(),
170                        );
171                        i += 1;
172                    } else if sibling < nr_nodes {
173                        new_hashes.push(
174                            D::new()
175                                .chain(&leaves[i])
176                                .chain(values.first().ok_or(MerkleTreeError::SerializationError)?)
177                                .finalize()
178                                .to_vec(),
179                        );
180                        values.remove(0);
181                    } else {
182                        new_hashes.push(
183                            D::new().chain(&leaves[i]).chain(D::digest([0u8])).finalize().to_vec(),
184                        );
185                    }
186                }
187                i += 1;
188            }
189            leaves.clone_from(&new_hashes);
190            ordered_indices.clone_from(&new_indices);
191        }
192
193        if leaves.len() == 1 && leaves[0] == self.root {
194            return Ok(());
195        }
196
197        Err(MerkleTreeError::BatchPathInvalid(proof.clone()))
198    }
199}
200
201impl<D: Digest> PartialEq for MerkleTreeBatchCommitment<D> {
202    fn eq(&self, other: &Self) -> bool {
203        self.root == other.root && self.nr_leaves == other.nr_leaves
204    }
205}
206
207impl<D: Digest> Eq for MerkleTreeBatchCommitment<D> {}
mithril_stm/merkle_tree/commitment.rs

mithril_stm/merkle_tree/
commitment.rs
