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