Skip to main content

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}