ab_merkle_tree/
unbalanced.rs

1use crate::{MerkleHasher, NewMerkleHasherForTreeSize};
2#[cfg(feature = "alloc")]
3use alloc::boxed::Box;
4#[cfg(feature = "alloc")]
5use alloc::vec::Vec;
6use core::iter::TrustedLen;
7use core::marker::PhantomCovariant;
8use core::mem::MaybeUninit;
9
10/// Number of elements in a proof for a tree with `MAX_N` leaves
11pub const PROOF_ELEMENTS<const MAX_N: u64>: usize = MAX_N.next_power_of_two().ilog2() as usize;
12const STACK_SIZE<const MAX_N: u64>: usize = MAX_N.next_power_of_two().ilog2() as usize + 1;
13const SUPPORTED_LEAVES_ARRAY_SIZE<const N: usize>: usize = {
14    assert!(N != 0, "Number of leaves must be greater than 0");
15    0
16};
17const USIZE_TO_U64<const N: usize>: u64 = N as u64;
18
19/// Merkle Tree variant that has pre-hashed leaves with arbitrary number of elements.
20///
21/// This can be considered a general case of [`BalancedMerkleTree`]. The root and proofs are
22/// identical for both in case the number of leaves is a power of two. [`BalancedMerkleTree`] is
23/// more efficient and should be preferred when possible.
24///
25/// [`BalancedMerkleTree`]: crate::balanced::BalancedMerkleTree
26///
27/// The unbalanced tree is not padded, it is created the same way Merkle Mountain Range would be:
28/// ```text
29///               Root
30///         /--------------\
31///        H3              H4
32///    /-------\         /----\
33///   H0       H1       H2     \
34///  /  \     /  \     /  \     \
35/// L0  L1   L2  L3   L4  L5    L6
36/// ```
37///
38/// Inner nodes are hashed with `H`, which is created from the number of leaves. That is why methods
39/// that take an iterator of leaves require a [`TrustedLen`] iterator. Variants with `dynamic_len`
40/// in their name accept any iterator instead, but only support hashers that implement [`Default`].
41#[derive(Debug)]
42pub struct UnbalancedMerkleTree<H>(PhantomCovariant<H>);
43
44// TODO: Experiment with replacing a single pass with splitting the whole data set with a sequence
45//  of power-of-two elements that can be processed in parallel and do it recursively until a single
46//  element is left. This can be done for both root creation and proof generation.
47impl<H> UnbalancedMerkleTree<H>
48where
49    H: MerkleHasher,
50{
51    /// Compute Merkle Tree Root.
52    ///
53    /// `MAX_N` generic constant defines the maximum number of elements supported and controls stack
54    /// usage.
55    ///
56    /// Returns `None` for an empty list of leaves, or if the number of leaves is larger than
57    /// `MAX_N`.
58    #[inline]
59    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
60    pub fn compute_root_only<'a, const MAX_N: u64, Item, Iter>(
61        leaves: Iter,
62    ) -> Option<[u8; H::OUT_LEN]>
63    where
64        H: NewMerkleHasherForTreeSize,
65        Item: Into<[u8; H::OUT_LEN]>,
66        Iter: IntoIterator<Item = Item, IntoIter: TrustedLen> + 'a,
67    {
68        let leaves = leaves.into_iter();
69        let hasher = Self::hasher_for_leaves(&leaves)?;
70
71        Self::compute_root_only_with::<MAX_N, _, _>(&hasher, leaves)
72    }
73
74    /// Similar to [`Self::compute_root_only()`], but for iterators that are not [`TrustedLen`]
75    #[inline]
76    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
77    pub fn compute_root_only_dynamic_len<'a, const MAX_N: u64, Item, Iter>(
78        leaves: Iter,
79    ) -> Option<[u8; H::OUT_LEN]>
80    where
81        H: Default,
82        Item: Into<[u8; H::OUT_LEN]>,
83        Iter: IntoIterator<Item = Item> + 'a,
84    {
85        let hasher = H::default();
86
87        Self::compute_root_only_with::<MAX_N, _, _>(&hasher, leaves)
88    }
89
90    #[inline]
91    fn compute_root_only_with<const MAX_N: u64, Item, Iter>(
92        hasher: &H,
93        leaves: Iter,
94    ) -> Option<[u8; H::OUT_LEN]>
95    where
96        Item: Into<[u8; H::OUT_LEN]>,
97        Iter: IntoIterator<Item = Item>,
98    {
99        // Stack of intermediate nodes per tree level
100        let mut stack = [[0u8; H::OUT_LEN]; STACK_SIZE::<MAX_N>];
101        let mut num_leaves = 0u64;
102
103        for hash in leaves {
104            // How many leaves were processed so far. Should have been `num_leaves == MAX_N`, but
105            // `>=` helps compiler with panic safety checks.
106            // TODO: Use `==` once https://github.com/rust-lang/rust/issues/162834 is resolved
107            if num_leaves >= MAX_N {
108                return None;
109            }
110
111            let mut current = hash.into();
112
113            // Every bit set to `1` corresponds to an active Merkle Tree level
114            let lowest_active_levels = num_leaves.trailing_ones() as usize;
115            for item in stack.iter().take(lowest_active_levels) {
116                current = hasher.hash_pair(&[*item, current]);
117            }
118
119            // Place the current hash at the first inactive level
120            stack[lowest_active_levels] = current;
121            num_leaves += 1;
122        }
123
124        if num_leaves == 0 {
125            // If no leaves were provided
126            return None;
127        }
128
129        // SAFETY: `num_leaves` is not zero, not larger than `MAX_N`, and peaks of all `num_leaves`
130        // leaves are on the stack
131        Some(unsafe { Self::root_from_peaks::<MAX_N>(hasher, &mut stack, num_leaves) })
132    }
133
134    /// Hash peaks of `num_leaves` leaves on the `stack` together into the root.
135    ///
136    /// # Safety
137    /// `num_leaves` must not be zero or larger than `MAX_N`, and peaks of all `num_leaves` leaves
138    /// must be on the `stack`.
139    #[inline(always)]
140    unsafe fn root_from_peaks<const MAX_N: u64>(
141        hasher: &H,
142        stack: &mut [[u8; H::OUT_LEN]; STACK_SIZE::<MAX_N>],
143        num_leaves: u64,
144    ) -> [u8; H::OUT_LEN] {
145        let mut stack_bits = num_leaves;
146
147        {
148            let lowest_active_level = stack_bits.trailing_zeros() as usize;
149            // Reuse `stack[0]` for resulting value
150            // SAFETY: Active level must have been set successfully before, hence it exists
151            stack[0] = *unsafe { stack.get_unchecked(lowest_active_level) };
152            // Clear lowest active level
153            stack_bits &= !(1 << lowest_active_level);
154        }
155
156        // Hash remaining peaks (if any) of the potentially unbalanced tree together
157        loop {
158            let lowest_active_level = stack_bits.trailing_zeros() as usize;
159
160            if lowest_active_level == u64::BITS as usize {
161                break;
162            }
163
164            // Clear lowest active level for next iteration
165            stack_bits &= !(1 << lowest_active_level);
166
167            // SAFETY: Active level must have been set successfully before, hence it exists
168            let lowest_active_level_item = unsafe { stack.get_unchecked(lowest_active_level) };
169
170            stack[0] = hasher.hash_pair(&[*lowest_active_level_item, stack[0]]);
171        }
172
173        if num_leaves == 1 {
174            stack[0] = hasher.hash_single_leaf(&stack[0]);
175        }
176
177        stack[0]
178    }
179
180    /// Compute Merkle Tree Root of `num_leaves` leaves, which are all zeroes except `leaves`.
181    ///
182    /// `leaves` are pairs of leaf index and leaf, with strictly increasing indices. Produces the
183    /// same root as [`Self::compute_root_only()`] with zeroes in all other positions, but hashes
184    /// runs of zero leaves as whole subtrees, so the number of hashes is proportional to the number
185    /// of `leaves` and the height of the tree rather than to `num_leaves`.
186    ///
187    /// `MAX_N` generic constant defines the maximum number of elements supported and controls stack
188    /// usage.
189    ///
190    /// Returns `None` if `num_leaves` is zero or larger than `MAX_N`, or if indices aren't strictly
191    /// increasing or not smaller than `num_leaves`.
192    #[inline]
193    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
194    pub fn compute_root_only_sparse<const MAX_N: u64, Item, Iter>(
195        num_leaves: u64,
196        leaves: Iter,
197    ) -> Option<[u8; H::OUT_LEN]>
198    where
199        H: NewMerkleHasherForTreeSize,
200        Item: Into<[u8; H::OUT_LEN]>,
201        Iter: IntoIterator<Item = (u64, Item)>,
202    {
203        if num_leaves == 0 || num_leaves > MAX_N {
204            return None;
205        }
206
207        let hasher = H::new_for_tree_size(num_leaves);
208        // Stack of intermediate nodes per tree level
209        let mut stack = [[0u8; H::OUT_LEN]; STACK_SIZE::<MAX_N>];
210        // Roots of perfect subtrees of zero leaves per height, filled up to `zeros_height`
211        let mut zeros = [[0u8; H::OUT_LEN]; STACK_SIZE::<MAX_N>];
212        let mut zeros_height = 0;
213        // Number of leaves processed so far, including zero leaves
214        let mut processed = 0u64;
215
216        for (index, leaf) in leaves {
217            if index < processed || index >= num_leaves {
218                return None;
219            }
220
221            // SAFETY: `processed <= index < num_leaves <= MAX_N`
222            unsafe {
223                Self::push_zero_leaves::<MAX_N>(
224                    &hasher,
225                    &mut stack,
226                    &mut zeros,
227                    &mut zeros_height,
228                    &mut processed,
229                    index,
230                );
231            }
232            // SAFETY: Any `processed` is a multiple of `2^0` and `processed + 1 <= num_leaves`,
233            // which is not larger than `MAX_N`
234            unsafe {
235                Self::push_subtree::<MAX_N>(&hasher, &mut stack, processed, 0, leaf.into());
236            }
237            processed += 1;
238        }
239
240        // SAFETY: `processed <= num_leaves <= MAX_N`
241        unsafe {
242            Self::push_zero_leaves::<MAX_N>(
243                &hasher,
244                &mut stack,
245                &mut zeros,
246                &mut zeros_height,
247                &mut processed,
248                num_leaves,
249            );
250        }
251
252        // SAFETY: `num_leaves` is not zero, not larger than `MAX_N`, and peaks of all `num_leaves`
253        // leaves are on the stack
254        Some(unsafe { Self::root_from_peaks::<MAX_N>(&hasher, &mut stack, num_leaves) })
255    }
256
257    /// Push zero leaves until `processed` reaches `target`, as perfect subtrees that are as large
258    /// as possible.
259    ///
260    /// # Safety
261    /// `processed` must not be larger than `target`, which must not be larger than `MAX_N`.
262    #[inline(always)]
263    unsafe fn push_zero_leaves<const MAX_N: u64>(
264        hasher: &H,
265        stack: &mut [[u8; H::OUT_LEN]; STACK_SIZE::<MAX_N>],
266        zeros: &mut [[u8; H::OUT_LEN]; STACK_SIZE::<MAX_N>],
267        zeros_height: &mut usize,
268        processed: &mut u64,
269        target: u64,
270    ) {
271        while let Some(height) = (target - *processed).checked_ilog2() {
272            // The largest subtree that fits into the remaining zero leaves and starts at a multiple
273            // of its size
274            let height = height.min(processed.trailing_zeros()) as usize;
275
276            while *zeros_height < height {
277                // SAFETY: `height` is at most `log2(MAX_N)`, which is within `STACK_SIZE`
278                let zero = unsafe { zeros.get_unchecked(*zeros_height) };
279                let zero = hasher.hash_pair(&[*zero, *zero]);
280                *zeros_height += 1;
281                // SAFETY: `height` is at most `log2(MAX_N)`, which is within `STACK_SIZE`
282                *unsafe { zeros.get_unchecked_mut(*zeros_height) } = zero;
283            }
284
285            // SAFETY: `height` is at most `log2(MAX_N)`, which is within `STACK_SIZE`
286            let zero = *unsafe { zeros.get_unchecked(height) };
287            // SAFETY: `processed` is a multiple of `2^height` and `processed + 2^height <= target`,
288            // which is not larger than `MAX_N`
289            unsafe {
290                Self::push_subtree::<MAX_N>(hasher, stack, *processed, height, zero);
291            }
292            *processed += 1 << height;
293        }
294    }
295
296    /// Push the root of a perfect subtree of height `height` after `processed` leaves, merging it
297    /// with the peaks it completes the same way as [`Self::compute_root_only()`] merges a leaf.
298    ///
299    /// # Safety
300    /// `processed` must be a multiple of `2^height`, and `processed + 2^height` must not be larger
301    /// than `MAX_N`.
302    #[inline(always)]
303    unsafe fn push_subtree<const MAX_N: u64>(
304        hasher: &H,
305        stack: &mut [[u8; H::OUT_LEN]; STACK_SIZE::<MAX_N>],
306        processed: u64,
307        height: usize,
308        mut node: [u8; H::OUT_LEN],
309    ) {
310        // Every bit set to `1` corresponds to an active Merkle Tree level
311        let active_levels = (processed >> height).trailing_ones() as usize;
312        for item in stack.iter().skip(height).take(active_levels) {
313            node = hasher.hash_pair(&[*item, node]);
314        }
315
316        // Place the node at the first inactive level.
317        //
318        // SAFETY: Bit `height + active_levels` is set in `processed + 2^height`, which is not
319        // larger than `MAX_N`, so the level is at most `log2(MAX_N)`, which is within `STACK_SIZE`
320        *unsafe { stack.get_unchecked_mut(height + active_levels) } = node;
321    }
322
323    /// Compute Merkle Tree Root for an array of leaves.
324    ///
325    /// Similar to [`Self::compute_root_only()`], but for inputs that are small arrays where `MAX_N`
326    /// is the same as the array length.
327    #[inline(always)]
328    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
329    pub fn compute_root_only_array<const N: usize, Item>(leaves: &[Item; N]) -> [u8; H::OUT_LEN]
330    where
331        H: const NewMerkleHasherForTreeSize,
332        [(); SUPPORTED_LEAVES_ARRAY_SIZE::<N>]:,
333        Item: Into<[u8; H::OUT_LEN]> + Copy,
334    {
335        // TODO: Technically redundant with the where-clause, which is not enforced in generic code
336        //  due to https://github.com/rust-lang/rust/issues/164019, remove once fixed
337        const { SUPPORTED_LEAVES_ARRAY_SIZE::<N> };
338
339        let hasher = const { H::new_for_tree_size(N as u64) };
340        // Iterating over indices rather than with `leaves.iter().copied()` allows the compiler to
341        // prove that the root always exists on aarch64 and riscv64 too
342        let leaves = (0..N).map(|index| *leaves.get(index).expect("Index is within array; qed"));
343        let maybe_root =
344            Self::compute_root_only_with::<{ USIZE_TO_U64::<N> }, _, _>(&hasher, leaves);
345        maybe_root.expect("Non-empty array with length matching `MAX_N` always has a root; qed")
346    }
347
348    /// Compute Merkle Tree root and generate a proof for the `leaf` at `target_index`.
349    ///
350    /// Returns `Some((root, proof))` on success, `None` if index is outside of list of leaves.
351    ///
352    /// `MAX_N` generic constant defines the maximum number of elements supported and controls stack
353    /// usage.
354    #[expect(clippy::type_complexity, reason = "Return type")]
355    #[inline]
356    #[cfg(feature = "alloc")]
357    pub fn compute_root_and_proof<'a, const MAX_N: u64, Item, Iter>(
358        leaves: Iter,
359        target_index: usize,
360    ) -> Option<([u8; H::OUT_LEN], Vec<[u8; H::OUT_LEN]>)>
361    where
362        H: NewMerkleHasherForTreeSize,
363        Item: Into<[u8; H::OUT_LEN]>,
364        Iter: IntoIterator<Item = Item, IntoIter: TrustedLen> + 'a,
365    {
366        let leaves = leaves.into_iter();
367        let hasher = Self::hasher_for_leaves(&leaves)?;
368
369        Self::compute_root_and_proof_with::<MAX_N, _, _>(&hasher, leaves, target_index)
370    }
371
372    /// Similar to [`Self::compute_root_and_proof()`], but for iterators that are not [`TrustedLen`]
373    #[expect(clippy::type_complexity, reason = "Return type")]
374    #[inline]
375    #[cfg(feature = "alloc")]
376    pub fn compute_root_and_proof_dynamic_len<'a, const MAX_N: u64, Item, Iter>(
377        leaves: Iter,
378        target_index: usize,
379    ) -> Option<([u8; H::OUT_LEN], Vec<[u8; H::OUT_LEN]>)>
380    where
381        H: Default,
382        Item: Into<[u8; H::OUT_LEN]>,
383        Iter: IntoIterator<Item = Item> + 'a,
384    {
385        let hasher = H::default();
386
387        Self::compute_root_and_proof_with::<MAX_N, _, _>(&hasher, leaves, target_index)
388    }
389
390    #[expect(clippy::type_complexity, reason = "Return type")]
391    #[inline(always)]
392    #[cfg(feature = "alloc")]
393    fn compute_root_and_proof_with<const MAX_N: u64, Item, Iter>(
394        hasher: &H,
395        leaves: Iter,
396        target_index: usize,
397    ) -> Option<([u8; H::OUT_LEN], Vec<[u8; H::OUT_LEN]>)>
398    where
399        Item: Into<[u8; H::OUT_LEN]>,
400        Iter: IntoIterator<Item = Item>,
401    {
402        // Stack of intermediate nodes per tree level
403        let mut stack = [[0u8; H::OUT_LEN]; _];
404        // SAFETY: Inner value is `MaybeUninit`
405        let mut proof =
406            unsafe { Box::<[MaybeUninit<[u8; H::OUT_LEN]>; _]>::new_uninit().assume_init() };
407
408        let (root, proof_length) = Self::compute_root_and_proof_inner::<MAX_N, _, _>(
409            hasher,
410            leaves,
411            target_index,
412            &mut stack,
413            &mut proof,
414        )?;
415
416        let proof_capacity = proof.len();
417        let proof = Box::into_raw(proof);
418        // SAFETY: Points to correctly allocated memory where `proof_length` elements were
419        // initialized
420        let proof = unsafe {
421            Vec::from_raw_parts(
422                proof.cast::<[u8; H::OUT_LEN]>(),
423                proof_length,
424                proof_capacity,
425            )
426        };
427
428        Some((root, proof))
429    }
430
431    /// Compute Merkle Tree root and generate a proof for the `leaf` at `target_index`.
432    ///
433    /// Returns `Some((root, proof))` on success, `None` if index is outside of list of leaves.
434    ///
435    /// `MAX_N` generic constant defines the maximum number of elements supported and controls stack
436    /// usage.
437    #[expect(clippy::type_complexity, reason = "Return type")]
438    #[inline]
439    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
440    pub fn compute_root_and_proof_in<'a, 'proof, const MAX_N: u64, Item, Iter>(
441        leaves: Iter,
442        target_index: usize,
443        proof: &'proof mut [MaybeUninit<[u8; H::OUT_LEN]>; PROOF_ELEMENTS::<MAX_N>],
444    ) -> Option<([u8; H::OUT_LEN], &'proof mut [[u8; H::OUT_LEN]])>
445    where
446        H: NewMerkleHasherForTreeSize,
447        Item: Into<[u8; H::OUT_LEN]>,
448        Iter: IntoIterator<Item = Item, IntoIter: TrustedLen> + 'a,
449    {
450        let leaves = leaves.into_iter();
451        let hasher = Self::hasher_for_leaves(&leaves)?;
452
453        Self::compute_root_and_proof_in_with::<MAX_N, _, _>(&hasher, leaves, target_index, proof)
454    }
455
456    /// Similar to [`Self::compute_root_and_proof_in()`], but for iterators that are not
457    /// [`TrustedLen`]
458    #[expect(clippy::type_complexity, reason = "Return type")]
459    #[inline]
460    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
461    pub fn compute_root_and_proof_dynamic_len_in<'a, 'proof, const MAX_N: u64, Item, Iter>(
462        leaves: Iter,
463        target_index: usize,
464        proof: &'proof mut [MaybeUninit<[u8; H::OUT_LEN]>; PROOF_ELEMENTS::<MAX_N>],
465    ) -> Option<([u8; H::OUT_LEN], &'proof mut [[u8; H::OUT_LEN]])>
466    where
467        H: Default,
468        Item: Into<[u8; H::OUT_LEN]>,
469        Iter: IntoIterator<Item = Item> + 'a,
470    {
471        let hasher = H::default();
472
473        Self::compute_root_and_proof_in_with::<MAX_N, _, _>(&hasher, leaves, target_index, proof)
474    }
475
476    #[expect(clippy::type_complexity, reason = "Return type")]
477    #[inline(always)]
478    fn compute_root_and_proof_in_with<'proof, const MAX_N: u64, Item, Iter>(
479        hasher: &H,
480        leaves: Iter,
481        target_index: usize,
482        proof: &'proof mut [MaybeUninit<[u8; H::OUT_LEN]>; PROOF_ELEMENTS::<MAX_N>],
483    ) -> Option<([u8; H::OUT_LEN], &'proof mut [[u8; H::OUT_LEN]])>
484    where
485        Item: Into<[u8; H::OUT_LEN]>,
486        Iter: IntoIterator<Item = Item>,
487    {
488        // Stack of intermediate nodes per tree level
489        let mut stack = [[0u8; H::OUT_LEN]; _];
490
491        let (root, proof_length) = Self::compute_root_and_proof_inner::<MAX_N, _, _>(
492            hasher,
493            leaves,
494            target_index,
495            &mut stack,
496            proof,
497        )?;
498        // SAFETY: Just correctly initialized `proof_length` elements
499        let proof = unsafe { proof.get_unchecked_mut(..proof_length).assume_init_mut() };
500
501        Some((root, proof))
502    }
503
504    #[inline]
505    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
506    fn compute_root_and_proof_inner<const MAX_N: u64, Item, Iter>(
507        hasher: &H,
508        leaves: Iter,
509        target_index: usize,
510        stack: &mut [[u8; H::OUT_LEN]; STACK_SIZE::<MAX_N>],
511        proof: &mut [MaybeUninit<[u8; H::OUT_LEN]>; PROOF_ELEMENTS::<MAX_N>],
512    ) -> Option<([u8; H::OUT_LEN], usize)>
513    where
514        Item: Into<[u8; H::OUT_LEN]>,
515        Iter: IntoIterator<Item = Item>,
516    {
517        let mut proof_length = 0;
518        let mut num_leaves = 0u64;
519
520        let mut current_target_level = None;
521        let mut position = target_index;
522
523        for (current_index, hash) in leaves.into_iter().enumerate() {
524            // How many leaves were processed so far. Should have been `num_leaves == MAX_N`, but
525            // `>=` helps compiler with panic safety checks.
526            // TODO: Use `==` once https://github.com/rust-lang/rust/issues/162834 is resolved
527            if num_leaves >= MAX_N {
528                return None;
529            }
530
531            let mut current = hash.into();
532
533            // Every bit set to `1` corresponds to an active Merkle Tree level
534            let lowest_active_levels = num_leaves.trailing_ones() as usize;
535
536            if current_index == target_index {
537                for item in stack.iter().take(lowest_active_levels) {
538                    // If at the target leaf index, need to collect the proof
539                    // SAFETY: Method signature guarantees upper bound of the proof length
540                    unsafe { proof.get_unchecked_mut(proof_length) }.write(*item);
541                    proof_length += 1;
542
543                    current = hasher.hash_pair(&[*item, current]);
544
545                    // Move up the tree
546                    position /= 2;
547                }
548
549                current_target_level = Some(lowest_active_levels);
550            } else {
551                for (level, item) in stack.iter().enumerate().take(lowest_active_levels) {
552                    if current_target_level == Some(level) {
553                        // SAFETY: Method signature guarantees upper bound of the proof length
554                        unsafe { proof.get_unchecked_mut(proof_length) }.write(
555                            if position.is_multiple_of(2) {
556                                current
557                            } else {
558                                *item
559                            },
560                        );
561                        proof_length += 1;
562
563                        current_target_level = Some(level + 1);
564
565                        // Move up the tree
566                        position /= 2;
567                    }
568
569                    current = hasher.hash_pair(&[*item, current]);
570                }
571            }
572
573            // Place the current hash at the first inactive level
574            stack[lowest_active_levels] = current;
575            num_leaves += 1;
576        }
577
578        // `active_levels` here contains the number of leaves after above loop
579        if target_index >= num_leaves as usize {
580            // If no leaves were provided
581            return None;
582        }
583
584        let Some(current_target_level) = current_target_level else {
585            // Index not found
586            return None;
587        };
588
589        let mut stack_bits = num_leaves;
590
591        {
592            let lowest_active_level = stack_bits.trailing_zeros() as usize;
593            // Reuse `stack[0]` for resulting value
594            // SAFETY: Active level must have been set successfully before, hence it exists
595            stack[0] = *unsafe { stack.get_unchecked(lowest_active_level) };
596            // Clear lowest active level
597            stack_bits &= !(1 << lowest_active_level);
598        }
599
600        // Hash remaining peaks (if any) of the potentially unbalanced tree together and collect
601        // proof hashes
602        let mut merged_peaks = false;
603        loop {
604            let lowest_active_level = stack_bits.trailing_zeros() as usize;
605
606            if lowest_active_level == u64::BITS as usize {
607                break;
608            }
609
610            // Clear lowest active level for next iteration
611            stack_bits &= !(1 << lowest_active_level);
612
613            // SAFETY: Active level must have been set successfully before, hence it exists
614            let lowest_active_level_item = unsafe { stack.get_unchecked(lowest_active_level) };
615
616            if lowest_active_level > current_target_level
617                || (lowest_active_level == current_target_level
618                    && !position.is_multiple_of(2)
619                    && !merged_peaks)
620            {
621                // SAFETY: Method signature guarantees upper bound of the proof length
622                unsafe { proof.get_unchecked_mut(proof_length) }.write(*lowest_active_level_item);
623                proof_length += 1;
624                merged_peaks = false;
625            } else if lowest_active_level == current_target_level {
626                // SAFETY: Method signature guarantees upper bound of the proof length
627                unsafe { proof.get_unchecked_mut(proof_length) }.write(stack[0]);
628                proof_length += 1;
629                merged_peaks = false;
630            } else {
631                // Not collecting proof because of the need to merge peaks of an unbalanced tree
632                merged_peaks = true;
633            }
634
635            // Collect the lowest peak into the proof
636            stack[0] = hasher.hash_pair(&[*lowest_active_level_item, stack[0]]);
637
638            position /= 2;
639        }
640
641        if num_leaves == 1 {
642            stack[0] = hasher.hash_single_leaf(&stack[0]);
643        }
644
645        Some((stack[0], proof_length))
646    }
647
648    /// Verify a Merkle proof for a leaf at the given index
649    #[inline]
650    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
651    pub fn verify(
652        root: &[u8; H::OUT_LEN],
653        proof: &[[u8; H::OUT_LEN]],
654        leaf_index: u64,
655        leaf: [u8; H::OUT_LEN],
656        num_leaves: u64,
657    ) -> bool
658    where
659        H: NewMerkleHasherForTreeSize,
660    {
661        if leaf_index >= num_leaves {
662            return false;
663        }
664
665        let hasher = H::new_for_tree_size(num_leaves);
666        let mut current = leaf;
667        let mut position = leaf_index;
668        let mut proof_pos = 0;
669        let mut level_size = num_leaves;
670
671        // Rebuild the path to the root
672        while level_size > 1 {
673            let is_left = position.is_multiple_of(2);
674            let is_last = position == level_size - 1;
675
676            if is_left && !is_last {
677                // Left node with a right sibling
678                if proof_pos >= proof.len() {
679                    // Missing sibling
680                    return false;
681                }
682                current = hasher.hash_pair(&[current, proof[proof_pos]]);
683                proof_pos += 1;
684            } else if !is_left {
685                // Right node with a left sibling
686                if proof_pos >= proof.len() {
687                    // Missing sibling
688                    return false;
689                }
690                current = hasher.hash_pair(&[proof[proof_pos], current]);
691                proof_pos += 1;
692            } else {
693                // Last node, no sibling, keep current
694            }
695
696            position /= 2;
697            // Size of next level
698            level_size = level_size.div_ceil(2);
699        }
700
701        if num_leaves == 1 {
702            current = hasher.hash_single_leaf(&current);
703        }
704
705        // Check if proof is fully used and matches root
706        proof_pos == proof.len() && current == *root
707    }
708
709    /// Create a hasher for a tree with `leaves`.
710    ///
711    /// Returns `None` if there are more than `usize::MAX` leaves.
712    #[inline(always)]
713    fn hasher_for_leaves<Iter>(leaves: &Iter) -> Option<H>
714    where
715        H: NewMerkleHasherForTreeSize,
716        Iter: TrustedLen,
717    {
718        // `TrustedLen` guarantees that the iterator yields exactly this many leaves. An empty list
719        // or more than `MAX_N` leaves are rejected while iterating, which avoids extra checks for
720        // hashers that ignore the number of leaves.
721        let (_lower, upper) = leaves.size_hint();
722
723        Some(H::new_for_tree_size(upper? as u64))
724    }
725}
ab_merkle_tree/unbalanced.rs

ab_merkle_tree/
unbalanced.rs
