miden-crypto/src/merkle/tiered_smt/mod.rs
src/merkle/tiered_smt/mod.rsRaw
1use super::{
2 BTreeMap, BTreeSet, EmptySubtreeRoots, Felt, InnerNodeInfo, MerkleError, MerklePath, NodeIndex,
3 Rpo256, RpoDigest, StarkField, Vec, Word,
4};
5use core::cmp;
6
7#[cfg(test)]
8mod tests;
9
10// TIERED SPARSE MERKLE TREE
11// ================================================================================================
12
13/// Tiered (compacted) Sparse Merkle tree mapping 256-bit keys to 256-bit values. Both keys and
14/// values are represented by 4 field elements.
15///
16/// Leaves in the tree can exist only on specific depths called "tiers". These depths are: 16, 32,
17/// 48, and 64. Initially, when a tree is empty, it is equivalent to an empty Sparse Merkle tree
18/// of depth 64 (i.e., leaves at depth 64 are set to [ZERO; 4]). As non-empty values are inserted
19/// into the tree they are added to the first available tier.
20///
21/// For example, when the first key-value is inserted, it will be stored in a node at depth 16
22/// such that the first 16 bits of the key determine the position of the node at depth 16. If
23/// another value with a key sharing the same 16-bit prefix is inserted, both values move into
24/// the next tier (depth 32). This process is repeated until values end up at tier 64. If multiple
25/// values have keys with a common 64-bit prefix, such key-value pairs are stored in a sorted list
26/// at the last tier (depth = 64).
27///
28/// To differentiate between internal and leaf nodes, node values are computed as follows:
29/// - Internal nodes: hash(left_child, right_child).
30/// - Leaf node at depths 16, 32, or 64: hash(key, value, domain=depth).
31/// - Leaf node at depth 64: hash([key_0, value_0, ..., key_n, value_n, domain=64]).
32///
33/// Where rem_key is computed by replacing d most significant bits of the key with zeros where d
34/// is depth (i.e., for a leaf at depth 16, we replace 16 most significant bits of the key with 0).
35#[derive(Debug, Clone, PartialEq, Eq)]
36pub struct TieredSmt {
37 root: RpoDigest,
38 nodes: BTreeMap<NodeIndex, RpoDigest>,
39 upper_leaves: BTreeMap<NodeIndex, RpoDigest>, // node_index |-> key
40 bottom_leaves: BTreeMap<u64, BottomLeaf>, // leaves of depth 64
41 values: BTreeMap<RpoDigest, Word>,
42}
43
44impl TieredSmt {
45 // CONSTANTS
46 // --------------------------------------------------------------------------------------------
47
48 /// The number of levels between tiers.
49 const TIER_SIZE: u8 = 16;
50
51 /// Depths at which leaves can exist in a tiered SMT.
52 const TIER_DEPTHS: [u8; 4] = [16, 32, 48, 64];
53
54 /// Maximum node depth. This is also the bottom tier of the tree.
55 const MAX_DEPTH: u8 = 64;
56
57 /// Value of an empty leaf.
58 pub const EMPTY_VALUE: Word = super::empty_roots::EMPTY_WORD;
59
60 // CONSTRUCTORS
61 // --------------------------------------------------------------------------------------------
62
63 /// Returns a new [TieredSmt] instantiated with the specified key-value pairs.
64 ///
65 /// # Errors
66 /// Returns an error if the provided entries contain multiple values for the same key.
67 pub fn with_leaves<R, I>(entries: R) -> Result<Self, MerkleError>
68 where
69 R: IntoIterator<IntoIter = I>,
70 I: Iterator<Item = (RpoDigest, Word)> + ExactSizeIterator,
71 {
72 // create an empty tree
73 let mut tree = Self::default();
74
75 // append leaves to the tree returning an error if a duplicate entry for the same key
76 // is found
77 let mut empty_entries = BTreeSet::new();
78 for (key, value) in entries {
79 let old_value = tree.insert(key, value);
80 if old_value != Self::EMPTY_VALUE || empty_entries.contains(&key) {
81 return Err(MerkleError::DuplicateValuesForKey(key));
82 }
83 // if we've processed an empty entry, add the key to the set of empty entry keys, and
84 // if this key was already in the set, return an error
85 if value == Self::EMPTY_VALUE && !empty_entries.insert(key) {
86 return Err(MerkleError::DuplicateValuesForKey(key));
87 }
88 }
89 Ok(tree)
90 }
91
92 // PUBLIC ACCESSORS
93 // --------------------------------------------------------------------------------------------
94
95 /// Returns the root of this Merkle tree.
96 pub const fn root(&self) -> RpoDigest {
97 self.root
98 }
99
100 /// Returns a node at the specified index.
101 ///
102 /// # Errors
103 /// Returns an error if:
104 /// - The specified index depth is 0 or greater than 64.
105 /// - The node with the specified index does not exists in the Merkle tree. This is possible
106 /// when a leaf node with the same index prefix exists at a tier higher than the requested
107 /// node.
108 pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
109 self.validate_node_access(index)?;
110 Ok(self.get_node_unchecked(&index))
111 }
112
113 /// Returns a Merkle path from the node at the specified index to the root.
114 ///
115 /// The node itself is not included in the path.
116 ///
117 /// # Errors
118 /// Returns an error if:
119 /// - The specified index depth is 0 or greater than 64.
120 /// - The node with the specified index does not exists in the Merkle tree. This is possible
121 /// when a leaf node with the same index prefix exists at a tier higher than the node to
122 /// which the path is requested.
123 pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
124 self.validate_node_access(index)?;
125
126 let mut path = Vec::with_capacity(index.depth() as usize);
127 for _ in 0..index.depth() {
128 let node = self.get_node_unchecked(&index.sibling());
129 path.push(node);
130 index.move_up();
131 }
132
133 Ok(path.into())
134 }
135
136 /// Returns the value associated with the specified key.
137 ///
138 /// If nothing was inserted into this tree for the specified key, [ZERO; 4] is returned.
139 pub fn get_value(&self, key: RpoDigest) -> Word {
140 match self.values.get(&key) {
141 Some(value) => *value,
142 None => Self::EMPTY_VALUE,
143 }
144 }
145
146 // STATE MUTATORS
147 // --------------------------------------------------------------------------------------------
148
149 /// Inserts the provided value into the tree under the specified key and returns the value
150 /// previously stored under this key.
151 ///
152 /// If the value for the specified key was not previously set, [ZERO; 4] is returned.
153 pub fn insert(&mut self, key: RpoDigest, value: Word) -> Word {
154 // insert the value into the key-value map, and if nothing has changed, return
155 let old_value = self.values.insert(key, value).unwrap_or(Self::EMPTY_VALUE);
156 if old_value == value {
157 return old_value;
158 }
159
160 // determine the index for the value node; this index could have 3 different meanings:
161 // - it points to a root of an empty subtree (excluding depth = 64); in this case, we can
162 // replace the node with the value node immediately.
163 // - it points to a node at the bottom tier (i.e., depth = 64); in this case, we need to
164 // process bottom-tier insertion which will be handled by insert_node().
165 // - it points to a leaf node; this node could be a node with the same key or a different
166 // key with a common prefix; in the latter case, we'll need to move the leaf to a lower
167 // tier; for this scenario the `leaf_key` will contain the key of the leaf node
168 let (mut index, leaf_key) = self.get_insert_location(&key);
169
170 // if the returned index points to a leaf, and this leaf is for a different key, we need
171 // to move the leaf to a lower tier
172 if let Some(other_key) = leaf_key {
173 if other_key != key {
174 // determine how far down the tree should we move the existing leaf
175 let common_prefix_len = get_common_prefix_tier(&key, &other_key);
176 let depth = cmp::min(common_prefix_len + Self::TIER_SIZE, Self::MAX_DEPTH);
177
178 // move the leaf to the new location; this requires first removing the existing
179 // index, re-computing node value, and inserting the node at a new location
180 let other_index = key_to_index(&other_key, depth);
181 let other_value = *self.values.get(&other_key).expect("no value for other key");
182 self.upper_leaves.remove(&index).expect("other node key not in map");
183 self.insert_leaf_node(other_index, other_key, other_value);
184
185 // the new leaf also needs to move down to the same tier
186 index = key_to_index(&key, depth);
187 }
188 }
189
190 // insert the node and return the old value
191 self.insert_leaf_node(index, key, value);
192 old_value
193 }
194
195 // ITERATORS
196 // --------------------------------------------------------------------------------------------
197
198 /// Returns an iterator over all inner nodes of this [TieredSmt] (i.e., nodes not at depths 16
199 /// 32, 48, or 64).
200 ///
201 /// The iterator order is unspecified.
202 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
203 self.nodes.iter().filter_map(|(index, node)| {
204 if is_inner_node(index) {
205 Some(InnerNodeInfo {
206 value: *node,
207 left: self.get_node_unchecked(&index.left_child()),
208 right: self.get_node_unchecked(&index.right_child()),
209 })
210 } else {
211 None
212 }
213 })
214 }
215
216 /// Returns an iterator over upper leaves (i.e., depth = 16, 32, or 48) for this [TieredSmt].
217 ///
218 /// Each yielded item is a (node, key, value) tuple where key is a full un-truncated key (i.e.,
219 /// with key[3] element unmodified).
220 ///
221 /// The iterator order is unspecified.
222 pub fn upper_leaves(&self) -> impl Iterator<Item = (RpoDigest, RpoDigest, Word)> + '_ {
223 self.upper_leaves.iter().map(|(index, key)| {
224 let node = self.get_node_unchecked(index);
225 let value = self.get_value(*key);
226 (node, *key, value)
227 })
228 }
229
230 /// Returns an iterator over bottom leaves (i.e., depth = 64) of this [TieredSmt].
231 ///
232 /// Each yielded item consists of the hash of the leaf and its contents, where contents is
233 /// a vector containing key-value pairs of entries storied in this leaf. Note that keys are
234 /// un-truncated keys (i.e., with key[3] element unmodified).
235 ///
236 /// The iterator order is unspecified.
237 pub fn bottom_leaves(&self) -> impl Iterator<Item = (RpoDigest, Vec<(RpoDigest, Word)>)> + '_ {
238 self.bottom_leaves.values().map(|leaf| (leaf.hash(), leaf.contents()))
239 }
240
241 // HELPER METHODS
242 // --------------------------------------------------------------------------------------------
243
244 /// Checks if the specified index is valid in the context of this Merkle tree.
245 ///
246 /// # Errors
247 /// Returns an error if:
248 /// - The specified index depth is 0 or greater than 64.
249 /// - The node for the specified index does not exists in the Merkle tree. This is possible
250 /// when an ancestors of the specified index is a leaf node.
251 fn validate_node_access(&self, index: NodeIndex) -> Result<(), MerkleError> {
252 if index.is_root() {
253 return Err(MerkleError::DepthTooSmall(index.depth()));
254 } else if index.depth() > Self::MAX_DEPTH {
255 return Err(MerkleError::DepthTooBig(index.depth() as u64));
256 } else {
257 // make sure that there are no leaf nodes in the ancestors of the index; since leaf
258 // nodes can live at specific depth, we just need to check these depths.
259 let tier = get_index_tier(&index);
260 let mut tier_index = index;
261 for &depth in Self::TIER_DEPTHS[..tier].iter().rev() {
262 tier_index.move_up_to(depth);
263 if self.upper_leaves.contains_key(&tier_index) {
264 return Err(MerkleError::NodeNotInSet(index));
265 }
266 }
267 }
268
269 Ok(())
270 }
271
272 /// Returns a node at the specified index. If the node does not exist at this index, a root
273 /// for an empty subtree at the index's depth is returned.
274 ///
275 /// Unlike [TieredSmt::get_node()] this does not perform any checks to verify that the returned
276 /// node is valid in the context of this tree.
277 fn get_node_unchecked(&self, index: &NodeIndex) -> RpoDigest {
278 match self.nodes.get(index) {
279 Some(node) => *node,
280 None => EmptySubtreeRoots::empty_hashes(Self::MAX_DEPTH)[index.depth() as usize],
281 }
282 }
283
284 /// Returns an index at which a node for the specified key should be inserted. If a leaf node
285 /// already exists at that index, returns the key associated with that leaf node.
286 ///
287 /// In case the index falls into the bottom tier (depth = 64), leaf node key is not returned
288 /// as the bottom tier may contain multiple key-value pairs in the same leaf.
289 fn get_insert_location(&self, key: &RpoDigest) -> (NodeIndex, Option<RpoDigest>) {
290 // traverse the tree from the root down checking nodes at tiers 16, 32, and 48. Return if
291 // a node at any of the tiers is either a leaf or a root of an empty subtree.
292 let mse = Word::from(key)[3].as_int();
293 for depth in (Self::TIER_DEPTHS[0]..Self::MAX_DEPTH).step_by(Self::TIER_SIZE as usize) {
294 let index = NodeIndex::new_unchecked(depth, mse >> (Self::MAX_DEPTH - depth));
295 if let Some(leaf_key) = self.upper_leaves.get(&index) {
296 return (index, Some(*leaf_key));
297 } else if !self.nodes.contains_key(&index) {
298 return (index, None);
299 }
300 }
301
302 // if we got here, that means all of the nodes checked so far are internal nodes, and
303 // the new node would need to be inserted in the bottom tier.
304 let index = NodeIndex::new_unchecked(Self::MAX_DEPTH, mse);
305 (index, None)
306 }
307
308 /// Inserts the provided key-value pair at the specified index and updates the root of this
309 /// Merkle tree by recomputing the path to the root.
310 fn insert_leaf_node(&mut self, mut index: NodeIndex, key: RpoDigest, value: Word) {
311 let depth = index.depth();
312 debug_assert!(Self::TIER_DEPTHS.contains(&depth));
313
314 // insert the key into index-key map and compute the new value of the node
315 let mut node = if index.depth() == Self::MAX_DEPTH {
316 // for the bottom tier, we add the key-value pair to the existing leaf, or create a
317 // new leaf with this key-value pair
318 self.bottom_leaves
319 .entry(index.value())
320 .and_modify(|leaves| leaves.add_value(key, value))
321 .or_insert(BottomLeaf::new(key, value))
322 .hash()
323 } else {
324 // for the upper tiers, we just update the index-key map and compute the value of the
325 // node
326 self.upper_leaves.insert(index, key);
327 // the node value is computed as: hash(key || value, domain = depth)
328 Rpo256::merge_in_domain(&[key, value.into()], depth.into())
329 };
330
331 // insert the node and update the path from the node to the root
332 for _ in 0..index.depth() {
333 self.nodes.insert(index, node);
334 let sibling = self.get_node_unchecked(&index.sibling());
335 node = Rpo256::merge(&index.build_node(node, sibling));
336 index.move_up();
337 }
338
339 // update the root
340 self.nodes.insert(NodeIndex::root(), node);
341 self.root = node;
342 }
343}
344
345impl Default for TieredSmt {
346 fn default() -> Self {
347 Self {
348 root: EmptySubtreeRoots::empty_hashes(Self::MAX_DEPTH)[0],
349 nodes: BTreeMap::new(),
350 upper_leaves: BTreeMap::new(),
351 bottom_leaves: BTreeMap::new(),
352 values: BTreeMap::new(),
353 }
354 }
355}
356
357// HELPER FUNCTIONS
358// ================================================================================================
359
360/// Returns index for the specified key inserted at the specified depth.
361///
362/// The value for the key is computed by taking n most significant bits from the most significant
363/// element of the key, where n is the specified depth.
364fn key_to_index(key: &RpoDigest, depth: u8) -> NodeIndex {
365 let mse = Word::from(key)[3].as_int();
366 let value = match depth {
367 16 | 32 | 48 | 64 => mse >> ((TieredSmt::MAX_DEPTH - depth) as u32),
368 _ => unreachable!("invalid depth: {depth}"),
369 };
370 NodeIndex::new_unchecked(depth, value)
371}
372
373/// Returns tiered common prefix length between the most significant elements of the provided keys.
374///
375/// Specifically:
376/// - returns 64 if the most significant elements are equal.
377/// - returns 48 if the common prefix is between 48 and 63 bits.
378/// - returns 32 if the common prefix is between 32 and 47 bits.
379/// - returns 16 if the common prefix is between 16 and 31 bits.
380/// - returns 0 if the common prefix is fewer than 16 bits.
381fn get_common_prefix_tier(key1: &RpoDigest, key2: &RpoDigest) -> u8 {
382 let e1 = Word::from(key1)[3].as_int();
383 let e2 = Word::from(key2)[3].as_int();
384 let ex = (e1 ^ e2).leading_zeros() as u8;
385 (ex / 16) * 16
386}
387
388/// Returns a tier for the specified index.
389///
390/// The tiers are defined as follows:
391/// - Tier 0: depth 0 through 16 (inclusive).
392/// - Tier 1: depth 17 through 32 (inclusive).
393/// - Tier 2: depth 33 through 48 (inclusive).
394/// - Tier 3: depth 49 through 64 (inclusive).
395const fn get_index_tier(index: &NodeIndex) -> usize {
396 debug_assert!(index.depth() <= TieredSmt::MAX_DEPTH, "invalid depth");
397 match index.depth() {
398 0..=16 => 0,
399 17..=32 => 1,
400 33..=48 => 2,
401 _ => 3,
402 }
403}
404
405/// Returns true if the specified index is an index for an inner node (i.e., the depth is not 16,
406/// 32, 48, or 64).
407const fn is_inner_node(index: &NodeIndex) -> bool {
408 !matches!(index.depth(), 16 | 32 | 48 | 64)
409}
410
411// BOTTOM LEAF
412// ================================================================================================
413
414/// Stores contents of the bottom leaf (i.e., leaf at depth = 64) in a [TieredSmt].
415///
416/// Bottom leaf can contain one or more key-value pairs all sharing the same 64-bit key prefix.
417/// The values are sorted by key to make sure the structure of the leaf is independent of the
418/// insertion order. This guarantees that a leaf with the same set of key-value pairs always has
419/// the same hash value.
420#[derive(Debug, Clone, PartialEq, Eq)]
421struct BottomLeaf {
422 prefix: u64,
423 values: BTreeMap<[u64; 4], Word>,
424}
425
426impl BottomLeaf {
427 /// Returns a new [BottomLeaf] with a single key-value pair added.
428 pub fn new(key: RpoDigest, value: Word) -> Self {
429 let prefix = Word::from(key)[3].as_int();
430 let mut values = BTreeMap::new();
431 values.insert(key.into(), value);
432 Self { prefix, values }
433 }
434
435 /// Adds a new key-value pair to this leaf.
436 pub fn add_value(&mut self, key: RpoDigest, value: Word) {
437 self.values.insert(key.into(), value);
438 }
439
440 /// Computes a hash of this leaf.
441 pub fn hash(&self) -> RpoDigest {
442 let mut elements = Vec::with_capacity(self.values.len() * 2);
443 for (key, val) in self.values.iter() {
444 key.iter().for_each(|&v| elements.push(Felt::new(v)));
445 elements.extend_from_slice(val.as_slice());
446 }
447 // TODO: hash in domain
448 Rpo256::hash_elements(&elements)
449 }
450
451 /// Returns contents of this leaf as a vector of (key, value) pairs.
452 ///
453 /// The keys are returned in their un-truncated form.
454 pub fn contents(&self) -> Vec<(RpoDigest, Word)> {
455 self.values
456 .iter()
457 .map(|(key, val)| {
458 let key = RpoDigest::from([
459 Felt::new(key[0]),
460 Felt::new(key[1]),
461 Felt::new(key[2]),
462 Felt::new(key[3]),
463 ]);
464 (key, *val)
465 })
466 .collect()
467 }
468}