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) + Send + Sync>;
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 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}