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