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