snarkvm_console_collections/merkle_tree/mod.rs
1// Copyright 2024 Aleo Network Foundation
2// This file is part of the snarkVM library.
3
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at:
7
8// http://www.apache.org/licenses/LICENSE-2.0
9
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16mod helpers;
17pub use helpers::*;
18
19mod path;
20pub use path::*;
21
22#[cfg(test)]
23mod tests;
24
25use snarkvm_console_types::prelude::*;
26
27use aleo_std::prelude::*;
28
29use std::collections::BTreeMap;
30
31#[cfg(not(feature = "serial"))]
32use rayon::prelude::*;
33
34#[derive(Clone)]
35pub struct MerkleTree<E: Environment, LH: LeafHash<Hash = PH::Hash>, PH: PathHash<Hash = Field<E>>, const DEPTH: u8> {
36 /// The leaf hasher for the Merkle tree.
37 leaf_hasher: LH,
38 /// The path hasher for the Merkle tree.
39 path_hasher: PH,
40 /// The computed root of the full Merkle tree.
41 root: PH::Hash,
42 /// The internal hashes, from root to hashed leaves, of the full Merkle tree.
43 tree: Vec<PH::Hash>,
44 /// The canonical empty hash.
45 empty_hash: Field<E>,
46 /// The number of hashed leaves in the tree.
47 number_of_leaves: usize,
48}
49
50impl<E: Environment, LH: LeafHash<Hash = PH::Hash>, PH: PathHash<Hash = Field<E>>, const DEPTH: u8>
51 MerkleTree<E, LH, PH, DEPTH>
52{
53 #[inline]
54 /// Initializes a new Merkle tree with the given leaves.
55 pub fn new(leaf_hasher: &LH, path_hasher: &PH, leaves: &[LH::Leaf]) -> Result<Self> {
56 let timer = timer!("MerkleTree::new");
57
58 // Ensure the Merkle tree depth is greater than 0.
59 ensure!(DEPTH > 0, "Merkle tree depth must be greater than 0");
60 // Ensure the Merkle tree depth is less than or equal to 64.
61 ensure!(DEPTH <= 64u8, "Merkle tree depth must be less than or equal to 64");
62
63 // Compute the maximum number of leaves.
64 let max_leaves = match leaves.len().checked_next_power_of_two() {
65 Some(num_leaves) => num_leaves,
66 None => bail!("Integer overflow when computing the maximum number of leaves in the Merkle tree"),
67 };
68
69 // Compute the number of nodes.
70 let num_nodes = max_leaves - 1;
71 // Compute the tree size as the maximum number of leaves plus the number of nodes.
72 let tree_size = max_leaves + num_nodes;
73 // Compute the number of levels in the Merkle tree (i.e. log2(tree_size)).
74 let tree_depth = tree_depth::<DEPTH>(tree_size)?;
75 // Compute the number of padded levels.
76 let padding_depth = DEPTH - tree_depth;
77
78 // Compute the empty hash.
79 let empty_hash = path_hasher.hash_empty()?;
80
81 // Calculate the size of the tree which excludes leafless nodes.
82 // The minimum tree size is either a single root node or the calculated number of nodes plus
83 // the supplied leaves; if the number of leaves is odd, an empty hash is added for padding.
84 let minimum_tree_size =
85 std::cmp::max(1, num_nodes + leaves.len() + if leaves.len() > 1 { leaves.len() % 2 } else { 0 });
86
87 // Initialize the Merkle tree.
88 let mut tree = vec![empty_hash; minimum_tree_size];
89
90 // Compute and store each leaf hash.
91 tree[num_nodes..num_nodes + leaves.len()].copy_from_slice(&leaf_hasher.hash_leaves(leaves)?);
92 lap!(timer, "Hashed {} leaves", leaves.len());
93
94 // Compute and store the hashes for each level, iterating from the penultimate level to the root level.
95 let mut start_index = num_nodes;
96 // Compute the start index of the current level.
97 while let Some(start) = parent(start_index) {
98 // Compute the end index of the current level.
99 let end = left_child(start);
100 // Construct the children for each node in the current level; the leaves are padded, which means
101 // that there either are 2 children, or there are none, at which point we may stop iterating.
102 let tuples = (start..end)
103 .take_while(|&i| tree.get(left_child(i)).is_some())
104 .map(|i| (tree[left_child(i)], tree[right_child(i)]))
105 .collect::<Vec<_>>();
106 // Compute and store the hashes for each node in the current level.
107 let num_full_nodes = tuples.len();
108 tree[start..][..num_full_nodes].copy_from_slice(&path_hasher.hash_all_children(&tuples)?);
109 // Use the precomputed empty node hash for every empty node, if there are any.
110 if start + num_full_nodes < end {
111 let empty_node_hash = path_hasher.hash_children(&empty_hash, &empty_hash)?;
112 for node in tree.iter_mut().take(end).skip(start + num_full_nodes) {
113 *node = empty_node_hash;
114 }
115 }
116 // Update the start index for the next level.
117 start_index = start;
118 }
119 lap!(timer, "Hashed {} levels", tree_depth);
120
121 // Compute the root hash, by iterating from the root level up to `DEPTH`.
122 let mut root_hash = tree[0];
123 for _ in 0..padding_depth {
124 // Update the root hash, by hashing the current root hash with the empty hash.
125 root_hash = path_hasher.hash_children(&root_hash, &empty_hash)?;
126 }
127 lap!(timer, "Hashed {} padding levels", padding_depth);
128
129 finish!(timer);
130
131 Ok(Self {
132 leaf_hasher: leaf_hasher.clone(),
133 path_hasher: path_hasher.clone(),
134 root: root_hash,
135 tree,
136 empty_hash,
137 number_of_leaves: leaves.len(),
138 })
139 }
140
141 #[inline]
142 /// Returns a new Merkle tree with the given new leaves appended to it.
143 pub fn prepare_append(&self, new_leaves: &[LH::Leaf]) -> Result<Self> {
144 let timer = timer!("MerkleTree::prepare_append");
145
146 // Compute the maximum number of leaves.
147 let max_leaves = match (self.number_of_leaves + new_leaves.len()).checked_next_power_of_two() {
148 Some(num_leaves) => num_leaves,
149 None => bail!("Integer overflow when computing the maximum number of leaves in the Merkle tree"),
150 };
151 // Compute the number of nodes.
152 let num_nodes = max_leaves - 1;
153 // Compute the tree size as the maximum number of leaves plus the number of nodes.
154 let tree_size = num_nodes + max_leaves;
155 // Compute the number of levels in the Merkle tree (i.e. log2(tree_size)).
156 let tree_depth = tree_depth::<DEPTH>(tree_size)?;
157 // Compute the number of padded levels.
158 let padding_depth = DEPTH - tree_depth;
159
160 // Initialize the Merkle tree.
161 let mut tree = vec![self.empty_hash; num_nodes];
162 // Extend the new Merkle tree with the existing leaf hashes.
163 tree.extend(self.leaf_hashes()?);
164 // Extend the new Merkle tree with the new leaf hashes.
165 tree.extend(&self.leaf_hasher.hash_leaves(new_leaves)?);
166
167 // Calculate the size of the tree which excludes leafless nodes.
168 let new_number_of_leaves = self.number_of_leaves + new_leaves.len();
169 let minimum_tree_size = std::cmp::max(
170 1,
171 num_nodes + new_number_of_leaves + if new_number_of_leaves > 1 { new_number_of_leaves % 2 } else { 0 },
172 );
173
174 // Resize the new Merkle tree with empty hashes to pad up to `tree_size`.
175 tree.resize(minimum_tree_size, self.empty_hash);
176 lap!(timer, "Hashed {} new leaves", new_leaves.len());
177
178 // Initialize a start index to track the starting index of the current level.
179 let start_index = num_nodes;
180 // Initialize a middle index to separate the precomputed indices from the new indices that need to be computed.
181 let middle_index = num_nodes + self.number_of_leaves;
182 // Initialize a precompute index to track the starting index of each precomputed level.
183 let start_precompute_index = match self.number_of_leaves.checked_next_power_of_two() {
184 Some(num_leaves) => num_leaves - 1,
185 None => bail!("Integer overflow when computing the Merkle tree precompute index"),
186 };
187 // Initialize a precompute index to track the middle index of each precomputed level.
188 let middle_precompute_index = match num_nodes == start_precompute_index {
189 // If the old tree and new tree are of the same size, then we can copy over the right half of the old tree.
190 true => Some(start_precompute_index + self.number_of_leaves + new_leaves.len() + 1),
191 // Otherwise, we need to compute the right half of the new tree.
192 false => None,
193 };
194
195 // Compute and store the hashes for each level, iterating from the penultimate level to the root level.
196 self.compute_updated_tree(
197 &mut tree,
198 start_index,
199 middle_index,
200 start_precompute_index,
201 middle_precompute_index,
202 )?;
203
204 // Compute the root hash, by iterating from the root level up to `DEPTH`.
205 let mut root_hash = tree[0];
206 for _ in 0..padding_depth {
207 // Update the root hash, by hashing the current root hash with the empty hash.
208 root_hash = self.path_hasher.hash_children(&root_hash, &self.empty_hash)?;
209 }
210 lap!(timer, "Hashed {} padding levels", padding_depth);
211
212 finish!(timer);
213
214 Ok(Self {
215 leaf_hasher: self.leaf_hasher.clone(),
216 path_hasher: self.path_hasher.clone(),
217 root: root_hash,
218 tree,
219 empty_hash: self.empty_hash,
220 number_of_leaves: self.number_of_leaves + new_leaves.len(),
221 })
222 }
223
224 #[inline]
225 /// Updates the Merkle tree with the given new leaves appended to it.
226 pub fn append(&mut self, new_leaves: &[LH::Leaf]) -> Result<()> {
227 let timer = timer!("MerkleTree::append");
228
229 // Compute the updated Merkle tree with the new leaves.
230 let updated_tree = self.prepare_append(new_leaves)?;
231 // Update the tree at the very end, so the original tree is not altered in case of failure.
232 *self = updated_tree;
233
234 finish!(timer);
235 Ok(())
236 }
237
238 #[inline]
239 /// Updates the Merkle tree at the location of the given leaf index with the new leaf.
240 pub fn update(&mut self, leaf_index: usize, new_leaf: &LH::Leaf) -> Result<()> {
241 let timer = timer!("MerkleTree::update");
242
243 // Compute the updated Merkle tree with the new leaves.
244 let updated_tree = self.prepare_update(leaf_index, new_leaf)?;
245 // Update the tree at the very end, so the original tree is not altered in case of failure.
246 *self = updated_tree;
247
248 finish!(timer);
249 Ok(())
250 }
251
252 #[inline]
253 /// Returns a new Merkle tree with updates at the location of the given leaf index with the new leaf.
254 pub fn prepare_update(&self, leaf_index: usize, new_leaf: &LH::Leaf) -> Result<Self> {
255 let timer = timer!("MerkleTree::prepare_update");
256
257 // Check that the leaf index is within the bounds of the Merkle tree.
258 ensure!(
259 leaf_index < self.number_of_leaves,
260 "Leaf index must be less than the number of leaves in the Merkle tree {leaf_index} , {}",
261 self.number_of_leaves
262 );
263
264 // Allocate a vector to store the path hashes.
265 let mut path_hashes = Vec::with_capacity(DEPTH as usize);
266
267 // Compute and add the new leaf hash to the path hashes.
268 path_hashes.push(self.leaf_hasher.hash_leaf(new_leaf)?);
269 lap!(timer, "Hashed 1 new leaf");
270
271 // Compute the start index (on the left) for the leaf hashes level in the Merkle tree.
272 let start = match self.number_of_leaves.checked_next_power_of_two() {
273 Some(num_leaves) => num_leaves - 1,
274 None => bail!("Integer overflow when computing the Merkle tree start index"),
275 };
276
277 // Compute the new hashes for the path from the leaf to the root.
278 let mut index = start + leaf_index;
279 while let Some(parent) = parent(index) {
280 // Get the left and right child hashes of the parent.
281 let (left, right) = match is_left_child(index) {
282 true => (path_hashes.last().unwrap(), &self.tree[right_child(parent)]),
283 false => (&self.tree[left_child(parent)], path_hashes.last().unwrap()),
284 };
285 // Compute and add the new parent hash to the path hashes.
286 path_hashes.push(self.path_hasher.hash_children(left, right)?);
287 // Update the index to the parent.
288 index = parent;
289 }
290
291 // Compute the number of levels in the Merkle tree (i.e. log2(tree_size)).
292 let tree_depth = tree_depth::<DEPTH>(self.tree.len())?;
293 // Compute the padding depth.
294 let padding_depth = DEPTH - tree_depth;
295
296 // Update the root hash.
297 // This unwrap is safe, as the path hashes vector is guaranteed to have at least one element.
298 let mut root_hash = *path_hashes.last().unwrap();
299 for _ in 0..padding_depth {
300 // Update the root hash, by hashing the current root hash with the empty hash.
301 root_hash = self.path_hasher.hash_children(&root_hash, &self.empty_hash)?;
302 }
303 lap!(timer, "Hashed {} padding levels", padding_depth);
304
305 // Initialize the Merkle tree.
306 let mut tree = Vec::with_capacity(self.tree.len());
307 // Extend the new Merkle tree with the existing leaf hashes.
308 tree.extend(&self.tree);
309
310 // Update the rest of the tree with the new path hashes.
311 let mut index = Some(start + leaf_index);
312 for path_hash in path_hashes {
313 tree[index.unwrap()] = path_hash;
314 index = parent(index.unwrap());
315 }
316
317 finish!(timer);
318
319 Ok(Self {
320 leaf_hasher: self.leaf_hasher.clone(),
321 path_hasher: self.path_hasher.clone(),
322 root: root_hash,
323 tree,
324 empty_hash: self.empty_hash,
325 number_of_leaves: self.number_of_leaves,
326 })
327 }
328
329 #[inline]
330 /// Updates the Merkle tree at the location of the given leaf indices with the new leaves.
331 pub fn update_many(&mut self, updates: &BTreeMap<usize, LH::Leaf>) -> Result<()> {
332 let timer = timer!("MerkleTree::update_many");
333
334 // Check that there are updates to perform.
335 ensure!(!updates.is_empty(), "There must be at least one leaf to update in the Merkle tree");
336
337 // Check that the latest leaf index is less than number of leaves in the Merkle tree.
338 // Note: This unwrap is safe since updates is guaranteed to be non-empty.
339 ensure!(
340 *updates.last_key_value().unwrap().0 < self.number_of_leaves,
341 "Leaf index must be less than the number of leaves in the Merkle tree"
342 );
343
344 // Compute the start index (on the left) for the leaf hashes level in the Merkle tree.
345 let start = match self.number_of_leaves.checked_next_power_of_two() {
346 Some(num_leaves) => num_leaves - 1,
347 None => bail!("Integer overflow when computing the Merkle tree start index"),
348 };
349
350 // A helper to compute the leaf hash.
351 let hash_update = |(leaf_index, leaf): &(&usize, &LH::Leaf)| {
352 self.leaf_hasher.hash_leaf(leaf).map(|hash| (start + **leaf_index, hash))
353 };
354
355 // Hash the leaves and add them to the updated hashes.
356 let leaf_hashes: Vec<(usize, LH::Hash)> = match updates.len() {
357 0..=100 => updates.iter().map(|update| hash_update(&update)).collect::<Result<Vec<_>>>()?,
358 _ => cfg_iter!(updates).map(|update| hash_update(&update)).collect::<Result<Vec<_>>>()?,
359 };
360 lap!(timer, "Hashed {} new leaves", leaf_hashes.len());
361
362 // Store the updated hashes by level.
363 let mut updated_hashes = Vec::new();
364 updated_hashes.push(leaf_hashes);
365
366 // A helper function to compute the path hashes for a given level.
367 type Update<PH> = (usize, (<PH as PathHash>::Hash, <PH as PathHash>::Hash));
368 let compute_path_hashes = |inputs: &[Update<PH>]| match inputs.len() {
369 0..=100 => inputs
370 .iter()
371 .map(|(index, (left, right))| self.path_hasher.hash_children(left, right).map(|hash| (*index, hash)))
372 .collect::<Result<Vec<_>>>(),
373 _ => cfg_iter!(inputs)
374 .map(|(index, (left, right))| self.path_hasher.hash_children(left, right).map(|hash| (*index, hash)))
375 .collect::<Result<Vec<_>>>(),
376 };
377
378 // Compute the depth of the tree. This corresponds to the number of levels of hashes in the tree.
379 let tree_depth = tree_depth::<DEPTH>(self.tree.len())?;
380 // Allocate a vector to store the inputs to the path hasher.
381 let mut inputs = Vec::with_capacity(updated_hashes[0].len());
382 // For each level in the tree, compute the path hashes.
383 // In the first iteration, we compute the path hashes for the updated leaf hashes.
384 // In the subsequent iterations, we compute the path hashes for the updated path hashes, until we reach the root.
385 for level in 0..tree_depth as usize {
386 let mut current = 0;
387 while current < updated_hashes[level].len() {
388 let (current_leaf_index, current_leaf_hash) = updated_hashes[level][current];
389 // Get the sibling of the current leaf.
390 let sibling_leaf_index = match sibling(current_leaf_index) {
391 Some(sibling_index) => sibling_index,
392 // If there is no sibling, then we have reached the root.
393 None => break,
394 };
395 // Check if the sibling hash is the next hash in the vector.
396 let sibling_is_next_hash = match current + 1 < updated_hashes[level].len() {
397 true => updated_hashes[level][current + 1].0 == sibling_leaf_index,
398 false => false,
399 };
400 // Get the sibling hash.
401 // Note: This algorithm assumes that the sibling hash is either the next hash in the vector,
402 // or in the original Merkle tree. Consequently, updates need to be provided in sequential order.
403 // This is enforced by the type of `updates: `BTreeMap<usize, LH::Leaf>`.
404 // If this assumption is violated, then the algorithm will compute incorrect path hashes in the Merkle tree.
405 let sibling_leaf_hash = match sibling_is_next_hash {
406 true => updated_hashes[level][current + 1].1,
407 false => self.tree[sibling_leaf_index],
408 };
409 // Order the current and sibling hashes.
410 let (left, right) = match is_left_child(current_leaf_index) {
411 true => (current_leaf_hash, sibling_leaf_hash),
412 false => (sibling_leaf_hash, current_leaf_hash),
413 };
414 // Compute the parent index.
415 // Note that this unwrap is safe, since we check that the `current_leaf_index` is not the root.
416 let parent_index = parent(current_leaf_index).unwrap();
417 // Add the parent hash to the updated hashes.
418 inputs.push((parent_index, (left, right)));
419 // Update the current index.
420 match sibling_is_next_hash {
421 true => current += 2,
422 false => current += 1,
423 }
424 }
425 // Compute the path hashes for the current level.
426 let path_hashes = compute_path_hashes(&inputs)?;
427 // Add the path hashes to the updated hashes.
428 updated_hashes.push(path_hashes);
429 // Clear the inputs.
430 inputs.clear();
431 }
432
433 // Compute the padding depth.
434 let padding_depth = DEPTH - tree_depth;
435
436 // Update the root hash.
437 // This unwrap is safe, as the updated hashes is guaranteed to have at least one element.
438 let mut root_hash = updated_hashes.last().unwrap()[0].1;
439 for _ in 0..padding_depth {
440 // Update the root hash, by hashing the current root hash with the empty hash.
441 root_hash = self.path_hasher.hash_children(&root_hash, &self.empty_hash)?;
442 }
443 lap!(timer, "Hashed {} padding levels", padding_depth);
444
445 // Update the root hash.
446 self.root = root_hash;
447
448 // Update the rest of the tree with the updated hashes.
449 for (index, hash) in updated_hashes.into_iter().flatten() {
450 self.tree[index] = hash;
451 }
452
453 finish!(timer);
454 Ok(())
455 }
456
457 #[inline]
458 /// Returns a new Merkle tree with the last 'n' leaves removed from it.
459 pub fn prepare_remove_last_n(&self, n: usize) -> Result<Self> {
460 let timer = timer!("MerkleTree::prepare_remove_last_n");
461
462 ensure!(n > 0, "Cannot remove zero leaves from the Merkle tree");
463
464 // Determine the updated number of leaves, after removing the last 'n' leaves.
465 let updated_number_of_leaves = self.number_of_leaves.checked_sub(n).ok_or_else(|| {
466 anyhow!("Failed to remove '{n}' leaves from the Merkle tree, as it only contains {}", self.number_of_leaves)
467 })?;
468
469 // Compute the maximum number of leaves.
470 let max_leaves = match (updated_number_of_leaves).checked_next_power_of_two() {
471 Some(num_leaves) => num_leaves,
472 None => bail!("Integer overflow when computing the maximum number of leaves in the Merkle tree"),
473 };
474 // Compute the number of nodes.
475 let num_nodes = max_leaves - 1;
476 // Compute the tree size as the maximum number of leaves plus the number of nodes.
477 let tree_size = num_nodes + max_leaves;
478 // Compute the number of levels in the Merkle tree (i.e. log2(tree_size)).
479 let tree_depth = tree_depth::<DEPTH>(tree_size)?;
480 // Compute the number of padded levels.
481 let padding_depth = DEPTH - tree_depth;
482
483 // Calculate the size of the tree which excludes leafless nodes.
484 let minimum_tree_size = std::cmp::max(
485 1,
486 num_nodes
487 + updated_number_of_leaves
488 + if updated_number_of_leaves > 1 { updated_number_of_leaves % 2 } else { 0 },
489 );
490
491 // Initialize the Merkle tree.
492 let mut tree = vec![self.empty_hash; num_nodes];
493 // Extend the new Merkle tree with the existing leaf hashes, excluding the last 'n' leaves.
494 tree.extend(&self.leaf_hashes()?[..updated_number_of_leaves]);
495 // Resize the new Merkle tree with empty hashes to pad up to `tree_size`.
496 tree.resize(minimum_tree_size, self.empty_hash);
497 lap!(timer, "Resizing to {} leaves", updated_number_of_leaves);
498
499 // Initialize a start index to track the starting index of the current level.
500 let start_index = num_nodes;
501 // Initialize a middle index to separate the precomputed indices from the new indices that need to be computed.
502 let middle_index = num_nodes + updated_number_of_leaves;
503 // Initialize a precompute index to track the starting index of each precomputed level.
504 let start_precompute_index = match self.number_of_leaves.checked_next_power_of_two() {
505 Some(num_leaves) => num_leaves - 1,
506 None => bail!("Integer overflow when computing the Merkle tree precompute index"),
507 };
508 // Initialize a precompute index to track the middle index of each precomputed level.
509 let middle_precompute_index = match num_nodes == start_precompute_index {
510 // If the old tree and new tree are of the same size, then we can copy over the right half of the old tree.
511 true => Some(start_precompute_index + self.number_of_leaves + 1),
512 // true => None,
513 // Otherwise, do nothing, since shrinking the tree is already free.
514 false => None,
515 };
516
517 // Compute and store the hashes for each level, iterating from the penultimate level to the root level.
518 self.compute_updated_tree(
519 &mut tree,
520 start_index,
521 middle_index,
522 start_precompute_index,
523 middle_precompute_index,
524 )?;
525
526 // Compute the root hash, by iterating from the root level up to `DEPTH`.
527 let mut root_hash = tree[0];
528 for _ in 0..padding_depth {
529 // Update the root hash, by hashing the current root hash with the empty hash.
530 root_hash = self.path_hasher.hash_children(&root_hash, &self.empty_hash)?;
531 }
532 lap!(timer, "Hashed {} padding levels", padding_depth);
533
534 finish!(timer);
535
536 Ok(Self {
537 leaf_hasher: self.leaf_hasher.clone(),
538 path_hasher: self.path_hasher.clone(),
539 root: root_hash,
540 tree,
541 empty_hash: self.empty_hash,
542 number_of_leaves: updated_number_of_leaves,
543 })
544 }
545
546 #[inline]
547 /// Updates the Merkle tree with the last 'n' leaves removed from it.
548 pub fn remove_last_n(&mut self, n: usize) -> Result<()> {
549 let timer = timer!("MerkleTree::remove_last_n");
550
551 // Compute the updated Merkle tree with the last 'n' leaves removed.
552 let updated_tree = self.prepare_remove_last_n(n)?;
553 // Update the tree at the very end, so the original tree is not altered in case of failure.
554 *self = updated_tree;
555
556 finish!(timer);
557 Ok(())
558 }
559
560 #[inline]
561 /// Returns the Merkle path for the given leaf index and leaf.
562 pub fn prove(&self, leaf_index: usize, leaf: &LH::Leaf) -> Result<MerklePath<E, DEPTH>> {
563 // Ensure the leaf index is valid.
564 ensure!(leaf_index < self.number_of_leaves, "The given Merkle leaf index is out of bounds");
565
566 // Compute the leaf hash.
567 let leaf_hash = self.leaf_hasher.hash_leaf(leaf)?;
568
569 // Compute the start index (on the left) for the leaf hashes level in the Merkle tree.
570 let start = match self.number_of_leaves.checked_next_power_of_two() {
571 Some(num_leaves) => num_leaves - 1,
572 None => bail!("Integer overflow when computing the Merkle tree start index"),
573 };
574
575 // Compute the absolute index of the leaf in the Merkle tree.
576 let mut index = start + leaf_index;
577 // Ensure the leaf index is valid.
578 ensure!(index < self.tree.len(), "The given Merkle leaf index is out of bounds");
579 // Ensure the leaf hash matches the one in the tree.
580 ensure!(self.tree[index] == leaf_hash, "The given Merkle leaf does not match the one in the Merkle tree");
581
582 // Initialize a vector for the Merkle path.
583 let mut path = Vec::with_capacity(DEPTH as usize);
584
585 // Iterate from the leaf hash to the root level, storing the sibling hashes along the path.
586 for _ in 0..DEPTH {
587 // Compute the index of the sibling hash, if it exists.
588 if let Some(sibling) = sibling(index) {
589 // Append the sibling hash to the path.
590 path.push(self.tree[sibling]);
591 // Compute the index of the parent hash, if it exists.
592 match parent(index) {
593 // Update the index to the parent index.
594 Some(parent) => index = parent,
595 // If the parent does not exist, the path is complete.
596 None => break,
597 }
598 }
599 }
600
601 // If the Merkle path length is not equal to `DEPTH`, pad the path with the empty hash.
602 path.resize(DEPTH as usize, self.empty_hash);
603
604 // Return the Merkle path.
605 MerklePath::try_from((U64::new(leaf_index as u64), path))
606 }
607
608 /// Returns `true` if the given Merkle path is valid for the given root and leaf.
609 pub fn verify(&self, path: &MerklePath<E, DEPTH>, root: &PH::Hash, leaf: &LH::Leaf) -> bool {
610 path.verify(&self.leaf_hasher, &self.path_hasher, root, leaf)
611 }
612
613 /// Returns the Merkle root of the tree.
614 pub const fn root(&self) -> &PH::Hash {
615 &self.root
616 }
617
618 /// Returns the Merkle tree (excluding the hashes of the leaves).
619 pub fn tree(&self) -> &[PH::Hash] {
620 &self.tree
621 }
622
623 /// Returns the empty hash.
624 pub const fn empty_hash(&self) -> &PH::Hash {
625 &self.empty_hash
626 }
627
628 /// Returns the leaf hashes from the Merkle tree.
629 pub fn leaf_hashes(&self) -> Result<&[LH::Hash]> {
630 // Compute the start index (on the left) for the leaf hashes level in the Merkle tree.
631 let start = match self.number_of_leaves.checked_next_power_of_two() {
632 Some(num_leaves) => num_leaves - 1,
633 None => bail!("Integer overflow when computing the Merkle tree start index"),
634 };
635 // Compute the end index (on the right) for the leaf hashes level in the Merkle tree.
636 let end = start + self.number_of_leaves;
637 // Return the leaf hashes.
638 Ok(&self.tree[start..end])
639 }
640
641 /// Returns the number of leaves in the Merkle tree.
642 pub const fn number_of_leaves(&self) -> usize {
643 self.number_of_leaves
644 }
645
646 /// Compute and store the hashes for each level, iterating from the penultimate level to the root level.
647 ///
648 /// ```ignore
649 /// start_index middle_index end_index
650 /// start_precompute_index middle_precompute_index end_index
651 /// ```
652 #[inline]
653 fn compute_updated_tree(
654 &self,
655 tree: &mut [Field<E>],
656 mut start_index: usize,
657 mut middle_index: usize,
658 mut start_precompute_index: usize,
659 mut middle_precompute_index: Option<usize>,
660 ) -> Result<()> {
661 // Initialize a timer for the while loop.
662 let timer = timer!("MerkleTree::compute_updated_tree");
663
664 // Compute and store the hashes for each level, iterating from the penultimate level to the root level.
665 let empty_hash = self.path_hasher.hash_empty()?;
666 while let (Some(start), Some(middle)) = (parent(start_index), parent(middle_index)) {
667 // Compute the end index of the current level.
668 let end = left_child(start);
669
670 // If the current level has precomputed indices, copy them instead of recomputing them.
671 if let Some(start_precompute) = parent(start_precompute_index) {
672 // Compute the end index of the precomputed level.
673 let end_precompute = start_precompute + (middle - start);
674 // Copy the hashes for each node in the current level.
675 tree[start..middle].copy_from_slice(&self.tree[start_precompute..end_precompute]);
676 // Update the precompute index for the next level.
677 start_precompute_index = start_precompute;
678 } else {
679 // Ensure the start index is equal to the middle index, as all precomputed indices have been processed.
680 ensure!(start == middle, "Failed to process all left precomputed indices in the Merkle tree");
681 }
682 lap!(timer, "Precompute (Left): {start} -> {middle}");
683
684 // If the current level has precomputed indices, copy them instead of recomputing them.
685 // Note: This logic works because the old tree and new tree are the same power of two.
686 if let Some(middle_precompute) = middle_precompute_index {
687 if let Some(middle_precompute) = parent(middle_precompute) {
688 // Construct the children for the new indices in the current level.
689 let tuples = (middle..middle_precompute)
690 .map(|i| {
691 (
692 tree.get(left_child(i)).copied().unwrap_or(empty_hash),
693 tree.get(right_child(i)).copied().unwrap_or(empty_hash),
694 )
695 })
696 .collect::<Vec<_>>();
697 // Process the indices that need to be computed for the current level.
698 // If any level requires computing more than 100 nodes, borrow the tree for performance.
699 match tuples.len() >= 100 {
700 // Option 1: Borrow the tree to compute and store the hashes for the new indices in the current level.
701 true => cfg_iter_mut!(tree[middle..middle_precompute]).zip_eq(cfg_iter!(tuples)).try_for_each(
702 |(node, (left, right))| {
703 *node = self.path_hasher.hash_children(left, right)?;
704 Ok::<_, Error>(())
705 },
706 )?,
707 // Option 2: Compute and store the hashes for the new indices in the current level.
708 false => tree[middle..middle_precompute].iter_mut().zip_eq(&tuples).try_for_each(
709 |(node, (left, right))| {
710 *node = self.path_hasher.hash_children(left, right)?;
711 Ok::<_, Error>(())
712 },
713 )?,
714 }
715 lap!(timer, "Compute: {middle} -> {middle_precompute}");
716
717 // Copy the hashes for each node in the current level.
718 tree[middle_precompute..end].copy_from_slice(&self.tree[middle_precompute..end]);
719 // Update the precompute index for the next level.
720 middle_precompute_index = Some(middle_precompute + 1);
721 lap!(timer, "Precompute (Right): {middle_precompute} -> {end}");
722 } else {
723 // Ensure the middle precompute index is equal to the end index, as all precomputed indices have been processed.
724 ensure!(
725 middle_precompute == end,
726 "Failed to process all right precomputed indices in the Merkle tree"
727 );
728 }
729 } else {
730 // Construct the children for the new indices in the current level.
731 let tuples = (middle..end)
732 .map(|i| {
733 (
734 tree.get(left_child(i)).copied().unwrap_or(empty_hash),
735 tree.get(right_child(i)).copied().unwrap_or(empty_hash),
736 )
737 })
738 .collect::<Vec<_>>();
739 // Process the indices that need to be computed for the current level.
740 // If any level requires computing more than 100 nodes, borrow the tree for performance.
741 match tuples.len() >= 100 {
742 // Option 1: Borrow the tree to compute and store the hashes for the new indices in the current level.
743 true => cfg_iter_mut!(tree[middle..end]).zip_eq(cfg_iter!(tuples)).try_for_each(
744 |(node, (left, right))| {
745 *node = self.path_hasher.hash_children(left, right)?;
746 Ok::<_, Error>(())
747 },
748 )?,
749 // Option 2: Compute and store the hashes for the new indices in the current level.
750 false => tree[middle..end].iter_mut().zip_eq(&tuples).try_for_each(|(node, (left, right))| {
751 *node = self.path_hasher.hash_children(left, right)?;
752 Ok::<_, Error>(())
753 })?,
754 }
755 lap!(timer, "Compute: {middle} -> {end}");
756 }
757
758 // Update the start index for the next level.
759 start_index = start;
760 // Update the middle index for the next level.
761 middle_index = middle;
762 }
763
764 // End the timer for the while loop.
765 finish!(timer);
766
767 Ok(())
768 }
769}
770
771/// Returns the depth of the tree, given the size of the tree.
772#[inline]
773fn tree_depth<const DEPTH: u8>(tree_size: usize) -> Result<u8> {
774 let tree_size = u64::try_from(tree_size)?;
775 // Since we only allow tree sizes up to u64::MAX, the maximum possible depth is 63.
776 let tree_depth = u8::try_from(tree_size.checked_ilog2().unwrap_or(0))?;
777 // Ensure the tree depth is within the depth bound.
778 match tree_depth <= DEPTH {
779 // Return the tree depth.
780 true => Ok(tree_depth),
781 false => bail!("Merkle tree cannot exceed depth {DEPTH}: attempted to reach depth {tree_depth}"),
782 }
783}
784
785/// Returns the index of the left child, given an index.
786#[inline]
787const fn left_child(index: usize) -> usize {
788 2 * index + 1
789}
790
791/// Returns the index of the right child, given an index.
792#[inline]
793const fn right_child(index: usize) -> usize {
794 2 * index + 2
795}
796
797/// Returns the index of the sibling, given an index.
798#[inline]
799const fn sibling(index: usize) -> Option<usize> {
800 if is_root(index) {
801 None
802 } else if is_left_child(index) {
803 Some(index + 1)
804 } else {
805 Some(index - 1)
806 }
807}
808
809/// Returns true iff the index represents the root.
810#[inline]
811const fn is_root(index: usize) -> bool {
812 index == 0
813}
814
815/// Returns true iff the given index represents a left child.
816#[inline]
817const fn is_left_child(index: usize) -> bool {
818 index % 2 == 1
819}
820
821/// Returns the index of the parent, given the index of a child.
822#[inline]
823const fn parent(index: usize) -> Option<usize> {
824 if index > 0 { Some((index - 1) >> 1) } else { None }
825}