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}