miden-crypto/src/merkle/merkle_tree.rs
src/merkle/merkle_tree.rsRaw
1use super::{
2 Felt, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, Word,
3};
4use crate::{
5 utils::{string::String, uninit_vector, word_to_hex},
6 FieldElement,
7};
8use core::{fmt, slice};
9use winter_math::log2;
10
11// MERKLE TREE
12// ================================================================================================
13
14/// A fully-balanced binary Merkle tree (i.e., a tree where the number of leaves is a power of two).
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct MerkleTree {
17 nodes: Vec<Word>,
18}
19
20impl MerkleTree {
21 // CONSTRUCTOR
22 // --------------------------------------------------------------------------------------------
23 /// Returns a Merkle tree instantiated from the provided leaves.
24 ///
25 /// # Errors
26 /// Returns an error if the number of leaves is smaller than two or is not a power of two.
27 pub fn new(leaves: Vec<Word>) -> Result<Self, MerkleError> {
28 let n = leaves.len();
29 if n <= 1 {
30 return Err(MerkleError::DepthTooSmall(n as u8));
31 } else if !n.is_power_of_two() {
32 return Err(MerkleError::NumLeavesNotPowerOfTwo(n));
33 }
34
35 // create un-initialized vector to hold all tree nodes
36 let mut nodes = unsafe { uninit_vector(2 * n) };
37 nodes[0] = [Felt::ZERO; 4];
38
39 // copy leaves into the second part of the nodes vector
40 nodes[n..].copy_from_slice(&leaves);
41
42 // re-interpret nodes as an array of two nodes fused together
43 // Safety: `nodes` will never move here as it is not bound to an external lifetime (i.e.
44 // `self`).
45 let ptr = nodes.as_ptr() as *const [RpoDigest; 2];
46 let pairs = unsafe { slice::from_raw_parts(ptr, n) };
47
48 // calculate all internal tree nodes
49 for i in (1..n).rev() {
50 nodes[i] = Rpo256::merge(&pairs[i]).into();
51 }
52
53 Ok(Self { nodes })
54 }
55
56 // PUBLIC ACCESSORS
57 // --------------------------------------------------------------------------------------------
58
59 /// Returns the root of this Merkle tree.
60 pub fn root(&self) -> Word {
61 self.nodes[1]
62 }
63
64 /// Returns the depth of this Merkle tree.
65 ///
66 /// Merkle tree of depth 1 has two leaves, depth 2 has four leaves etc.
67 pub fn depth(&self) -> u8 {
68 log2(self.nodes.len() / 2) as u8
69 }
70
71 /// Returns a node at the specified depth and index value.
72 ///
73 /// # Errors
74 /// Returns an error if:
75 /// * The specified depth is greater than the depth of the tree.
76 /// * The specified index is not valid for the specified depth.
77 pub fn get_node(&self, index: NodeIndex) -> Result<Word, MerkleError> {
78 if index.is_root() {
79 return Err(MerkleError::DepthTooSmall(index.depth()));
80 } else if index.depth() > self.depth() {
81 return Err(MerkleError::DepthTooBig(index.depth() as u64));
82 } else if !index.is_valid() {
83 return Err(MerkleError::InvalidIndex(index));
84 }
85
86 let pos = index.to_scalar_index() as usize;
87 Ok(self.nodes[pos])
88 }
89
90 /// Returns a Merkle path to the node at the specified depth and index value. The node itself
91 /// is not included in the path.
92 ///
93 /// # Errors
94 /// Returns an error if:
95 /// * The specified depth is greater than the depth of the tree.
96 /// * The specified value is not valid for the specified depth.
97 pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
98 if index.is_root() {
99 return Err(MerkleError::DepthTooSmall(index.depth()));
100 } else if index.depth() > self.depth() {
101 return Err(MerkleError::DepthTooBig(index.depth() as u64));
102 } else if !index.is_valid() {
103 return Err(MerkleError::InvalidIndex(index));
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!(
117 index.is_root(),
118 "the path walk must go all the way to the root"
119 );
120
121 Ok(path.into())
122 }
123
124 /// Replaces the leaf at the specified index with the provided value.
125 ///
126 /// # Errors
127 /// Returns an error if the specified index value is not a valid leaf value for this tree.
128 pub fn update_leaf<'a>(&'a mut self, index_value: u64, value: Word) -> Result<(), MerkleError> {
129 let depth = self.depth();
130 let mut index = NodeIndex::new(depth, index_value);
131 if !index.is_valid() {
132 return Err(MerkleError::InvalidIndex(index));
133 }
134
135 // we don't need to copy the pairs into a new address as we are logically guaranteed to not
136 // overlap write instructions. however, it's important to bind the lifetime of pairs to
137 // `self.nodes` so the compiler will never move one without moving the other.
138 debug_assert_eq!(self.nodes.len() & 1, 0);
139 let n = self.nodes.len() / 2;
140
141 // Safety: the length of nodes is guaranteed to contain pairs of words; hence, pairs of
142 // digests. we explicitly bind the lifetime here so we add an extra layer of guarantee that
143 // `self.nodes` will be moved only if `pairs` is moved as well. also, the algorithm is
144 // logically guaranteed to not overlap write positions as the write index is always half
145 // the index from which we read the digest input.
146 let ptr = self.nodes.as_ptr() as *const [RpoDigest; 2];
147 let pairs: &'a [[RpoDigest; 2]] = unsafe { slice::from_raw_parts(ptr, n) };
148
149 // update the current node
150 let pos = index.to_scalar_index() as usize;
151 self.nodes[pos] = value;
152
153 // traverse to the root, updating each node with the merged values of its parents
154 for _ in 0..index.depth() {
155 index.move_up();
156 let pos = index.to_scalar_index() as usize;
157 let value = Rpo256::merge(&pairs[pos]).into();
158 self.nodes[pos] = value;
159 }
160
161 Ok(())
162 }
163
164 /// An iterator over every inner node in the tree. The iterator order is unspecified.
165 pub fn inner_nodes(&self) -> MerkleTreeNodes<'_> {
166 MerkleTreeNodes {
167 nodes: &self.nodes,
168 index: 1, // index 0 is just padding, start at 1
169 }
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 MerkleTreeNodes<'a> {
180 nodes: &'a Vec<Word>,
181 index: usize,
182}
183
184impl<'a> Iterator for MerkleTreeNodes<'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 to vizualize a [MerkleTree] in text.
207pub fn tree_to_text(tree: &MerkleTree) -> Result<String, fmt::Error> {
208 let indent = " ";
209 let mut s = String::new();
210 s.push_str(&word_to_hex(&tree.root())?);
211 s.push('\n');
212 for d in 1..=tree.depth() {
213 let entries = 2u64.pow(d.into());
214 for i in 0..entries {
215 let index = NodeIndex::new(d, i);
216
217 let node = tree
218 .get_node(index)
219 .expect("The index must always be valid");
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 vizualize 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::merkle::{int_to_node, InnerNodeInfo};
259 use core::mem::size_of;
260 use proptest::prelude::*;
261
262 const LEAVES4: [Word; 4] = [
263 int_to_node(1),
264 int_to_node(2),
265 int_to_node(3),
266 int_to_node(4),
267 ];
268
269 const LEAVES8: [Word; 8] = [
270 int_to_node(1),
271 int_to_node(2),
272 int_to_node(3),
273 int_to_node(4),
274 int_to_node(5),
275 int_to_node(6),
276 int_to_node(7),
277 int_to_node(8),
278 ];
279
280 #[test]
281 fn build_merkle_tree() {
282 let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap();
283 assert_eq!(8, tree.nodes.len());
284
285 // leaves were copied correctly
286 for (a, b) in tree.nodes.iter().skip(4).zip(LEAVES4.iter()) {
287 assert_eq!(a, b);
288 }
289
290 let (root, node2, node3) = compute_internal_nodes();
291
292 assert_eq!(root, tree.nodes[1]);
293 assert_eq!(node2, tree.nodes[2]);
294 assert_eq!(node3, tree.nodes[3]);
295
296 assert_eq!(root, tree.root());
297 }
298
299 #[test]
300 fn get_leaf() {
301 let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap();
302
303 // check depth 2
304 assert_eq!(LEAVES4[0], tree.get_node(NodeIndex::new(2, 0)).unwrap());
305 assert_eq!(LEAVES4[1], tree.get_node(NodeIndex::new(2, 1)).unwrap());
306 assert_eq!(LEAVES4[2], tree.get_node(NodeIndex::new(2, 2)).unwrap());
307 assert_eq!(LEAVES4[3], tree.get_node(NodeIndex::new(2, 3)).unwrap());
308
309 // check depth 1
310 let (_, node2, node3) = compute_internal_nodes();
311
312 assert_eq!(node2, tree.get_node(NodeIndex::new(1, 0)).unwrap());
313 assert_eq!(node3, tree.get_node(NodeIndex::new(1, 1)).unwrap());
314 }
315
316 #[test]
317 fn get_path() {
318 let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap();
319
320 let (_, node2, node3) = compute_internal_nodes();
321
322 // check depth 2
323 assert_eq!(
324 vec![LEAVES4[1], node3],
325 *tree.get_path(NodeIndex::new(2, 0)).unwrap()
326 );
327 assert_eq!(
328 vec![LEAVES4[0], node3],
329 *tree.get_path(NodeIndex::new(2, 1)).unwrap()
330 );
331 assert_eq!(
332 vec![LEAVES4[3], node2],
333 *tree.get_path(NodeIndex::new(2, 2)).unwrap()
334 );
335 assert_eq!(
336 vec![LEAVES4[2], node2],
337 *tree.get_path(NodeIndex::new(2, 3)).unwrap()
338 );
339
340 // check depth 1
341 assert_eq!(vec![node3], *tree.get_path(NodeIndex::new(1, 0)).unwrap());
342 assert_eq!(vec![node2], *tree.get_path(NodeIndex::new(1, 1)).unwrap());
343 }
344
345 #[test]
346 fn update_leaf() {
347 let mut tree = super::MerkleTree::new(LEAVES8.to_vec()).unwrap();
348
349 // update one leaf
350 let value = 3;
351 let new_node = int_to_node(9);
352 let mut expected_leaves = LEAVES8.to_vec();
353 expected_leaves[value as usize] = new_node;
354 let expected_tree = super::MerkleTree::new(expected_leaves.clone()).unwrap();
355
356 tree.update_leaf(value, new_node).unwrap();
357 assert_eq!(expected_tree.nodes, tree.nodes);
358
359 // update another leaf
360 let value = 6;
361 let new_node = int_to_node(10);
362 expected_leaves[value as usize] = new_node;
363 let expected_tree = super::MerkleTree::new(expected_leaves.clone()).unwrap();
364
365 tree.update_leaf(value, new_node).unwrap();
366 assert_eq!(expected_tree.nodes, tree.nodes);
367 }
368
369 #[test]
370 fn nodes() -> Result<(), MerkleError> {
371 let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap();
372 let root = tree.root();
373 let l1n0 = tree.get_node(NodeIndex::new(1, 0))?;
374 let l1n1 = tree.get_node(NodeIndex::new(1, 1))?;
375 let l2n0 = tree.get_node(NodeIndex::new(2, 0))?;
376 let l2n1 = tree.get_node(NodeIndex::new(2, 1))?;
377 let l2n2 = tree.get_node(NodeIndex::new(2, 2))?;
378 let l2n3 = tree.get_node(NodeIndex::new(2, 3))?;
379
380 let nodes: Vec<InnerNodeInfo> = tree.inner_nodes().collect();
381 let expected = vec![
382 InnerNodeInfo {
383 value: root,
384 left: l1n0,
385 right: l1n1,
386 },
387 InnerNodeInfo {
388 value: l1n0,
389 left: l2n0,
390 right: l2n1,
391 },
392 InnerNodeInfo {
393 value: l1n1,
394 left: l2n2,
395 right: l2n3,
396 },
397 ];
398 assert_eq!(nodes, expected);
399
400 Ok(())
401 }
402
403 proptest! {
404 #[test]
405 fn arbitrary_word_can_be_represented_as_digest(
406 a in prop::num::u64::ANY,
407 b in prop::num::u64::ANY,
408 c in prop::num::u64::ANY,
409 d in prop::num::u64::ANY,
410 ) {
411 // this test will assert the memory equivalence between word and digest.
412 // it is used to safeguard the `[MerkleTee::update_leaf]` implementation
413 // that assumes this equivalence.
414
415 // build a word and copy it to another address as digest
416 let word = [Felt::new(a), Felt::new(b), Felt::new(c), Felt::new(d)];
417 let digest = RpoDigest::from(word);
418
419 // assert the addresses are different
420 let word_ptr = (&word).as_ptr() as *const u8;
421 let digest_ptr = (&digest).as_ptr() as *const u8;
422 assert_ne!(word_ptr, digest_ptr);
423
424 // compare the bytes representation
425 let word_bytes = unsafe { slice::from_raw_parts(word_ptr, size_of::<Word>()) };
426 let digest_bytes = unsafe { slice::from_raw_parts(digest_ptr, size_of::<RpoDigest>()) };
427 assert_eq!(word_bytes, digest_bytes);
428 }
429 }
430
431 // HELPER FUNCTIONS
432 // --------------------------------------------------------------------------------------------
433
434 fn compute_internal_nodes() -> (Word, Word, Word) {
435 let node2 = Rpo256::hash_elements(&[LEAVES4[0], LEAVES4[1]].concat());
436 let node3 = Rpo256::hash_elements(&[LEAVES4[2], LEAVES4[3]].concat());
437 let root = Rpo256::merge(&[node2, node3]);
438
439 (root.into(), node2.into(), node3.into())
440 }
441}