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