nmt_rs/simple_merkle/
tree.rs

1use super::db::{LeafWithHash, Node, PreimageDb};
2use super::error::RangeProofError;
3use super::proof::Proof;
4use super::utils::{compute_num_left_siblings, compute_tree_size};
5use crate::maybestd::{boxed::Box, fmt::Debug, hash::Hash, ops::Range, vec::Vec};
6
7/// Manually implement the method we need from #[feature(slice_take)] to
8/// allow building with stable;
9trait TakeLast<T> {
10    fn slice_take_last<'a>(self: &mut &'a Self) -> Option<&'a T>;
11}
12
13impl<T> TakeLast<T> for [T] {
14    fn slice_take_last<'a>(self: &mut &'a Self) -> Option<&'a T> {
15        let (last, rem) = self.split_last()?;
16        *self = rem;
17        Some(last)
18    }
19}
20
21trait TakeFirst<T> {
22    fn slice_take_first<'a>(self: &mut &'a Self) -> Option<&'a T>;
23}
24
25impl<T> TakeFirst<T> for [T] {
26    fn slice_take_first<'a>(self: &mut &'a Self) -> Option<&'a T> {
27        let (first, rem) = self.split_first()?;
28        *self = rem;
29        Some(first)
30    }
31}
32
33type BoxedVisitor<M> = Box<dyn Fn(&<M as MerkleHash>::Output)>;
34
35/// Helper data structure for immutable data used during proof narrowing recursion.
36/// All indices are relative to the leaves of the entire tree.
37struct ProofNarrowingParams<'a, M: MerkleHash> {
38    /// All the leaves inside the old proof range, but to the left of the new (desired) proof range
39    left_extra_leaves: &'a [M::Output],
40    /// The start and end indices of the final, narrower proven range.
41    narrowed_leaf_range: Range<usize>,
42    /// All the leaves inside the old proof range, but to the right of the new (desired) proof range
43    right_extra_leaves: &'a [M::Output],
44    /// The starting index (w.r.t. the tree's leaves) of the old proof; equivalently, the index of
45    /// the first leaf in left_extra_leaves
46    leaves_start_idx: usize,
47}
48
49/// Implements an RFC 6962 compatible merkle tree over an in-memory data store which maps preimages to hashes.
50pub struct MerkleTree<Db, M>
51where
52    M: MerkleHash,
53{
54    leaves: Vec<LeafWithHash<M>>,
55    db: Db,
56    root: Option<M::Output>,
57    visitor: BoxedVisitor<M>,
58    hasher: M,
59}
60
61impl<Db: PreimageDb<<M as MerkleHash>::Output>, M: MerkleHash + Default> Default
62    for MerkleTree<Db, M>
63{
64    fn default() -> Self {
65        Self {
66            leaves: Default::default(),
67            db: Default::default(),
68            root: Default::default(),
69            visitor: Box::new(|_| {}),
70            hasher: Default::default(),
71        }
72    }
73}
74
75/// A trait for hashing data into a merkle tree
76pub trait MerkleHash {
77    // --- no-std
78    /// The output of this hasher.
79    #[cfg(all(not(feature = "serde"), not(feature = "borsh"), not(feature = "std")))]
80    type Output: Debug + PartialEq + Eq + Clone + Default + Hash;
81
82    /// The output of this hasher.
83    #[cfg(all(feature = "serde", not(feature = "borsh"), not(feature = "std")))]
84    type Output: Debug
85        + PartialEq
86        + Eq
87        + Clone
88        + Default
89        + Hash
90        + serde::Serialize
91        + serde::de::DeserializeOwned;
92
93    /// The output of this hasher.
94    #[cfg(all(feature = "borsh", not(feature = "serde"), not(feature = "std")))]
95    type Output: Debug
96        + PartialEq
97        + Eq
98        + Clone
99        + Default
100        + Hash
101        + borsh::BorshSerialize
102        + borsh::BorshDeserialize;
103
104    /// The output of this hasher.
105    #[cfg(all(feature = "borsh", feature = "serde", not(feature = "std")))]
106    type Output: Debug
107        + PartialEq
108        + Eq
109        + Clone
110        + Default
111        + Hash
112        + borsh::BorshSerialize
113        + borsh::BorshDeserialize
114        + serde::Serialize
115        + serde::de::DeserializeOwned;
116
117    // --- std
118    /// The output of this hasher.
119    #[cfg(all(not(feature = "serde"), not(feature = "borsh"), feature = "std"))]
120    type Output: Debug + PartialEq + Eq + Clone + Default + Hash + Ord;
121
122    /// The output of this hasher.
123    #[cfg(all(feature = "serde", not(feature = "borsh"), feature = "std"))]
124    type Output: Debug
125        + PartialEq
126        + Eq
127        + Clone
128        + Default
129        + Hash
130        + Ord
131        + serde::Serialize
132        + serde::de::DeserializeOwned;
133
134    /// The output of this hasher.
135    #[cfg(all(not(feature = "serde"), feature = "borsh", feature = "std"))]
136    type Output: Debug
137        + PartialEq
138        + Eq
139        + Clone
140        + Default
141        + Hash
142        + Ord
143        + borsh::BorshSerialize
144        + borsh::BorshDeserialize;
145
146    /// The output of this hasher.
147    #[cfg(all(feature = "serde", feature = "borsh", feature = "std"))]
148    type Output: Debug
149        + PartialEq
150        + Eq
151        + Clone
152        + Default
153        + Hash
154        + Ord
155        + serde::Serialize
156        + serde::de::DeserializeOwned
157        + borsh::BorshSerialize
158        + borsh::BorshDeserialize;
159
160    /// The hash of the empty tree. This is often defined as the hash of the empty string.
161    const EMPTY_ROOT: Self::Output;
162
163    /// Hashes data as a "leaf" of the tree. This operation *should* be domain separated.
164    fn hash_leaf(&self, data: &[u8]) -> Self::Output;
165    /// Hashes two digests into one. This operation *should* be domain separated.
166    fn hash_nodes(&self, l: &Self::Output, r: &Self::Output) -> Self::Output;
167}
168
169impl<Db, M> MerkleTree<Db, M>
170where
171    Db: PreimageDb<M::Output>,
172    M: MerkleHash + Default,
173{
174    /// Constructs an empty merkle tree with a default hasher
175    pub fn new() -> Self {
176        Self::with_hasher(Default::default())
177    }
178}
179
180impl<Db, M> MerkleTree<Db, M>
181where
182    Db: PreimageDb<M::Output>,
183    M: MerkleHash,
184{
185    /// Constructs an empty merkle tree with the given hasher
186    pub fn with_hasher(hasher: M) -> Self {
187        Self {
188            leaves: Vec::new(),
189            db: Default::default(),
190            root: Some(M::EMPTY_ROOT),
191            visitor: Box::new(|_| {}),
192            hasher,
193        }
194    }
195
196    /// Appends the given leaf to the tree
197    pub fn push_raw_leaf(&mut self, raw_leaf: &[u8]) {
198        let leaf = LeafWithHash::with_hasher(raw_leaf.to_vec(), &self.hasher);
199        self.push_leaf_with_hash(leaf);
200    }
201
202    /// Appends a pre-hashed leaf to the tree
203    pub fn push_leaf_with_hash(&mut self, leaf_with_hash: LeafWithHash<M>) {
204        self.root = None;
205        self.leaves.push(leaf_with_hash);
206    }
207
208    /// Returns the root of the tree, computing it if necessary. Repeated queries return a cached result.
209    pub fn root(&mut self) -> M::Output {
210        if let Some(inner) = &self.root {
211            return inner.clone();
212        }
213        let inner = self.compute_root(0..self.leaves.len());
214        self.root = Some(inner.clone());
215        inner
216    }
217
218    /// Returns the requested range of leaves
219    pub fn get_leaves(&self, range: Range<usize>) -> Vec<Vec<u8>> {
220        let leaves = &self.leaves[range];
221        leaves.iter().map(|leaf| leaf.data().to_vec()).collect()
222    }
223
224    /// Returns all leaves in the tree
225    pub fn leaves(&self) -> &[LeafWithHash<M>] {
226        &self.leaves[..]
227    }
228
229    fn compute_root(&mut self, leaf_range: Range<usize>) -> M::Output {
230        match leaf_range.len() {
231            0 => {
232                let root = M::EMPTY_ROOT;
233                (self.visitor)(&root);
234                root
235            }
236            1 => {
237                let leaf_with_hash = &self.leaves[leaf_range.start];
238                let root = leaf_with_hash.hash().clone();
239                (self.visitor)(&root);
240                self.db
241                    .put(root.clone(), Node::Leaf(leaf_with_hash.data().to_vec()));
242                root
243            }
244            _ => {
245                let split_point = next_smaller_po2(leaf_range.len()) + leaf_range.start;
246                let left = self.compute_root(leaf_range.start..split_point);
247                let right = self.compute_root(split_point..leaf_range.end);
248                let root = self.hasher.hash_nodes(&left, &right);
249                (self.visitor)(&root);
250                self.db.put(root.clone(), Node::Inner(left, right));
251                root
252            }
253        }
254    }
255
256    fn build_range_proof_inner(
257        &self,
258        range_to_prove: Range<usize>,
259        subtrie_root: M::Output,
260        subtrie_range: Range<usize>,
261        out: &mut Vec<M::Output>,
262    ) {
263        if let Some(inner_node) = self.db.get(&subtrie_root) {
264            match inner_node {
265                // If we've bottomed out, return the leaf hash
266                Node::Leaf(_) => {
267                    if !range_to_prove.contains(&subtrie_range.start) {
268                        out.push(subtrie_root.clone())
269                    }
270                }
271                Node::Inner(l, r) => {
272                    let split_point = next_smaller_po2(subtrie_range.len()) + subtrie_range.start;
273                    // If the range to prove, doesn't overlap with the left subtrie, add the left subtrie root to the proof.
274                    // We're now done with the left subtrie
275                    if range_to_prove.start >= split_point {
276                        out.push(l.clone())
277                        //  If the range of nodes to prove completely contains the left subtrie, then we don't need to recurse.
278                    } else if range_to_prove.start > subtrie_range.start
279                        || range_to_prove.end < split_point
280                    {
281                        self.build_range_proof_inner(
282                            range_to_prove.clone(),
283                            l.clone(),
284                            subtrie_range.start..split_point,
285                            out,
286                        );
287                    }
288
289                    // Similarly, if the range to prove, doesn't overlap with the right subtrie, add the right subtrie root to the proof and return
290                    if range_to_prove.end <= split_point {
291                        out.push(r.clone())
292                    } else if range_to_prove.start > split_point
293                        || range_to_prove.end < subtrie_range.end
294                    {
295                        self.build_range_proof_inner(
296                            range_to_prove,
297                            r.clone(),
298                            split_point..subtrie_range.end,
299                            out,
300                        );
301                    }
302                }
303            }
304        } else {
305            assert_eq!(&subtrie_root, &M::EMPTY_ROOT);
306            out.push(subtrie_root)
307        }
308    }
309
310    fn check_range_proof_inner(
311        &self,
312        leaves: &mut &[M::Output],
313        proof: &mut &[M::Output],
314        leaves_start_idx: usize,
315        subtrie_size: usize,
316        offset: usize,
317    ) -> Result<M::Output, RangeProofError> {
318        let split_point = next_smaller_po2(subtrie_size);
319
320        let leaves_end_idx = (leaves.len() + leaves_start_idx) - 1;
321
322        // If the leaf range overlaps with the right subtree
323        let right = if leaves_end_idx >= (split_point + offset) {
324            let right_subtrie_size = subtrie_size - split_point;
325            // If the right subtree contains only a single node, it must be the last remaining leaf
326            if right_subtrie_size == 1 {
327                leaves
328                    .slice_take_last()
329                    .ok_or(RangeProofError::MissingLeaf)?
330                    .clone()
331            } else {
332                // Recurse right
333                self.check_range_proof_inner(
334                    leaves,
335                    proof,
336                    leaves_start_idx,
337                    right_subtrie_size,
338                    offset + split_point,
339                )?
340            }
341        } else {
342            // Otherwise (if the leaf range doesn't overlap with the right subtree),
343            // the sibling node must have been included in the range proof
344            proof
345                .slice_take_last()
346                .ok_or(RangeProofError::MissingProofNode)?
347                .clone()
348        };
349
350        // Similarly, // If the leaf range overlaps with the left subtree
351        let left = if leaves_start_idx < (split_point + offset) {
352            let left_subtrie_size = split_point;
353            // If the right subtree contains only a single node, it must be the last remaining leaf
354            if left_subtrie_size == 1 {
355                leaves
356                    .slice_take_last()
357                    .ok_or(RangeProofError::MissingLeaf)?
358                    .clone()
359            } else {
360                // Recurse left
361                self.check_range_proof_inner(
362                    leaves,
363                    proof,
364                    leaves_start_idx,
365                    left_subtrie_size,
366                    offset,
367                )?
368            }
369        } else {
370            // Otherwise (if the leaf range doesn't overlap with the right subtree),
371            // the sibling node must have been included in the range proof
372            proof
373                .slice_take_last()
374                .ok_or(RangeProofError::MissingProofNode)?
375                .clone()
376        };
377
378        Ok(self.hasher.hash_nodes(&left, &right))
379    }
380
381    /// Helper for the proof narrowing operation.
382    ///
383    /// # Arguments:
384    /// - params: the immutable data used during recursion
385    /// - working_range: The range of leaf indices, relative to the entire tree, being currently
386    ///   considered. Recursion starts with Range(0..tree_size).
387    /// - current_proof: A slice containing the proof of the current, wide range. The slice is
388    ///   mutable as the recursion consumes nodes from it and copies them to the output proof.
389    /// - out: will contain the new proof after recursion finishes
390    fn narrow_range_proof_inner(
391        &self,
392        params: &ProofNarrowingParams<M>,
393        working_range: Range<usize>,
394        current_proof: &mut &[M::Output],
395        out: &mut Vec<M::Output>,
396    ) -> Result<(), RangeProofError> {
397        // Sanity check. This will always be true because:
398        // - At the top level, the working_range is the tree size, and we handle sizes 0 and 1 as
399        // special cases
400        // - When recursing, working_range of length 1 is a base case (we just return the leaf),
401        // so we will never recurse on it
402        assert!(working_range.len() > 1);
403
404        let split_point = next_smaller_po2(working_range.len()) + working_range.start;
405
406        // If the left subtree doesn't overlap with the new leaf, get its root and add it to the proof
407        if params.narrowed_leaf_range.start >= (split_point) {
408            let sibling = self.partial_tree_subroot_inner(
409                working_range.start..split_point,
410                current_proof,
411                params.left_extra_leaves,
412                params.leaves_start_idx,
413            )?;
414            out.push(sibling.clone());
415        } else {
416            let subtree_size = split_point - working_range.start;
417            assert!(subtree_size > 0); // sanity check: since working_range > 1, each sub-tree must be >= 1
418            if subtree_size == 1 {
419                // If it's a leaf, do nothing
420                let index = working_range.start;
421                // Sanity check: if this fails, there's a bug in calculating the range limits and
422                // indices somewhere
423                assert!(params.narrowed_leaf_range.contains(&index));
424            } else {
425                // Else, recurse
426                self.narrow_range_proof_inner(
427                    params,
428                    working_range.start..split_point,
429                    current_proof,
430                    out,
431                )?;
432            }
433        }
434
435        // If the right subtree doesn't overlap with the new leaf, get its root and add it to the proof
436        if params.narrowed_leaf_range.end <= (split_point) {
437            let right_leaves_start_idx = params
438                .leaves_start_idx
439                .checked_add(params.left_extra_leaves.len())
440                .and_then(|i| i.checked_add(params.narrowed_leaf_range.len()))
441                .ok_or(RangeProofError::TreeTooLarge)?;
442            let sibling = self.partial_tree_subroot_inner(
443                split_point..working_range.end,
444                current_proof,
445                params.right_extra_leaves,
446                right_leaves_start_idx,
447            )?;
448            out.push(sibling.clone());
449        } else {
450            let subtree_size = working_range.end - split_point;
451            assert!(subtree_size > 0); // sanity check - see left subtree explanation
452            if subtree_size == 1 {
453                // If it's a leaf, do nothing
454                let index = split_point;
455                assert!(params.narrowed_leaf_range.contains(&index)); // sanity check - see left subtree explanation
456            } else {
457                // Else, recurse
458                self.narrow_range_proof_inner(
459                    params,
460                    split_point..working_range.end,
461                    current_proof,
462                    out,
463                )?;
464            }
465        }
466
467        Ok(())
468    }
469
470    /// To be used during the narrowing operation
471    /// Calculates a new subroot to be part of the narrowed proof,
472    /// in an area covered by the old proof and new leaves.
473    ///
474    /// All indices are relative to the entire tree.
475    ///
476    /// # Arguments
477    ///  - subtree_range: The indices (in the tree) of the leaves of the subtree that we're
478    ///    calculating the subroot of.
479    ///  - extra_leaves: One of the two sets of hashes supplied by the user to narrow down the
480    ///    proof range. Because the two sets are discontiguous, one on each side of the desired new
481    ///    narrower range, only one set at a time is relevant here.
482    ///  - leaves_start_idx: The start of the extra_leaves (relative to the tree). When calculating
483    ///    subroots to the left of the narrowed range (i.e. extra_leaves == left_extra_leaves), this will
484    ///    simply be the (original) proof's start_idx; when calculating subroots to the right, this will
485    ///    be offset correspondingly (i.e. original_start_idx + left_extra_leaves.len() + desired_range_size.len()).
486    fn partial_tree_subroot_inner(
487        &self,
488        subtree_range: Range<usize>,
489        current_proof: &mut &[M::Output],
490        extra_leaves: &[M::Output],
491        leaves_start_idx: usize,
492    ) -> Result<M::Output, RangeProofError> {
493        // Helper that essentially replicates `compute_root`, but with no side-effects and with
494        // only a partial leaf set
495        struct SubrootParams<'a, M: MerkleHash> {
496            extra_leaves: &'a [M::Output],
497            leaves_start_idx: usize,
498            hasher: &'a M,
499        }
500        fn local_subroot_from_leaves<M: MerkleHash>(
501            range: Range<usize>,
502            params: &SubrootParams<M>,
503        ) -> Result<M::Output, RangeProofError> {
504            if range.len() == 1 {
505                return params
506                    .extra_leaves
507                    .get(range.start - params.leaves_start_idx)
508                    .ok_or(RangeProofError::MissingLeaf)
509                    .cloned();
510            } else {
511                let split_point = next_smaller_po2(range.len()) + range.start;
512                let left = local_subroot_from_leaves(range.start..split_point, params)?;
513                let right = local_subroot_from_leaves(split_point..range.end, params)?;
514                Ok(params.hasher.hash_nodes(&left, &right))
515            }
516        }
517
518        // We are operating on a full subtree. So the base cases are (where _ is an unknown leaf,
519        // and # is a leaf included in extra_leaves):
520        //
521        // [####] - the added leaves are covering the entire range; use them to calculate the subroot
522        // [____] - there are no added leaves in the range; there is an existing proof node for this entire subtree
523        // In all other cases, we split as normal and recurse on both subtrees.
524        //
525        // For example:
526        // [___#] - We recurse on the two sub-trees [__] and [_#]. The left one will correspond to
527        // a single proof node hashing both leaves. On the right one, we recurse again
528        // into [_] and [#]. The left one is a single leaf and must also have been included in the
529        // proof; the right one was part of the old proved range, and now supplied as part of
530        // extra_leaves. Now we can hash these two together, and then hash it with the known parent of
531        // the unknown left two nodes to obtain the root for the 4-wide subtree.
532
533        let leaves_end_idx = leaves_start_idx + extra_leaves.len();
534        if leaves_start_idx <= subtree_range.start && leaves_end_idx >= subtree_range.end {
535            local_subroot_from_leaves(
536                subtree_range,
537                &SubrootParams {
538                    extra_leaves,
539                    leaves_start_idx,
540                    hasher: &self.hasher,
541                },
542            )
543        } else if leaves_start_idx >= subtree_range.end || leaves_end_idx <= subtree_range.start {
544            return current_proof
545                .slice_take_first()
546                .ok_or(RangeProofError::MissingProofNode)
547                .cloned();
548        } else {
549            // Sanity check. Both in narrow_range_proof_inner and here, we never recurse on ranges
550            // < 2, as those are base cases (we return the leaves directly).
551            assert!(subtree_range.len() > 1);
552
553            let split_point = next_smaller_po2(subtree_range.len()) + subtree_range.start;
554            let left = self.partial_tree_subroot_inner(
555                subtree_range.start..split_point,
556                current_proof,
557                extra_leaves,
558                leaves_start_idx,
559            )?;
560            let right = self.partial_tree_subroot_inner(
561                split_point..subtree_range.end,
562                current_proof,
563                extra_leaves,
564                leaves_start_idx,
565            )?;
566            return Ok(self.hasher.hash_nodes(&left, &right));
567        }
568    }
569
570    /// Checks a given range proof
571    pub fn check_range_proof(
572        &self,
573        root: &M::Output,
574        leaves: &[M::Output],
575        proof: &[M::Output],
576        leaves_start_idx: usize,
577    ) -> Result<(), RangeProofError> {
578        // As an optimization, the internal call doesn't recurse into subtrees of size smaller than 2
579        // so we need to ensure that the root has size 2 or greater.
580        match leaves.len() {
581            0 => {
582                if root == &M::EMPTY_ROOT && proof.is_empty() {
583                    return Ok(());
584                }
585                return Err(RangeProofError::NoLeavesProvided);
586            }
587            1 => {
588                if proof.is_empty() {
589                    if &leaves[0] == root && leaves_start_idx == 0 {
590                        return Ok(());
591                    }
592                    return Err(RangeProofError::TreeDoesNotContainLeaf);
593                }
594            }
595            _ => {}
596        };
597
598        let num_left_siblings = compute_num_left_siblings(leaves_start_idx);
599        let num_right_siblings = proof
600            .len()
601            .checked_sub(num_left_siblings)
602            .ok_or(RangeProofError::MissingProofNode)?;
603
604        let tree_size = compute_tree_size(num_right_siblings, leaves_start_idx + leaves.len() - 1)?;
605
606        let computed_root = self.check_range_proof_inner(
607            &mut &leaves[..],
608            &mut &proof[..],
609            leaves_start_idx,
610            tree_size,
611            0,
612        )?;
613        if &computed_root == root {
614            return Ok(());
615        }
616        Err(RangeProofError::InvalidRoot)
617    }
618
619    /// Creates a range proof providing the sibling hashes required to show that a set of values really does occur in
620    /// the merkle tree at some half-open range of indices. Intermediate hashes are identified by an in-order traversal
621    /// and are returned in that same order. Panics if the range to prove is larger than the tree's leaf array.
622    ///
623    /// Example: consider the following merkle tree with leaves [C, D, E, F]
624    /// ```ascii
625    ///          root
626    ///        /      \
627    ///       A        B
628    ///      / \      /  \
629    ///     C   D    E    F
630    ///
631    /// ```
632    ///
633    /// A range proof of build_range_proof(1..3) would return the vector [C, F], since those two hashes, together
634    /// with the two leaves in the range, are sufficient to reconstruct the tree
635    pub fn build_range_proof(&mut self, leaf_range: Range<usize>) -> Proof<M> {
636        // Calculate the root to ensure that the preimage db is populated
637        let root = self.root();
638        let mut proof = Vec::new();
639        let start = leaf_range.start as u32;
640        let end = leaf_range.end as u32;
641        if leaf_range.end > self.leaves.len() {
642            panic!(
643                "Index out of range: cannot access leaf {} in leaves array of size {}",
644                leaf_range.end,
645                self.leaves.len()
646            )
647        }
648        self.build_range_proof_inner(leaf_range, root, 0..self.leaves.len(), &mut proof);
649
650        Proof {
651            siblings: proof,
652            range: start..end,
653        }
654    }
655
656    /// Narrows the proof range: uses an existing proof to create
657    /// a new proof for a subrange of the original proof's range.
658    ///
659    /// Effectively, we have two ranges of leaves provided, which can make the range narrower from
660    /// the left or the right respectively (alongside the original proof). The high level logic of
661    /// building a proof out of that is very similar to the normal build_range_proof logic, with
662    /// two exceptions: we don't have the root (or most inner nodes), so we recurse based on the
663    /// leaves and calculate the intermediate hashes we need as we go; and we don't have all the
664    /// leaves either, so the partial_tree_subroot_inner function calculates inner node roots using
665    /// information from both the original proof and the leaves we do have.
666    ///
667    /// Example: consider the following merkle tree with eight leaves:
668    /// ```ascii
669    ///                    root
670    ///                /          \
671    ///            A                  B
672    ///         /    \             /    \
673    ///       C        D         E        F
674    ///      / \      /  \      / \      /  \
675    ///     G   H    I    J    K   L    M    N
676    ///
677    /// ```
678    /// A proof of [H, I, J, K] will contain nodes [G, L, F]. If we want to turn that into a proof
679    /// of [J], that would need nodes [I, C, B].
680    /// We recursively subdivide the total leaf range to find the subtrees that don't overlap the
681    /// final desired range, just as in the normal build_range_proof - in this case, [G, H], [I],
682    /// and [K, L, M, N]. We can then combine the information from the proof and the {left|right}_extra_leaves
683    /// to calculate the subroots of each of those trees - for example, B = hash(E | F), where F is
684    /// from the original proof, and E is calculated using K (from right_extra_leaves) and L (from
685    /// the original proof). Thus we arrive at the new proof for the narrower range.
686    pub fn narrow_range_proof(
687        &mut self,
688        left_extra_leaves: &[M::Output],
689        narrowed_leaf_range: Range<usize>,
690        right_extra_leaves: &[M::Output],
691        current_proof: &mut &[M::Output],
692        leaves_start_idx: usize,
693    ) -> Result<Proof<M>, RangeProofError> {
694        let num_left_siblings = compute_num_left_siblings(leaves_start_idx);
695        let num_right_siblings = current_proof
696            .len()
697            .checked_sub(num_left_siblings)
698            .ok_or(RangeProofError::MissingProofNode)?;
699
700        let current_leaf_size = left_extra_leaves
701            .len()
702            .checked_add(narrowed_leaf_range.len())
703            .ok_or(RangeProofError::TreeTooLarge)?
704            .checked_add(right_extra_leaves.len())
705            .ok_or(RangeProofError::TreeTooLarge)?;
706        let tree_size =
707            compute_tree_size(num_right_siblings, leaves_start_idx + current_leaf_size - 1)?;
708        let mut proof = Vec::new();
709        match tree_size {
710            0 => {
711                if !(current_proof.is_empty()
712                    && left_extra_leaves.is_empty()
713                    && right_extra_leaves.is_empty())
714                {
715                    return Err(RangeProofError::NoLeavesProvided);
716                }
717            }
718            1 => {
719                // For trees of size 1, the root is the only possible proof. An empty proof
720                // is also valid (as the root is provided anyway when verifying).
721                // As these are the only possible options and they are both valid,
722                // there is nothing to be done when narrowing.
723                proof = current_proof.to_vec();
724            }
725            _ => {
726                self.narrow_range_proof_inner(
727                    &ProofNarrowingParams {
728                        left_extra_leaves,
729                        narrowed_leaf_range: narrowed_leaf_range.clone(),
730                        right_extra_leaves,
731                        leaves_start_idx,
732                    },
733                    0..tree_size,
734                    current_proof,
735                    &mut proof,
736                )?;
737            }
738        };
739        Ok(Proof {
740            siblings: proof,
741            // TODO: is it really safe to convert usize to and from u32 everywhere in this library?
742            range: narrowed_leaf_range.start as u32..narrowed_leaf_range.end as u32,
743        })
744    }
745
746    /// Fetches the requested range of leaves, along with a proof of correctness.
747    pub fn get_range_with_proof(&mut self, leaf_range: Range<usize>) -> (Vec<Vec<u8>>, Proof<M>) {
748        let leaves = &self.leaves[leaf_range.clone()];
749        let leaves = leaves.iter().map(|leaf| leaf.data().to_vec()).collect();
750        (leaves, self.build_range_proof(leaf_range))
751    }
752
753    /// Fetches the leaf at the given index, along with a proof of inclusion.
754    pub fn get_index_with_proof(&mut self, idx: usize) -> (Vec<u8>, Proof<M>) {
755        (
756            self.leaves[idx].data().to_vec(),
757            self.build_range_proof(idx..idx + 1),
758        )
759    }
760}
761
762/// Calculates the largest power of two which is strictly less than the argument
763fn next_smaller_po2(int: usize) -> usize {
764    // Calculate the first power of two which is greater than or equal to the argument, then divide by two.
765    int.next_power_of_two() >> 1
766}