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