miden-crypto/src/merkle/tiered_smt/mod.rs
src/merkle/tiered_smt/mod.rsRaw
1use super::{
2 BTreeMap, BTreeSet, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex,
3 Rpo256, RpoDigest, StarkField, Vec, Word,
4};
5use crate::utils::vec;
6use core::{cmp, ops::Deref};
7
8mod nodes;
9use nodes::NodeStore;
10
11mod values;
12use values::ValueStore;
13
14mod proof;
15pub use proof::TieredSmtProof;
16
17mod error;
18pub use error::TieredSmtProofError;
19
20#[cfg(test)]
21mod tests;
22
23// TIERED SPARSE MERKLE TREE
24// ================================================================================================
25
26/// Tiered (compacted) Sparse Merkle tree mapping 256-bit keys to 256-bit values. Both keys and
27/// values are represented by 4 field elements.
28///
29/// Leaves in the tree can exist only on specific depths called "tiers". These depths are: 16, 32,
30/// 48, and 64. Initially, when a tree is empty, it is equivalent to an empty Sparse Merkle tree
31/// of depth 64 (i.e., leaves at depth 64 are set to [ZERO; 4]). As non-empty values are inserted
32/// into the tree they are added to the first available tier.
33///
34/// For example, when the first key-value pair is inserted, it will be stored in a node at depth
35/// 16 such that the 16 most significant bits of the key determine the position of the node at
36/// depth 16. If another value with a key sharing the same 16-bit prefix is inserted, both values
37/// move into the next tier (depth 32). This process is repeated until values end up at the bottom
38/// tier (depth 64). If multiple values have keys with a common 64-bit prefix, such key-value pairs
39/// are stored in a sorted list at the bottom tier.
40///
41/// To differentiate between internal and leaf nodes, node values are computed as follows:
42/// - Internal nodes: hash(left_child, right_child).
43/// - Leaf node at depths 16, 32, or 64: hash(key, value, domain=depth).
44/// - Leaf node at depth 64: hash([key_0, value_0, ..., key_n, value_n], domain=64).
45#[derive(Debug, Clone, PartialEq, Eq)]
46#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
47pub struct TieredSmt {
48 root: RpoDigest,
49 nodes: NodeStore,
50 values: ValueStore,
51}
52
53impl TieredSmt {
54 // CONSTANTS
55 // --------------------------------------------------------------------------------------------
56
57 /// The number of levels between tiers.
58 pub const TIER_SIZE: u8 = 16;
59
60 /// Depths at which leaves can exist in a tiered SMT.
61 pub const TIER_DEPTHS: [u8; 4] = [16, 32, 48, 64];
62
63 /// Maximum node depth. This is also the bottom tier of the tree.
64 pub const MAX_DEPTH: u8 = 64;
65
66 /// Value of an empty leaf.
67 pub const EMPTY_VALUE: Word = super::empty_roots::EMPTY_WORD;
68
69 // CONSTRUCTORS
70 // --------------------------------------------------------------------------------------------
71
72 /// Returns a new [TieredSmt] instantiated with the specified key-value pairs.
73 ///
74 /// # Errors
75 /// Returns an error if the provided entries contain multiple values for the same key.
76 pub fn with_leaves<R, I>(entries: R) -> Result<Self, MerkleError>
77 where
78 R: IntoIterator<IntoIter = I>,
79 I: Iterator<Item = (RpoDigest, Word)> + ExactSizeIterator,
80 {
81 // create an empty tree
82 let mut tree = Self::default();
83
84 // append leaves to the tree returning an error if a duplicate entry for the same key
85 // is found
86 let mut empty_entries = BTreeSet::new();
87 for (key, value) in entries {
88 let old_value = tree.insert(key, value);
89 if old_value != Self::EMPTY_VALUE || empty_entries.contains(&key) {
90 return Err(MerkleError::DuplicateValuesForKey(key));
91 }
92 // if we've processed an empty entry, add the key to the set of empty entry keys, and
93 // if this key was already in the set, return an error
94 if value == Self::EMPTY_VALUE && !empty_entries.insert(key) {
95 return Err(MerkleError::DuplicateValuesForKey(key));
96 }
97 }
98 Ok(tree)
99 }
100
101 // PUBLIC ACCESSORS
102 // --------------------------------------------------------------------------------------------
103
104 /// Returns the root of this Merkle tree.
105 pub const fn root(&self) -> RpoDigest {
106 self.root
107 }
108
109 /// Returns a node at the specified index.
110 ///
111 /// # Errors
112 /// Returns an error if:
113 /// - The specified index depth is 0 or greater than 64.
114 /// - The node with the specified index does not exists in the Merkle tree. This is possible
115 /// when a leaf node with the same index prefix exists at a tier higher than the requested
116 /// node.
117 pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
118 self.nodes.get_node(index)
119 }
120
121 /// Returns a Merkle path from the node at the specified index to the root.
122 ///
123 /// The node itself is not included in the path.
124 ///
125 /// # Errors
126 /// Returns an error if:
127 /// - The specified index depth is 0 or greater than 64.
128 /// - The node with the specified index does not exists in the Merkle tree. This is possible
129 /// when a leaf node with the same index prefix exists at a tier higher than the node to
130 /// which the path is requested.
131 pub fn get_path(&self, index: NodeIndex) -> Result<MerklePath, MerkleError> {
132 self.nodes.get_path(index)
133 }
134
135 /// Returns the value associated with the specified key.
136 ///
137 /// If nothing was inserted into this tree for the specified key, [ZERO; 4] is returned.
138 pub fn get_value(&self, key: RpoDigest) -> Word {
139 match self.values.get(&key) {
140 Some(value) => *value,
141 None => Self::EMPTY_VALUE,
142 }
143 }
144
145 /// Returns a proof for a key-value pair defined by the specified key.
146 ///
147 /// The proof can be used to attest membership of this key-value pair in a Tiered Sparse Merkle
148 /// Tree defined by the same root as this tree.
149 pub fn prove(&self, key: RpoDigest) -> TieredSmtProof {
150 let (path, index, leaf_exists) = self.nodes.get_proof(&key);
151
152 let entries = if index.depth() == Self::MAX_DEPTH {
153 match self.values.get_all(index.value()) {
154 Some(entries) => entries,
155 None => vec![(key, Self::EMPTY_VALUE)],
156 }
157 } else if leaf_exists {
158 let entry =
159 self.values.get_first(index_to_prefix(&index)).expect("leaf entry not found");
160 debug_assert_eq!(entry.0, key);
161 vec![*entry]
162 } else {
163 vec![(key, Self::EMPTY_VALUE)]
164 };
165
166 TieredSmtProof::new(path, entries).expect("Bug detected, TSMT produced invalid proof")
167 }
168
169 // STATE MUTATORS
170 // --------------------------------------------------------------------------------------------
171
172 /// Inserts the provided value into the tree under the specified key and returns the value
173 /// previously stored under this key.
174 ///
175 /// If the value for the specified key was not previously set, [ZERO; 4] is returned.
176 pub fn insert(&mut self, key: RpoDigest, value: Word) -> Word {
177 // if an empty value is being inserted, remove the leaf node to make it look as if the
178 // value was never inserted
179 if value == Self::EMPTY_VALUE {
180 return self.remove_leaf_node(key);
181 }
182
183 // insert the value into the value store, and if the key was already in the store, update
184 // it with the new value
185 if let Some(old_value) = self.values.insert(key, value) {
186 if old_value != value {
187 // if the new value is different from the old value, determine the location of
188 // the leaf node for this key, build the node, and update the root
189 let (index, leaf_exists) = self.nodes.get_leaf_index(&key);
190 debug_assert!(leaf_exists);
191 let node = self.build_leaf_node(index, key, value);
192 self.root = self.nodes.update_leaf_node(index, node);
193 }
194 return old_value;
195 };
196
197 // determine the location for the leaf node; this index could have 3 different meanings:
198 // - it points to a root of an empty subtree or an empty node at depth 64; in this case,
199 // we can replace the node with the value node immediately.
200 // - it points to an existing leaf at the bottom tier (i.e., depth = 64); in this case,
201 // we need to process update the bottom leaf.
202 // - it points to an existing leaf node for a different key with the same prefix (same
203 // key case was handled above); in this case, we need to move the leaf to a lower tier
204 let (index, leaf_exists) = self.nodes.get_leaf_index(&key);
205
206 self.root = if leaf_exists && index.depth() == Self::MAX_DEPTH {
207 // returned index points to a leaf at the bottom tier
208 let node = self.build_leaf_node(index, key, value);
209 self.nodes.update_leaf_node(index, node)
210 } else if leaf_exists {
211 // returned index pointes to a leaf for a different key with the same prefix
212
213 // get the key-value pair for the key with the same prefix; since the key-value
214 // pair has already been inserted into the value store, we need to filter it out
215 // when looking for the other key-value pair
216 let (other_key, other_value) = self
217 .values
218 .get_first_filtered(index_to_prefix(&index), &key)
219 .expect("other key-value pair not found");
220
221 // determine how far down the tree should we move the leaves
222 let common_prefix_len = get_common_prefix_tier_depth(&key, other_key);
223 let depth = cmp::min(common_prefix_len + Self::TIER_SIZE, Self::MAX_DEPTH);
224
225 // compute node locations for new and existing key-value paris
226 let new_index = LeafNodeIndex::from_key(&key, depth);
227 let other_index = LeafNodeIndex::from_key(other_key, depth);
228
229 // compute node values for the new and existing key-value pairs
230 let new_node = self.build_leaf_node(new_index, key, value);
231 let other_node = self.build_leaf_node(other_index, *other_key, *other_value);
232
233 // replace the leaf located at index with a subtree containing nodes for new and
234 // existing key-value paris
235 self.nodes.replace_leaf_with_subtree(
236 index,
237 [(new_index, new_node), (other_index, other_node)],
238 )
239 } else {
240 // returned index points to an empty subtree or an empty leaf at the bottom tier
241 let node = self.build_leaf_node(index, key, value);
242 self.nodes.insert_leaf_node(index, node)
243 };
244
245 Self::EMPTY_VALUE
246 }
247
248 // ITERATORS
249 // --------------------------------------------------------------------------------------------
250
251 /// Returns an iterator over all key-value pairs in this [TieredSmt].
252 pub fn iter(&self) -> impl Iterator<Item = &(RpoDigest, Word)> {
253 self.values.iter()
254 }
255
256 /// Returns an iterator over all inner nodes of this [TieredSmt] (i.e., nodes not at depths 16
257 /// 32, 48, or 64).
258 ///
259 /// The iterator order is unspecified.
260 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
261 self.nodes.inner_nodes()
262 }
263
264 /// Returns an iterator over upper leaves (i.e., depth = 16, 32, or 48) for this [TieredSmt]
265 /// where each yielded item is a (node, key, value) tuple.
266 ///
267 /// The iterator order is unspecified.
268 pub fn upper_leaves(&self) -> impl Iterator<Item = (RpoDigest, RpoDigest, Word)> + '_ {
269 self.nodes.upper_leaves().map(|(index, node)| {
270 let key_prefix = index_to_prefix(index);
271 let (key, value) = self.values.get_first(key_prefix).expect("upper leaf not found");
272 debug_assert_eq!(*index, LeafNodeIndex::from_key(key, index.depth()).into());
273 (*node, *key, *value)
274 })
275 }
276
277 /// Returns an iterator over upper leaves (i.e., depth = 16, 32, or 48) for this [TieredSmt]
278 /// where each yielded item is a (node_index, value) tuple.
279 pub fn upper_leaf_nodes(&self) -> impl Iterator<Item = (&NodeIndex, &RpoDigest)> {
280 self.nodes.upper_leaves()
281 }
282
283 /// Returns an iterator over bottom leaves (i.e., depth = 64) of this [TieredSmt].
284 ///
285 /// Each yielded item consists of the hash of the leaf and its contents, where contents is
286 /// a vector containing key-value pairs of entries storied in this leaf.
287 ///
288 /// The iterator order is unspecified.
289 pub fn bottom_leaves(&self) -> impl Iterator<Item = (RpoDigest, Vec<(RpoDigest, Word)>)> + '_ {
290 self.nodes.bottom_leaves().map(|(&prefix, node)| {
291 let values = self.values.get_all(prefix).expect("bottom leaf not found");
292 (*node, values)
293 })
294 }
295
296 // HELPER METHODS
297 // --------------------------------------------------------------------------------------------
298
299 /// Removes the node holding the key-value pair for the specified key from this tree, and
300 /// returns the value associated with the specified key.
301 ///
302 /// If no value was associated with the specified key, [ZERO; 4] is returned.
303 fn remove_leaf_node(&mut self, key: RpoDigest) -> Word {
304 // remove the key-value pair from the value store; if no value was associated with the
305 // specified key, return.
306 let old_value = match self.values.remove(&key) {
307 Some(old_value) => old_value,
308 None => return Self::EMPTY_VALUE,
309 };
310
311 // determine the location of the leaf holding the key-value pair to be removed
312 let (index, leaf_exists) = self.nodes.get_leaf_index(&key);
313 debug_assert!(leaf_exists);
314
315 // if the leaf is at the bottom tier and after removing the key-value pair from it, the
316 // leaf is still not empty, we either just update it, or move it up to a higher tier (if
317 // the leaf doesn't have siblings at lower tiers)
318 if index.depth() == Self::MAX_DEPTH {
319 if let Some(entries) = self.values.get_all(index.value()) {
320 // if there is only one key-value pair left at the bottom leaf, and it can be
321 // moved up to a higher tier, truncate the branch and return
322 if entries.len() == 1 {
323 let new_depth = self.nodes.get_last_single_child_parent_depth(index.value());
324 if new_depth != Self::MAX_DEPTH {
325 let node = hash_upper_leaf(entries[0].0, entries[0].1, new_depth);
326 self.root = self.nodes.truncate_branch(index.value(), new_depth, node);
327 return old_value;
328 }
329 }
330
331 // otherwise just recompute the leaf hash and update the leaf node
332 let node = hash_bottom_leaf(&entries);
333 self.root = self.nodes.update_leaf_node(index, node);
334 return old_value;
335 };
336 }
337
338 // if the removed key-value pair has a lone sibling at the current tier with a root at
339 // higher tier, we need to move the sibling to a higher tier
340 if let Some((sib_key, sib_val, new_sib_index)) = self.values.get_lone_sibling(index) {
341 // determine the current index of the sibling node
342 let sib_index = LeafNodeIndex::from_key(sib_key, index.depth());
343 debug_assert!(sib_index.depth() > new_sib_index.depth());
344
345 // compute node value for the new location of the sibling leaf and replace the subtree
346 // with this leaf node
347 let node = self.build_leaf_node(new_sib_index, *sib_key, *sib_val);
348 let new_sib_depth = new_sib_index.depth();
349 self.root = self.nodes.replace_subtree_with_leaf(index, sib_index, new_sib_depth, node);
350 } else {
351 // if the removed key-value pair did not have a sibling at the current tier with a
352 // root at higher tiers, just clear the leaf node
353 self.root = self.nodes.clear_leaf_node(index);
354 }
355
356 old_value
357 }
358
359 /// Builds and returns a leaf node value for the node located as the specified index.
360 ///
361 /// This method assumes that the key-value pair for the node has already been inserted into
362 /// the value store, however, for depths 16, 32, and 48, the node is computed directly from
363 /// the passed-in values (for depth 64, the value store is queried to get all the key-value
364 /// pairs located at the specified index).
365 fn build_leaf_node(&self, index: LeafNodeIndex, key: RpoDigest, value: Word) -> RpoDigest {
366 let depth = index.depth();
367
368 // insert the key into index-key map and compute the new value of the node
369 if index.depth() == Self::MAX_DEPTH {
370 // for the bottom tier, we add the key-value pair to the existing leaf, or create a
371 // new leaf with this key-value pair
372 let values = self.values.get_all(index.value()).unwrap();
373 hash_bottom_leaf(&values)
374 } else {
375 debug_assert_eq!(self.values.get_first(index_to_prefix(&index)), Some(&(key, value)));
376 hash_upper_leaf(key, value, depth)
377 }
378 }
379}
380
381impl Default for TieredSmt {
382 fn default() -> Self {
383 let root = EmptySubtreeRoots::empty_hashes(Self::MAX_DEPTH)[0];
384 Self {
385 root,
386 nodes: NodeStore::new(root),
387 values: ValueStore::default(),
388 }
389 }
390}
391
392// LEAF NODE INDEX
393// ================================================================================================
394/// A wrapper around [NodeIndex] to provide type-safe references to nodes at depths 16, 32, 48, and
395/// 64.
396#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
397pub struct LeafNodeIndex(NodeIndex);
398
399impl LeafNodeIndex {
400 /// Returns a new [LeafNodeIndex] instantiated from the provided [NodeIndex].
401 ///
402 /// In debug mode, panics if index depth is not 16, 32, 48, or 64.
403 pub fn new(index: NodeIndex) -> Self {
404 // check if the depth is 16, 32, 48, or 64; this works because for a valid depth,
405 // depth - 16, can be 0, 16, 32, or 48 - i.e., the value is either 0 or any of the 4th
406 // or 5th bits are set. We can test for this by computing a bitwise AND with a value
407 // which has all but the 4th and 5th bits set (which is !48).
408 debug_assert_eq!(((index.depth() - 16) & !48), 0, "invalid tier depth {}", index.depth());
409 Self(index)
410 }
411
412 /// Returns a new [LeafNodeIndex] instantiated from the specified key inserted at the specified
413 /// depth.
414 ///
415 /// The value for the key is computed by taking n most significant bits from the most significant
416 /// element of the key, where n is the specified depth.
417 pub fn from_key(key: &RpoDigest, depth: u8) -> Self {
418 let mse = get_key_prefix(key);
419 Self::new(NodeIndex::new_unchecked(depth, mse >> (TieredSmt::MAX_DEPTH - depth)))
420 }
421
422 /// Returns a new [LeafNodeIndex] instantiated for testing purposes.
423 #[cfg(test)]
424 pub fn make(depth: u8, value: u64) -> Self {
425 Self::new(NodeIndex::make(depth, value))
426 }
427
428 /// Traverses towards the root until the specified depth is reached.
429 ///
430 /// The new depth must be a valid tier depth - i.e., 16, 32, 48, or 64.
431 pub fn move_up_to(&mut self, depth: u8) {
432 debug_assert_eq!(((depth - 16) & !48), 0, "invalid tier depth: {depth}");
433 self.0.move_up_to(depth);
434 }
435}
436
437impl Deref for LeafNodeIndex {
438 type Target = NodeIndex;
439
440 fn deref(&self) -> &Self::Target {
441 &self.0
442 }
443}
444
445impl From<NodeIndex> for LeafNodeIndex {
446 fn from(value: NodeIndex) -> Self {
447 Self::new(value)
448 }
449}
450
451impl From<LeafNodeIndex> for NodeIndex {
452 fn from(value: LeafNodeIndex) -> Self {
453 value.0
454 }
455}
456
457// HELPER FUNCTIONS
458// ================================================================================================
459
460/// Returns the value representing the 64 most significant bits of the specified key.
461fn get_key_prefix(key: &RpoDigest) -> u64 {
462 Word::from(key)[3].as_int()
463}
464
465/// Returns the index value shifted to be in the most significant bit positions of the returned
466/// u64 value.
467fn index_to_prefix(index: &NodeIndex) -> u64 {
468 index.value() << (TieredSmt::MAX_DEPTH - index.depth())
469}
470
471/// Returns tiered common prefix length between the most significant elements of the provided keys.
472///
473/// Specifically:
474/// - returns 64 if the most significant elements are equal.
475/// - returns 48 if the common prefix is between 48 and 63 bits.
476/// - returns 32 if the common prefix is between 32 and 47 bits.
477/// - returns 16 if the common prefix is between 16 and 31 bits.
478/// - returns 0 if the common prefix is fewer than 16 bits.
479fn get_common_prefix_tier_depth(key1: &RpoDigest, key2: &RpoDigest) -> u8 {
480 let e1 = get_key_prefix(key1);
481 let e2 = get_key_prefix(key2);
482 let ex = (e1 ^ e2).leading_zeros() as u8;
483 (ex / 16) * 16
484}
485
486/// Computes node value for leaves at tiers 16, 32, or 48.
487///
488/// Node value is computed as: hash(key || value, domain = depth).
489pub fn hash_upper_leaf(key: RpoDigest, value: Word, depth: u8) -> RpoDigest {
490 const NUM_UPPER_TIERS: usize = TieredSmt::TIER_DEPTHS.len() - 1;
491 debug_assert!(TieredSmt::TIER_DEPTHS[..NUM_UPPER_TIERS].contains(&depth));
492 Rpo256::merge_in_domain(&[key, value.into()], depth.into())
493}
494
495/// Computes node value for leaves at the bottom tier (depth 64).
496///
497/// Node value is computed as: hash([key_0, value_0, ..., key_n, value_n], domain=64).
498///
499/// TODO: when hashing in domain is implemented for `hash_elements()`, combine this function with
500/// `hash_upper_leaf()` function.
501pub fn hash_bottom_leaf(values: &[(RpoDigest, Word)]) -> RpoDigest {
502 let mut elements = Vec::with_capacity(values.len() * 8);
503 for (key, val) in values.iter() {
504 elements.extend_from_slice(key.as_elements());
505 elements.extend_from_slice(val.as_slice());
506 }
507 // TODO: hash in domain
508 Rpo256::hash_elements(&elements)
509}