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