miden_crypto/merkle/
merkle_tree.rs

1use alloc::{string::String, vec::Vec};
2use core::{fmt, slice};
3
4use super::{InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Poseidon2, Word};
5use crate::utils::{assume_init_vec, uninit_vector, word_to_hex};
6
7// MERKLE TREE
8// ================================================================================================
9
10/// A fully-balanced binary Merkle tree (i.e., a tree where the number of leaves is a power of two).
11#[derive(Debug, Clone, PartialEq, Eq)]
12#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
13pub struct MerkleTree {
14    nodes: Vec<Word>,
15}
16
17impl MerkleTree {
18    // CONSTRUCTOR
19    // --------------------------------------------------------------------------------------------
20    /// Returns a Merkle tree instantiated from the provided leaves.
21    ///
22    /// # Errors
23    /// Returns an error if the number of leaves is smaller than two or is not a power of two.
24    pub fn new<T>(leaves: T) -> Result<Self, MerkleError>
25    where
26        T: AsRef<[Word]>,
27    {
28        let leaves = leaves.as_ref();
29        let n = leaves.len();
30        if n <= 1 {
31            return Err(MerkleError::DepthTooSmall(n as u8));
32        } else if !n.is_power_of_two() {
33            return Err(MerkleError::NumLeavesNotPowerOfTwo(n));
34        }
35
36        // Create an uninitialized vector to avoid eagerly zeroing `2 * n` words. This is a
37        // hot path during tree construction; see benches in `miden-crypto/benches/merkle.rs`
38        // (e.g. `merkle_tree_construction`) for performance motivation.
39        // SAFETY: All elements are written before being read (leaves copied, then computed).
40        let mut nodes = uninit_vector::<Word>(2 * n);
41        nodes[0].write(Word::default());
42
43        // copy leaves into the second part of the nodes vector
44        nodes[n..].iter_mut().zip(leaves).for_each(|(node, leaf)| {
45            node.write(*leaf);
46        });
47
48        // calculate all internal tree nodes
49        for i in (1..n).rev() {
50            // SAFETY: We fill leaves first, then iterate from the bottom up. At this point,
51            // nodes[2 * i] and nodes[2 * i + 1] have already been written.
52            let left = unsafe { nodes[2 * i].assume_init_read() };
53            let right = unsafe { nodes[2 * i + 1].assume_init_read() };
54            nodes[i].write(Poseidon2::merge(&[left, right]));
55        }
56
57        // SAFETY: all elements were written above.
58        let nodes = unsafe { assume_init_vec(nodes) };
59
60        Ok(Self { nodes })
61    }
62
63    // PUBLIC ACCESSORS
64    // --------------------------------------------------------------------------------------------
65
66    /// Returns the root of this Merkle tree.
67    pub fn root(&self) -> Word {
68        self.nodes[1]
69    }
70
71    /// Returns the depth of this Merkle tree.
72    ///
73    /// Merkle tree of depth 1 has two leaves, depth 2 has four leaves etc.
74    pub fn depth(&self) -> u8 {
75        (self.nodes.len() / 2).ilog2() as u8
76    }
77
78    /// Returns a node at the specified depth and index value.
79    ///
80    /// # Errors
81    /// Returns an error if:
82    /// * The specified depth is greater than the depth of the tree.
83    /// * The specified index is not valid for the specified depth.
84    pub fn get_node(&self, index: NodeIndex) -> Result<Word, MerkleError> {
85        if index.is_root() {
86            return Err(MerkleError::DepthTooSmall(index.depth()));
87        } else if index.depth() > self.depth() {
88            return Err(MerkleError::DepthTooBig(index.depth() as u64));
89        }
90
91        let pos = index.to_scalar_index()? as usize;
92        Ok(self.nodes[pos])
93    }
94
95    /// Returns a Merkle path to the node at the specified depth and index value. The node itself
96    /// is not included in the path.
97    ///
98    /// # Errors
99    /// Returns an error if:
100    /// * The specified depth is greater than the depth of the tree.
101    /// * The specified value is not valid for the specified depth.
102    pub fn get_path(&self, index: NodeIndex) -> Result<MerklePath, MerkleError> {
103        if index.is_root() {
104            return Err(MerkleError::DepthTooSmall(index.depth()));
105        } else if index.depth() > self.depth() {
106            return Err(MerkleError::DepthTooBig(index.depth() as u64));
107        }
108
109        Ok(MerklePath::from(Vec::from_iter(
110            index.proof_indices().map(|index| self.get_node(index).unwrap()),
111        )))
112    }
113
114    // ITERATORS
115    // --------------------------------------------------------------------------------------------
116
117    /// Returns an iterator over the leaves of this [MerkleTree].
118    pub fn leaves(&self) -> impl Iterator<Item = (u64, &Word)> {
119        let leaves_start = self.nodes.len() / 2;
120        self.nodes.iter().skip(leaves_start).enumerate().map(|(i, v)| (i as u64, v))
121    }
122
123    /// Returns n iterator over every inner node of this [MerkleTree].
124    ///
125    /// The iterator order is unspecified.
126    pub fn inner_nodes(&self) -> InnerNodeIterator<'_> {
127        InnerNodeIterator {
128            nodes: &self.nodes,
129            index: 1, // index 0 is just padding, start at 1
130        }
131    }
132
133    // STATE MUTATORS
134    // --------------------------------------------------------------------------------------------
135
136    /// Replaces the leaf at the specified index with the provided value.
137    ///
138    /// # Errors
139    /// Returns an error if the specified index value is not a valid leaf value for this tree.
140    pub fn update_leaf<'a>(&'a mut self, index_value: u64, value: Word) -> Result<(), MerkleError> {
141        let mut index = NodeIndex::new(self.depth(), index_value)?;
142
143        // Performance note: We use unsafe pointer casts here for ~2-2.5% performance improvement
144        // at scale. See benches/merkle.rs for benchmarks. The safe alternative uses index
145        // arithmetic (`nodes[pos*2]`, `nodes[pos*2+1]`) which is measurably slower on large
146        // trees.
147
148        // we don't need to copy the pairs into a new address as we are logically guaranteed to not
149        // overlap write instructions. however, it's important to bind the lifetime of pairs to
150        // `self.nodes` so the compiler will never move one without moving the other.
151        debug_assert_eq!(self.nodes.len() & 1, 0);
152        let n = self.nodes.len() / 2;
153
154        // Safety: the length of nodes is guaranteed to contain pairs of words; hence, pairs of
155        // digests. we explicitly bind the lifetime here so we add an extra layer of guarantee that
156        // `self.nodes` will be moved only if `pairs` is moved as well. also, the algorithm is
157        // logically guaranteed to not overlap write positions as the write index is always half
158        // the index from which we read the digest input.
159        let ptr = self.nodes.as_ptr() as *const [Word; 2];
160        let pairs: &'a [[Word; 2]] = unsafe { slice::from_raw_parts(ptr, n) };
161
162        // update the current node
163        let pos = index.to_scalar_index()? as usize;
164        self.nodes[pos] = value;
165
166        // traverse to the root, updating each node with the merged values of its parents
167        for _ in 0..index.depth() {
168            index.move_up();
169            let pos = index.to_scalar_index()? as usize;
170            let value = Poseidon2::merge(&pairs[pos]);
171            self.nodes[pos] = value;
172        }
173
174        Ok(())
175    }
176}
177
178// CONVERSIONS
179// ================================================================================================
180
181impl TryFrom<&[Word]> for MerkleTree {
182    type Error = MerkleError;
183
184    fn try_from(value: &[Word]) -> Result<Self, Self::Error> {
185        MerkleTree::new(value)
186    }
187}
188
189// ITERATORS
190// ================================================================================================
191
192/// An iterator over every inner node of the [MerkleTree].
193///
194/// Use this to extract the data of the tree, there is no guarantee on the order of the elements.
195pub struct InnerNodeIterator<'a> {
196    nodes: &'a Vec<Word>,
197    index: usize,
198}
199
200impl Iterator for InnerNodeIterator<'_> {
201    type Item = InnerNodeInfo;
202
203    fn next(&mut self) -> Option<Self::Item> {
204        if self.index < self.nodes.len() / 2 {
205            let value = self.index;
206            let left = self.index * 2;
207            let right = left + 1;
208
209            self.index += 1;
210
211            Some(InnerNodeInfo {
212                value: self.nodes[value],
213                left: self.nodes[left],
214                right: self.nodes[right],
215            })
216        } else {
217            None
218        }
219    }
220}
221
222// UTILITY FUNCTIONS
223// ================================================================================================
224
225/// Utility to visualize a [MerkleTree] in text.
226pub fn tree_to_text(tree: &MerkleTree) -> Result<String, fmt::Error> {
227    let indent = "  ";
228    let mut s = String::new();
229    s.push_str(&word_to_hex(&tree.root())?);
230    s.push('\n');
231    for d in 1..=tree.depth() {
232        let entries = 2u64.pow(d.into());
233        for i in 0..entries {
234            let index = NodeIndex::new(d, i).expect("The index must always be valid");
235            let node = tree.get_node(index).expect("The node must always be found");
236
237            for _ in 0..d {
238                s.push_str(indent);
239            }
240            s.push_str(&word_to_hex(&node)?);
241            s.push('\n');
242        }
243    }
244
245    Ok(s)
246}
247
248/// Utility to visualize a [MerklePath] in text.
249pub fn path_to_text(path: &MerklePath) -> Result<String, fmt::Error> {
250    let mut s = String::new();
251    s.push('[');
252
253    for el in path.iter() {
254        s.push_str(&word_to_hex(el)?);
255        s.push_str(", ");
256    }
257
258    // remove the last ", "
259    if !path.is_empty() {
260        s.pop();
261        s.pop();
262    }
263    s.push(']');
264
265    Ok(s)
266}
267
268// TESTS
269// ================================================================================================
270
271#[cfg(test)]
272mod tests {
273    use core::mem::size_of;
274
275    use proptest::prelude::*;
276
277    use super::*;
278    use crate::{
279        Felt,
280        merkle::{int_to_leaf, int_to_node},
281    };
282
283    const LEAVES4: [Word; Word::NUM_ELEMENTS] =
284        [int_to_node(1), int_to_node(2), int_to_node(3), int_to_node(4)];
285
286    const LEAVES8: [Word; 8] = [
287        int_to_node(1),
288        int_to_node(2),
289        int_to_node(3),
290        int_to_node(4),
291        int_to_node(5),
292        int_to_node(6),
293        int_to_node(7),
294        int_to_node(8),
295    ];
296
297    #[test]
298    fn build_merkle_tree() {
299        let tree = MerkleTree::new(LEAVES4).unwrap();
300        assert_eq!(8, tree.nodes.len());
301
302        // leaves were copied correctly
303        for (a, b) in tree.nodes.iter().skip(4).zip(LEAVES4.iter()) {
304            assert_eq!(a, b);
305        }
306
307        let (root, node2, node3) = compute_internal_nodes();
308
309        assert_eq!(root, tree.nodes[1]);
310        assert_eq!(node2, tree.nodes[2]);
311        assert_eq!(node3, tree.nodes[3]);
312
313        assert_eq!(root, tree.root());
314    }
315
316    #[test]
317    fn get_leaf() {
318        let tree = MerkleTree::new(LEAVES4).unwrap();
319
320        // check depth 2
321        assert_eq!(LEAVES4[0], tree.get_node(NodeIndex::make(2, 0)).unwrap());
322        assert_eq!(LEAVES4[1], tree.get_node(NodeIndex::make(2, 1)).unwrap());
323        assert_eq!(LEAVES4[2], tree.get_node(NodeIndex::make(2, 2)).unwrap());
324        assert_eq!(LEAVES4[3], tree.get_node(NodeIndex::make(2, 3)).unwrap());
325
326        // check depth 1
327        let (_, node2, node3) = compute_internal_nodes();
328
329        assert_eq!(node2, tree.get_node(NodeIndex::make(1, 0)).unwrap());
330        assert_eq!(node3, tree.get_node(NodeIndex::make(1, 1)).unwrap());
331    }
332
333    #[test]
334    fn get_path() {
335        let tree = MerkleTree::new(LEAVES4).unwrap();
336
337        let (_, node2, node3) = compute_internal_nodes();
338
339        // check depth 2
340        assert_eq!(vec![LEAVES4[1], node3], *tree.get_path(NodeIndex::make(2, 0)).unwrap());
341        assert_eq!(vec![LEAVES4[0], node3], *tree.get_path(NodeIndex::make(2, 1)).unwrap());
342        assert_eq!(vec![LEAVES4[3], node2], *tree.get_path(NodeIndex::make(2, 2)).unwrap());
343        assert_eq!(vec![LEAVES4[2], node2], *tree.get_path(NodeIndex::make(2, 3)).unwrap());
344
345        // check depth 1
346        assert_eq!(vec![node3], *tree.get_path(NodeIndex::make(1, 0)).unwrap());
347        assert_eq!(vec![node2], *tree.get_path(NodeIndex::make(1, 1)).unwrap());
348    }
349
350    #[test]
351    fn update_leaf() {
352        let mut tree = MerkleTree::new(LEAVES8).unwrap();
353
354        // update one leaf
355        let value = 3;
356        let new_node = int_to_leaf(9);
357        let mut expected_leaves = LEAVES8.to_vec();
358        expected_leaves[value as usize] = new_node;
359        let expected_tree = MerkleTree::new(expected_leaves.clone()).unwrap();
360
361        tree.update_leaf(value, new_node).unwrap();
362        assert_eq!(expected_tree.nodes, tree.nodes);
363
364        // update another leaf
365        let value = 6;
366        let new_node = int_to_leaf(10);
367        expected_leaves[value as usize] = new_node;
368        let expected_tree = MerkleTree::new(expected_leaves.clone()).unwrap();
369
370        tree.update_leaf(value, new_node).unwrap();
371        assert_eq!(expected_tree.nodes, tree.nodes);
372    }
373
374    #[test]
375    fn nodes() -> Result<(), MerkleError> {
376        let tree = MerkleTree::new(LEAVES4).unwrap();
377        let root = tree.root();
378        let l1n0 = tree.get_node(NodeIndex::make(1, 0))?;
379        let l1n1 = tree.get_node(NodeIndex::make(1, 1))?;
380        let l2n0 = tree.get_node(NodeIndex::make(2, 0))?;
381        let l2n1 = tree.get_node(NodeIndex::make(2, 1))?;
382        let l2n2 = tree.get_node(NodeIndex::make(2, 2))?;
383        let l2n3 = tree.get_node(NodeIndex::make(2, 3))?;
384
385        let nodes: Vec<InnerNodeInfo> = tree.inner_nodes().collect();
386        let expected = vec![
387            InnerNodeInfo { value: root, left: l1n0, right: l1n1 },
388            InnerNodeInfo { value: l1n0, left: l2n0, right: l2n1 },
389            InnerNodeInfo { value: l1n1, left: l2n2, right: l2n3 },
390        ];
391        assert_eq!(nodes, expected);
392
393        Ok(())
394    }
395
396    proptest! {
397        #[test]
398        fn arbitrary_word_can_be_represented_as_digest(
399            a in prop::num::u64::ANY,
400            b in prop::num::u64::ANY,
401            c in prop::num::u64::ANY,
402            d in prop::num::u64::ANY,
403        ) {
404            // this test will assert the memory equivalence between word and digest.
405            // it is used to safeguard the `[MerkleTee::update_leaf]` implementation
406            // that assumes this equivalence.
407
408            // build a word and copy it to another address as digest
409            let word = [Felt::new_unchecked(a), Felt::new_unchecked(b), Felt::new_unchecked(c), Felt::new_unchecked(d)];
410            let digest = Word::from(word);
411
412            // assert the addresses are different
413            let word_ptr = word.as_ptr() as *const u8;
414            let digest_ptr = digest.as_ptr() as *const u8;
415            assert_ne!(word_ptr, digest_ptr);
416
417            // compare the bytes representation
418            let word_bytes = unsafe { slice::from_raw_parts(word_ptr, size_of::<Word>()) };
419            let digest_bytes = unsafe { slice::from_raw_parts(digest_ptr, size_of::<Word>()) };
420            assert_eq!(word_bytes, digest_bytes);
421        }
422    }
423
424    // HELPER FUNCTIONS
425    // --------------------------------------------------------------------------------------------
426
427    fn compute_internal_nodes() -> (Word, Word, Word) {
428        let node2 = Poseidon2::hash_elements(&[*LEAVES4[0], *LEAVES4[1]].concat());
429        let node3 = Poseidon2::hash_elements(&[*LEAVES4[2], *LEAVES4[3]].concat());
430        let root = Poseidon2::merge(&[node2, node3]);
431
432        (root, node2, node3)
433    }
434}
miden_crypto/merkle/merkle_tree.rs

miden_crypto/merkle/
merkle_tree.rs
