miden-crypto/src/merkle/merkle_tree.rs
src/merkle/merkle_tree.rsRaw
1use super::{Felt, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, Word};
2use crate::{
3 utils::{string::String, uninit_vector, word_to_hex},
4 FieldElement,
5};
6use core::{fmt, slice};
7use winter_math::log2;
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)]
14pub struct MerkleTree {
15 pub(crate) nodes: Vec<Word>,
16}
17
18impl MerkleTree {
19 // CONSTRUCTOR
20 // --------------------------------------------------------------------------------------------
21 /// Returns a Merkle tree instantiated from the provided leaves.
22 ///
23 /// # Errors
24 /// Returns an error if the number of leaves is smaller than two or is not a power of two.
25 pub fn new(leaves: Vec<Word>) -> Result<Self, MerkleError> {
26 let n = leaves.len();
27 if n <= 1 {
28 return Err(MerkleError::DepthTooSmall(n as u8));
29 } else if !n.is_power_of_two() {
30 return Err(MerkleError::NumLeavesNotPowerOfTwo(n));
31 }
32
33 // create un-initialized vector to hold all tree nodes
34 let mut nodes = unsafe { uninit_vector(2 * n) };
35 nodes[0] = [Felt::ZERO; 4];
36
37 // copy leaves into the second part of the nodes vector
38 nodes[n..].copy_from_slice(&leaves);
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]).into();
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) -> Word {
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<Word, 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 } else if !index.is_valid() {
81 return Err(MerkleError::InvalidIndex(index));
82 }
83
84 let pos = index.to_scalar_index() as usize;
85 Ok(self.nodes[pos])
86 }
87
88 /// Returns a Merkle path to the node at the specified depth and index value. The node itself
89 /// is not included in the path.
90 ///
91 /// # Errors
92 /// Returns an error if:
93 /// * The specified depth is greater than the depth of the tree.
94 /// * The specified value is not valid for the specified depth.
95 pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
96 if index.is_root() {
97 return Err(MerkleError::DepthTooSmall(index.depth()));
98 } else if index.depth() > self.depth() {
99 return Err(MerkleError::DepthTooBig(index.depth() as u64));
100 } else if !index.is_valid() {
101 return Err(MerkleError::InvalidIndex(index));
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!(
115 index.is_root(),
116 "the path walk must go all the way to the root"
117 );
118
119 Ok(path.into())
120 }
121
122 /// Replaces the leaf at the specified index with the provided value.
123 ///
124 /// # Errors
125 /// Returns an error if the specified index value is not a valid leaf value for this tree.
126 pub fn update_leaf<'a>(&'a mut self, index_value: u64, value: Word) -> Result<(), MerkleError> {
127 let depth = self.depth();
128 let mut index = NodeIndex::new(depth, index_value);
129 if !index.is_valid() {
130 return Err(MerkleError::InvalidIndex(index));
131 }
132
133 // we don't need to copy the pairs into a new address as we are logically guaranteed to not
134 // overlap write instructions. however, it's important to bind the lifetime of pairs to
135 // `self.nodes` so the compiler will never move one without moving the other.
136 debug_assert_eq!(self.nodes.len() & 1, 0);
137 let n = self.nodes.len() / 2;
138
139 // Safety: the length of nodes is guaranteed to contain pairs of words; hence, pairs of
140 // digests. we explicitly bind the lifetime here so we add an extra layer of guarantee that
141 // `self.nodes` will be moved only if `pairs` is moved as well. also, the algorithm is
142 // logically guaranteed to not overlap write positions as the write index is always half
143 // the index from which we read the digest input.
144 let ptr = self.nodes.as_ptr() as *const [RpoDigest; 2];
145 let pairs: &'a [[RpoDigest; 2]] = unsafe { slice::from_raw_parts(ptr, n) };
146
147 // update the current node
148 let pos = index.to_scalar_index() as usize;
149 self.nodes[pos] = value;
150
151 // traverse to the root, updating each node with the merged values of its parents
152 for _ in 0..index.depth() {
153 index.move_up();
154 let pos = index.to_scalar_index() as usize;
155 let value = Rpo256::merge(&pairs[pos]).into();
156 self.nodes[pos] = value;
157 }
158
159 Ok(())
160 }
161}
162
163/// Utility to vizualize a [MerkleTree] in text.
164pub fn tree_to_text(tree: &MerkleTree) -> Result<String, fmt::Error> {
165 let indent = " ";
166 let mut s = String::new();
167 s.push_str(&word_to_hex(&tree.root())?);
168 s.push('\n');
169 for d in 1..=tree.depth() {
170 let entries = 2u64.pow(d.into());
171 for i in 0..entries {
172 let index = NodeIndex::new(d, i);
173
174 let node = tree
175 .get_node(index)
176 .expect("The index must always be valid");
177
178 for _ in 0..d {
179 s.push_str(indent);
180 }
181 s.push_str(&word_to_hex(&node)?);
182 s.push('\n');
183 }
184 }
185
186 Ok(s)
187}
188
189/// Utility to vizualize a [MerklePath] in text.
190pub fn path_to_text(path: &MerklePath) -> Result<String, fmt::Error> {
191 let mut s = String::new();
192 s.push('[');
193
194 for el in path.iter() {
195 s.push_str(&word_to_hex(el)?);
196 s.push_str(", ");
197 }
198
199 // remove the last ", "
200 if path.len() != 0 {
201 s.pop();
202 s.pop();
203 }
204 s.push(']');
205
206 Ok(s)
207}
208
209// TESTS
210// ================================================================================================
211
212#[cfg(test)]
213mod tests {
214 use super::*;
215 use crate::merkle::int_to_node;
216 use core::mem::size_of;
217 use proptest::prelude::*;
218
219 const LEAVES4: [Word; 4] = [
220 int_to_node(1),
221 int_to_node(2),
222 int_to_node(3),
223 int_to_node(4),
224 ];
225
226 const LEAVES8: [Word; 8] = [
227 int_to_node(1),
228 int_to_node(2),
229 int_to_node(3),
230 int_to_node(4),
231 int_to_node(5),
232 int_to_node(6),
233 int_to_node(7),
234 int_to_node(8),
235 ];
236
237 #[test]
238 fn build_merkle_tree() {
239 let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap();
240 assert_eq!(8, tree.nodes.len());
241
242 // leaves were copied correctly
243 for (a, b) in tree.nodes.iter().skip(4).zip(LEAVES4.iter()) {
244 assert_eq!(a, b);
245 }
246
247 let (root, node2, node3) = compute_internal_nodes();
248
249 assert_eq!(root, tree.nodes[1]);
250 assert_eq!(node2, tree.nodes[2]);
251 assert_eq!(node3, tree.nodes[3]);
252
253 assert_eq!(root, tree.root());
254 }
255
256 #[test]
257 fn get_leaf() {
258 let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap();
259
260 // check depth 2
261 assert_eq!(LEAVES4[0], tree.get_node(NodeIndex::new(2, 0)).unwrap());
262 assert_eq!(LEAVES4[1], tree.get_node(NodeIndex::new(2, 1)).unwrap());
263 assert_eq!(LEAVES4[2], tree.get_node(NodeIndex::new(2, 2)).unwrap());
264 assert_eq!(LEAVES4[3], tree.get_node(NodeIndex::new(2, 3)).unwrap());
265
266 // check depth 1
267 let (_, node2, node3) = compute_internal_nodes();
268
269 assert_eq!(node2, tree.get_node(NodeIndex::new(1, 0)).unwrap());
270 assert_eq!(node3, tree.get_node(NodeIndex::new(1, 1)).unwrap());
271 }
272
273 #[test]
274 fn get_path() {
275 let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap();
276
277 let (_, node2, node3) = compute_internal_nodes();
278
279 // check depth 2
280 assert_eq!(
281 vec![LEAVES4[1], node3],
282 *tree.get_path(NodeIndex::new(2, 0)).unwrap()
283 );
284 assert_eq!(
285 vec![LEAVES4[0], node3],
286 *tree.get_path(NodeIndex::new(2, 1)).unwrap()
287 );
288 assert_eq!(
289 vec![LEAVES4[3], node2],
290 *tree.get_path(NodeIndex::new(2, 2)).unwrap()
291 );
292 assert_eq!(
293 vec![LEAVES4[2], node2],
294 *tree.get_path(NodeIndex::new(2, 3)).unwrap()
295 );
296
297 // check depth 1
298 assert_eq!(vec![node3], *tree.get_path(NodeIndex::new(1, 0)).unwrap());
299 assert_eq!(vec![node2], *tree.get_path(NodeIndex::new(1, 1)).unwrap());
300 }
301
302 #[test]
303 fn update_leaf() {
304 let mut tree = super::MerkleTree::new(LEAVES8.to_vec()).unwrap();
305
306 // update one leaf
307 let value = 3;
308 let new_node = int_to_node(9);
309 let mut expected_leaves = LEAVES8.to_vec();
310 expected_leaves[value as usize] = new_node;
311 let expected_tree = super::MerkleTree::new(expected_leaves.clone()).unwrap();
312
313 tree.update_leaf(value, new_node).unwrap();
314 assert_eq!(expected_tree.nodes, tree.nodes);
315
316 // update another leaf
317 let value = 6;
318 let new_node = int_to_node(10);
319 expected_leaves[value as usize] = new_node;
320 let expected_tree = super::MerkleTree::new(expected_leaves.clone()).unwrap();
321
322 tree.update_leaf(value, new_node).unwrap();
323 assert_eq!(expected_tree.nodes, tree.nodes);
324 }
325
326 proptest! {
327 #[test]
328 fn arbitrary_word_can_be_represented_as_digest(
329 a in prop::num::u64::ANY,
330 b in prop::num::u64::ANY,
331 c in prop::num::u64::ANY,
332 d in prop::num::u64::ANY,
333 ) {
334 // this test will assert the memory equivalence between word and digest.
335 // it is used to safeguard the `[MerkleTee::update_leaf]` implementation
336 // that assumes this equivalence.
337
338 // build a word and copy it to another address as digest
339 let word = [Felt::new(a), Felt::new(b), Felt::new(c), Felt::new(d)];
340 let digest = RpoDigest::from(word);
341
342 // assert the addresses are different
343 let word_ptr = (&word).as_ptr() as *const u8;
344 let digest_ptr = (&digest).as_ptr() as *const u8;
345 assert_ne!(word_ptr, digest_ptr);
346
347 // compare the bytes representation
348 let word_bytes = unsafe { slice::from_raw_parts(word_ptr, size_of::<Word>()) };
349 let digest_bytes = unsafe { slice::from_raw_parts(digest_ptr, size_of::<RpoDigest>()) };
350 assert_eq!(word_bytes, digest_bytes);
351 }
352 }
353
354 // HELPER FUNCTIONS
355 // --------------------------------------------------------------------------------------------
356
357 fn compute_internal_nodes() -> (Word, Word, Word) {
358 let node2 = Rpo256::hash_elements(&[LEAVES4[0], LEAVES4[1]].concat());
359 let node3 = Rpo256::hash_elements(&[LEAVES4[2], LEAVES4[3]].concat());
360 let root = Rpo256::merge(&[node2, node3]);
361
362 (root.into(), node2.into(), node3.into())
363 }
364}