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