ab_merkle_tree/sparse.rs
1//! Sparse Merkle Tree and related data structures.
2//!
3//! Sparse Merkle Tree is essentially a huge Balanced Merkle Tree, where most of the leaves are
4//! empty. By "empty" here we mean `[0u8; 32]`. To optimize proofs and their verification, the
5//! hashing function is customized and returns `[0u8; 32]` when both left and right branch are
6//! `[0u8; 32]`, otherwise BLAKE3 hash is used like in a Balanced Merkle Tree.
7
8use crate::{OUT_LEN, hash_pair};
9use core::num::NonZeroU128;
10
11/// Number of elements in a proof for a tree with `BITS`
12pub const PROOF_ELEMENTS<const BITS: u8>: usize = usize::from(BITS);
13const STACK_SIZE<const BITS: u8>: usize = usize::from(BITS) + 1;
14
15/// Ensuring only supported `NUM_BITS` can be specified for [`SparseMerkleTree`].
16///
17/// This is essentially a workaround for the current Rust type system constraints that do not allow
18/// a nicer way to do the same thing at compile time.
19const ENSURE_SUPPORTED_BITS<const BITS: u8>: usize = {
20 assert!(
21 BITS <= 128,
22 "This Sparse Merkle Tree doesn't support more than 2^128 leaves"
23 );
24
25 assert!(
26 BITS != 0,
27 "This Sparse Merkle Tree must have more than one leaf"
28 );
29
30 0
31};
32
33/// Sparse Merkle Tree Leaf
34#[derive(Debug)]
35pub enum Leaf<'a> {
36 // TODO: Batch of leaves for efficiently, especially with SIMD?
37 /// Leaf contains a value
38 Occupied {
39 /// Leaf value
40 leaf: &'a [u8; OUT_LEN],
41 },
42 /// Leaf contains a value (owned)
43 OccupiedOwned {
44 /// Leaf value
45 leaf: [u8; OUT_LEN],
46 },
47 /// Leaf is empty
48 Empty {
49 /// Number of consecutive empty leaves
50 skip_count: NonZeroU128,
51 },
52}
53
54impl<'a> From<&'a [u8; OUT_LEN]> for Leaf<'a> {
55 #[inline(always)]
56 fn from(leaf: &'a [u8; OUT_LEN]) -> Self {
57 Self::Occupied { leaf }
58 }
59}
60
61// TODO: A version that can hold intermediate nodes in memory, efficiently update leaves, etc.
62/// Sparse Merkle Tree variant that has hash-sized leaves, with most leaves being empty
63/// (have value `[0u8; 32]`).
64///
65/// In contrast to a proper Balanced Merkle Tree, constant `BITS` here specifies the max number of
66/// leaves hypothetically possible in a tree (2^BITS, often intractable), rather than the number of
67/// non-empty leaves actually present.
68#[derive(Debug)]
69pub struct SparseMerkleTree<const BITS: u8>;
70
71// TODO: Optimize by implementing SIMD-accelerated hashing of multiple values:
72// https://github.com/BLAKE3-team/BLAKE3/issues/478
73impl<const BITS: u8> SparseMerkleTree<BITS>
74where
75 [(); ENSURE_SUPPORTED_BITS::<BITS>]:,
76{
77 // TODO: Method that generates not only root, but also proof, like Unbalanced Merkle Tree
78 /// Compute Merkle Tree root.
79 ///
80 /// If provided iterator ends early, it means the rest of the leaves are empty.
81 ///
82 /// There must be no [`Leaf::Occupied`] for empty/unoccupied leaves or else they may result in
83 /// invalid root, [`Leaf::Empty`] must be used instead.
84 ///
85 /// Returns `None` if too many leaves were provided.
86 #[inline]
87 #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
88 pub fn compute_root_only<'a, Iter>(leaves: Iter) -> Option<[u8; OUT_LEN]>
89 where
90 Iter: IntoIterator<Item = Leaf<'a>> + 'a,
91 {
92 // TODO: Technically redundant with the where-clause, which is not enforced in generic code
93 // due to https://github.com/rust-lang/rust/issues/164019, remove once fixed
94 const { ENSURE_SUPPORTED_BITS::<BITS> };
95
96 // Stack of intermediate nodes per tree level
97 let mut stack = [[0u8; OUT_LEN]; _];
98 let mut processed_some = false;
99 let mut num_leaves = 0u128;
100
101 for leaf in leaves {
102 if u32::from(BITS) < u128::BITS {
103 // How many leaves were processed so far. Should have been
104 // `num_leaves == 2^BITS`, but `>=` helps compiler with panic safety checks.
105 // TODO: Use `==` once https://github.com/rust-lang/rust/issues/162834 is resolved
106 if num_leaves >= 2u128.pow(u32::from(BITS)) {
107 return None;
108 }
109 } else {
110 // For `BITS == u128::BITS` `num_leaves` will wrap around back to zero right at the
111 // very end
112 if processed_some && num_leaves == 0 {
113 return None;
114 }
115 processed_some = true;
116 }
117
118 let leaf = match leaf {
119 Leaf::Occupied { leaf } => *leaf,
120 Leaf::OccupiedOwned { leaf } => leaf,
121 Leaf::Empty { skip_count } => {
122 num_leaves = Self::skip_leaves(
123 &mut stack,
124 &mut processed_some,
125 num_leaves,
126 skip_count.get(),
127 )?;
128 continue;
129 }
130 };
131
132 let mut current = leaf;
133
134 // Every bit set to `1` corresponds to an active Merkle Tree level
135 let lowest_active_levels = num_leaves.trailing_ones() as usize;
136 for item in stack.iter().take(lowest_active_levels) {
137 current = hash_pair(item, ¤t);
138 }
139
140 // Place the current hash at the first inactive level
141 *stack.get_mut(lowest_active_levels).expect(
142 "Stack has `BITS + 1` levels, enough for `num_leaves < 2^BITS` checked above; qed",
143 ) = current;
144 // Wrapping is needed for `BITS == u128::BITS`, where number of leaves narrowly
145 // doesn't fit into `u128` itself
146 num_leaves = num_leaves.wrapping_add(1);
147 }
148
149 if u32::from(BITS) < u128::BITS {
150 Self::skip_leaves(
151 &mut stack,
152 &mut processed_some,
153 num_leaves,
154 2u128.pow(u32::from(BITS)) - num_leaves,
155 )?;
156 } else if processed_some && num_leaves != 0 {
157 // For `BITS == u128::BITS` `num_leaves` will wrap around back to zero right at the
158 // very end, so we reverse the mechanism here
159 Self::skip_leaves(
160 &mut stack,
161 &mut processed_some,
162 num_leaves,
163 0u128.wrapping_sub(num_leaves),
164 )?;
165 }
166
167 Some(stack[BITS as usize])
168 }
169
170 /// Returns updated number of leaves
171 #[inline]
172 #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
173 fn skip_leaves(
174 stack: &mut [[u8; OUT_LEN]; STACK_SIZE::<BITS>],
175 processed_some: &mut bool,
176 mut num_leaves: u128,
177 mut skip_count: u128,
178 ) -> Option<u128> {
179 const ZERO: [u8; OUT_LEN] = [0; OUT_LEN];
180
181 if u32::from(BITS) < u128::BITS {
182 // How many leaves were processed so far
183 if num_leaves.checked_add(skip_count)? > 2u128.pow(u32::from(BITS)) {
184 return None;
185 }
186 } else {
187 // For `BITS == u128::BITS` `num_leaves` will wrap around back to zero right at the
188 // very end
189 let (overflow_amount, overflowed) = num_leaves.overflowing_add(skip_count);
190 if *processed_some && overflowed && overflow_amount > 0 {
191 return None;
192 }
193 *processed_some = true;
194 }
195
196 while skip_count > 0 {
197 // Find the largest aligned chunk to skip for the current state of the tree
198 let max_levels_to_skip = skip_count.ilog2().min(num_leaves.trailing_zeros());
199 let chunk_size = 1u128 << max_levels_to_skip;
200
201 let mut level = max_levels_to_skip;
202 let mut current = ZERO;
203 for item in stack.iter().skip(max_levels_to_skip as usize) {
204 // Check the active level for merging up the stack.
205 //
206 // `BITS == u128::BITS` condition is only added for better dead code elimination
207 // since that check is only relevant for 2^128 leaves case and nothing else.
208 if (u32::from(BITS) == u128::BITS && level == u128::BITS)
209 || num_leaves & (1 << level) == 0
210 {
211 // Level wasn't active before, stop here
212 break;
213 }
214
215 // Hash together unless both are zero
216 if !(item == &ZERO && current == ZERO) {
217 current = hash_pair(item, ¤t);
218 }
219
220 level += 1;
221 }
222 // SAFETY: Level is limited by the number of leaves, which in turn is checked above to
223 // fit into 2^BITS, while `BITS` generic in turn ensured sufficient stack size
224 *unsafe { stack.get_unchecked_mut(level as usize) } = current;
225
226 // Wrapping is needed for `BITS == u128::BITS`, where number of leaves narrowly
227 // doesn't fit into `u128` itself
228 num_leaves = num_leaves.wrapping_add(chunk_size);
229 skip_count -= chunk_size;
230 }
231
232 Some(num_leaves)
233 }
234
235 /// Verify previously generated proof.
236 ///
237 /// Leaf can either be leaf value for a leaf that is occupied or `[0; 32]` for a leaf that is
238 /// supposed to be empty.
239 #[inline]
240 #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
241 pub fn verify(
242 root: &[u8; OUT_LEN],
243 proof: &[[u8; OUT_LEN]; PROOF_ELEMENTS::<BITS>],
244 leaf_index: u128,
245 leaf: [u8; OUT_LEN],
246 ) -> bool {
247 // TODO: Technically redundant with the where-clause, which is not enforced in generic code
248 // due to https://github.com/rust-lang/rust/issues/164019, remove once fixed
249 const { ENSURE_SUPPORTED_BITS::<BITS> };
250
251 // For `BITS == u128::BITS` any index is valid by definition
252 if u32::from(BITS) < u128::BITS && leaf_index >= 2u128.pow(u32::from(BITS)) {
253 return false;
254 }
255
256 let mut computed_root = leaf;
257
258 let mut position = leaf_index;
259 for hash in proof {
260 computed_root = if position.is_multiple_of(2) {
261 hash_pair(&computed_root, hash)
262 } else {
263 hash_pair(hash, &computed_root)
264 };
265
266 position /= 2;
267 }
268
269 root == &computed_root
270 }
271}