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 core::cmp;
6
7mod nodes;
8use nodes::NodeStore;
9
10mod values;
11use values::ValueStore;
12
13#[cfg(test)]
14mod tests;
15
16// TIERED SPARSE MERKLE TREE
17// ================================================================================================
18
19/// Tiered (compacted) Sparse Merkle tree mapping 256-bit keys to 256-bit values. Both keys and
20/// values are represented by 4 field elements.
21///
22/// Leaves in the tree can exist only on specific depths called "tiers". These depths are: 16, 32,
23/// 48, and 64. Initially, when a tree is empty, it is equivalent to an empty Sparse Merkle tree
24/// of depth 64 (i.e., leaves at depth 64 are set to [ZERO; 4]). As non-empty values are inserted
25/// into the tree they are added to the first available tier.
26///
27/// For example, when the first key-value pair is inserted, it will be stored in a node at depth
28/// 16 such that the 16 most significant bits of the key determine the position of the node at
29/// depth 16. If another value with a key sharing the same 16-bit prefix is inserted, both values
30/// move into the next tier (depth 32). This process is repeated until values end up at the bottom
31/// tier (depth 64). If multiple values have keys with a common 64-bit prefix, such key-value pairs
32/// are stored in a sorted list at the bottom tier.
33///
34/// To differentiate between internal and leaf nodes, node values are computed as follows:
35/// - Internal nodes: hash(left_child, right_child).
36/// - Leaf node at depths 16, 32, or 64: hash(key, value, domain=depth).
37/// - Leaf node at depth 64: hash([key_0, value_0, ..., key_n, value_n, domain=64]).
38#[derive(Debug, Clone, PartialEq, Eq)]
39pub struct TieredSmt {
40 root: RpoDigest,
41 nodes: NodeStore,
42 values: ValueStore,
43}
44
45impl TieredSmt {
46 // CONSTANTS
47 // --------------------------------------------------------------------------------------------
48
49 /// The number of levels between tiers.
50 const TIER_SIZE: u8 = 16;
51
52 /// Depths at which leaves can exist in a tiered SMT.
53 const TIER_DEPTHS: [u8; 4] = [16, 32, 48, 64];
54
55 /// Maximum node depth. This is also the bottom tier of the tree.
56 const MAX_DEPTH: u8 = 64;
57
58 /// Value of an empty leaf.
59 pub const EMPTY_VALUE: Word = super::empty_roots::EMPTY_WORD;
60
61 // CONSTRUCTORS
62 // --------------------------------------------------------------------------------------------
63
64 /// Returns a new [TieredSmt] instantiated with the specified key-value pairs.
65 ///
66 /// # Errors
67 /// Returns an error if the provided entries contain multiple values for the same key.
68 pub fn with_leaves<R, I>(entries: R) -> Result<Self, MerkleError>
69 where
70 R: IntoIterator<IntoIter = I>,
71 I: Iterator<Item = (RpoDigest, Word)> + ExactSizeIterator,
72 {
73 // create an empty tree
74 let mut tree = Self::default();
75
76 // append leaves to the tree returning an error if a duplicate entry for the same key
77 // is found
78 let mut empty_entries = BTreeSet::new();
79 for (key, value) in entries {
80 let old_value = tree.insert(key, value);
81 if old_value != Self::EMPTY_VALUE || empty_entries.contains(&key) {
82 return Err(MerkleError::DuplicateValuesForKey(key));
83 }
84 // if we've processed an empty entry, add the key to the set of empty entry keys, and
85 // if this key was already in the set, return an error
86 if value == Self::EMPTY_VALUE && !empty_entries.insert(key) {
87 return Err(MerkleError::DuplicateValuesForKey(key));
88 }
89 }
90 Ok(tree)
91 }
92
93 // PUBLIC ACCESSORS
94 // --------------------------------------------------------------------------------------------
95
96 /// Returns the root of this Merkle tree.
97 pub const fn root(&self) -> RpoDigest {
98 self.root
99 }
100
101 /// Returns a node at the specified index.
102 ///
103 /// # Errors
104 /// Returns an error if:
105 /// - The specified index depth is 0 or greater than 64.
106 /// - The node with the specified index does not exists in the Merkle tree. This is possible
107 /// when a leaf node with the same index prefix exists at a tier higher than the requested
108 /// node.
109 pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
110 self.nodes.get_node(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, index: NodeIndex) -> Result<MerklePath, MerkleError> {
124 self.nodes.get_path(index)
125 }
126
127 /// Returns the value associated with the specified key.
128 ///
129 /// If nothing was inserted into this tree for the specified key, [ZERO; 4] is returned.
130 pub fn get_value(&self, key: RpoDigest) -> Word {
131 match self.values.get(&key) {
132 Some(value) => *value,
133 None => Self::EMPTY_VALUE,
134 }
135 }
136
137 // STATE MUTATORS
138 // --------------------------------------------------------------------------------------------
139
140 /// Inserts the provided value into the tree under the specified key and returns the value
141 /// previously stored under this key.
142 ///
143 /// If the value for the specified key was not previously set, [ZERO; 4] is returned.
144 pub fn insert(&mut self, key: RpoDigest, value: Word) -> Word {
145 // if an empty value is being inserted, remove the leaf node to make it look as if the
146 // value was never inserted
147 if value == Self::EMPTY_VALUE {
148 return self.remove_leaf_node(key);
149 }
150
151 // insert the value into the value store, and if nothing has changed, return
152 let (old_value, is_update) = match self.values.insert(key, value) {
153 Some(old_value) => {
154 if old_value == value {
155 return old_value;
156 }
157 (old_value, true)
158 }
159 None => (Self::EMPTY_VALUE, false),
160 };
161
162 // determine the index for the value node; this index could have 3 different meanings:
163 // - it points to a root of an empty subtree (excluding depth = 64); in this case, we can
164 // replace the node with the value node immediately.
165 // - it points to a node at the bottom tier (i.e., depth = 64); in this case, we need to
166 // process bottom-tier insertion which will be handled by insert_leaf_node().
167 // - it points to an existing leaf node; this node could be a node with the same key or a
168 // different key with a common prefix; in the latter case, we'll need to move the leaf
169 // to a lower tier
170 let (index, leaf_exists) = self.nodes.get_insert_location(&key);
171 debug_assert!(!is_update || leaf_exists);
172
173 // if the returned index points to a leaf, and this leaf is for a different key (i.e., we
174 // are not updating a value for an existing key), we need to replace this leaf with a tree
175 // containing leaves for both the old and the new key-value pairs
176 if leaf_exists && !is_update {
177 // get the key-value pair for the key with the same prefix; since the key-value
178 // pair has already been inserted into the value store, we need to filter it out
179 // when looking for the other key-value pair
180 let (other_key, other_value) = self
181 .values
182 .get_first_filtered(index_to_prefix(&index), &key)
183 .expect("other key-value pair not found");
184
185 // determine how far down the tree should we move the leaves
186 let common_prefix_len = get_common_prefix_tier(&key, other_key);
187 let depth = cmp::min(common_prefix_len + Self::TIER_SIZE, Self::MAX_DEPTH);
188
189 // compute node locations for new and existing key-value paris
190 let new_index = key_to_index(&key, depth);
191 let other_index = key_to_index(other_key, depth);
192
193 // compute node values for the new and existing key-value pairs
194 let new_node = self.build_leaf_node(new_index, key, value);
195 let other_node = self.build_leaf_node(other_index, *other_key, *other_value);
196
197 // replace the leaf located at index with a subtree containing nodes for new and
198 // existing key-value paris
199 self.root = self.nodes.replace_leaf_with_subtree(
200 index,
201 [(new_index, new_node), (other_index, other_node)],
202 );
203 } else {
204 // if the returned index points to an empty subtree, or a leaf with the same key (i.e.,
205 // we are performing an update), or a leaf is at the bottom tier, compute its node
206 // value and do a simple insert
207 let node = self.build_leaf_node(index, key, value);
208 self.root = self.nodes.insert_leaf_node(index, node);
209 }
210
211 old_value
212 }
213
214 // ITERATORS
215 // --------------------------------------------------------------------------------------------
216
217 /// Returns an iterator over all key-value pairs in this [TieredSmt].
218 pub fn iter(&self) -> impl Iterator<Item = &(RpoDigest, Word)> {
219 self.values.iter()
220 }
221
222 /// Returns an iterator over all inner nodes of this [TieredSmt] (i.e., nodes not at depths 16
223 /// 32, 48, or 64).
224 ///
225 /// The iterator order is unspecified.
226 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
227 self.nodes.inner_nodes()
228 }
229
230 /// Returns an iterator over upper leaves (i.e., depth = 16, 32, or 48) for this [TieredSmt]
231 /// where each yielded item is a (node, key, value) tuple.
232 ///
233 /// The iterator order is unspecified.
234 pub fn upper_leaves(&self) -> impl Iterator<Item = (RpoDigest, RpoDigest, Word)> + '_ {
235 self.nodes.upper_leaves().map(|(index, node)| {
236 let key_prefix = index_to_prefix(index);
237 let (key, value) = self.values.get_first(key_prefix).expect("upper leaf not found");
238 debug_assert_eq!(key_to_index(key, index.depth()), *index);
239 (*node, *key, *value)
240 })
241 }
242
243 /// Returns an iterator over bottom leaves (i.e., depth = 64) of this [TieredSmt].
244 ///
245 /// Each yielded item consists of the hash of the leaf and its contents, where contents is
246 /// a vector containing key-value pairs of entries storied in this leaf.
247 ///
248 /// The iterator order is unspecified.
249 pub fn bottom_leaves(&self) -> impl Iterator<Item = (RpoDigest, Vec<(RpoDigest, Word)>)> + '_ {
250 self.nodes.bottom_leaves().map(|(&prefix, node)| {
251 let values = self.values.get_all(prefix).expect("bottom leaf not found");
252 (*node, values)
253 })
254 }
255
256 // HELPER METHODS
257 // --------------------------------------------------------------------------------------------
258
259 /// Removes the node holding the key-value pair for the specified key from this tree, and
260 /// returns the value associated with the specified key.
261 ///
262 /// If no value was associated with the specified key, [ZERO; 4] is returned.
263 fn remove_leaf_node(&mut self, key: RpoDigest) -> Word {
264 // remove the key-value pair from the value store; if no value was associated with the
265 // specified key, return.
266 let old_value = match self.values.remove(&key) {
267 Some(old_value) => old_value,
268 None => return Self::EMPTY_VALUE,
269 };
270
271 // determine the location of the leaf holding the key-value pair to be removed
272 let (index, leaf_exists) = self.nodes.get_insert_location(&key);
273 debug_assert!(index.depth() == Self::MAX_DEPTH || leaf_exists);
274
275 // if the leaf is at the bottom tier and after removing the key-value pair from it, the
276 // leaf is still not empty, just recompute its hash and update the leaf node.
277 if index.depth() == Self::MAX_DEPTH {
278 if let Some(values) = self.values.get_all(index.value()) {
279 let node = hash_bottom_leaf(&values);
280 self.root = self.nodes.update_leaf_node(index, node);
281 return old_value;
282 };
283 }
284
285 // if the removed key-value pair has a lone sibling at the current tier with a root at
286 // higher tier, we need to move the sibling to a higher tier
287 if let Some((sib_key, sib_val, new_sib_index)) = self.values.get_lone_sibling(index) {
288 // determine the current index of the sibling node
289 let sib_index = key_to_index(sib_key, index.depth());
290 debug_assert!(sib_index.depth() > new_sib_index.depth());
291
292 // compute node value for the new location of the sibling leaf and replace the subtree
293 // with this leaf node
294 let node = self.build_leaf_node(new_sib_index, *sib_key, *sib_val);
295 let new_sib_depth = new_sib_index.depth();
296 self.root = self.nodes.replace_subtree_with_leaf(index, sib_index, new_sib_depth, node);
297 } else {
298 // if the removed key-value pair did not have a sibling at the current tier with a
299 // root at higher tiers, just clear the leaf node
300 self.root = self.nodes.clear_leaf_node(index);
301 }
302
303 old_value
304 }
305
306 /// Builds and returns a leaf node value for the node located as the specified index.
307 ///
308 /// This method assumes that the key-value pair for the node has already been inserted into
309 /// the value store, however, for depths 16, 32, and 48, the node is computed directly from
310 /// the passed-in values (for depth 64, the value store is queried to get all the key-value
311 /// pairs located at the specified index).
312 fn build_leaf_node(&self, index: NodeIndex, key: RpoDigest, value: Word) -> RpoDigest {
313 let depth = index.depth();
314 debug_assert!(Self::TIER_DEPTHS.contains(&depth));
315
316 // insert the key into index-key map and compute the new value of the node
317 if index.depth() == Self::MAX_DEPTH {
318 // for the bottom tier, we add the key-value pair to the existing leaf, or create a
319 // new leaf with this key-value pair
320 let values = self.values.get_all(index.value()).unwrap();
321 hash_bottom_leaf(&values)
322 } else {
323 debug_assert_eq!(self.values.get_first(index_to_prefix(&index)), Some(&(key, value)));
324 hash_upper_leaf(key, value, depth)
325 }
326 }
327}
328
329impl Default for TieredSmt {
330 fn default() -> Self {
331 let root = EmptySubtreeRoots::empty_hashes(Self::MAX_DEPTH)[0];
332 Self {
333 root,
334 nodes: NodeStore::new(root),
335 values: ValueStore::default(),
336 }
337 }
338}
339
340// HELPER FUNCTIONS
341// ================================================================================================
342
343/// Returns the value representing the 64 most significant bits of the specified key.
344fn get_key_prefix(key: &RpoDigest) -> u64 {
345 Word::from(key)[3].as_int()
346}
347
348/// Returns the index value shifted to be in the most significant bit positions of the returned
349/// u64 value.
350fn index_to_prefix(index: &NodeIndex) -> u64 {
351 index.value() << (TieredSmt::MAX_DEPTH - index.depth())
352}
353
354/// Returns index for the specified key inserted at the specified depth.
355///
356/// The value for the key is computed by taking n most significant bits from the most significant
357/// element of the key, where n is the specified depth.
358fn key_to_index(key: &RpoDigest, depth: u8) -> NodeIndex {
359 let mse = get_key_prefix(key);
360 let value = match depth {
361 16 | 32 | 48 | 64 => mse >> ((TieredSmt::MAX_DEPTH - depth) as u32),
362 _ => unreachable!("invalid depth: {depth}"),
363 };
364 NodeIndex::new_unchecked(depth, value)
365}
366
367/// Returns tiered common prefix length between the most significant elements of the provided keys.
368///
369/// Specifically:
370/// - returns 64 if the most significant elements are equal.
371/// - returns 48 if the common prefix is between 48 and 63 bits.
372/// - returns 32 if the common prefix is between 32 and 47 bits.
373/// - returns 16 if the common prefix is between 16 and 31 bits.
374/// - returns 0 if the common prefix is fewer than 16 bits.
375fn get_common_prefix_tier(key1: &RpoDigest, key2: &RpoDigest) -> u8 {
376 let e1 = get_key_prefix(key1);
377 let e2 = get_key_prefix(key2);
378 let ex = (e1 ^ e2).leading_zeros() as u8;
379 (ex / 16) * 16
380}
381
382/// Returns a tier for the specified index.
383///
384/// The tiers are defined as follows:
385/// - Tier 0: depth 0 through 16 (inclusive).
386/// - Tier 1: depth 17 through 32 (inclusive).
387/// - Tier 2: depth 33 through 48 (inclusive).
388/// - Tier 3: depth 49 through 64 (inclusive).
389const fn get_index_tier(index: &NodeIndex) -> usize {
390 debug_assert!(index.depth() <= TieredSmt::MAX_DEPTH, "invalid depth");
391 match index.depth() {
392 0..=16 => 0,
393 17..=32 => 1,
394 33..=48 => 2,
395 _ => 3,
396 }
397}
398
399/// Returns true if the specified index is an index for an leaf node (i.e., the depth is 16, 32,
400/// 48, or 64).
401const fn is_leaf_node(index: &NodeIndex) -> bool {
402 matches!(index.depth(), 16 | 32 | 48 | 64)
403}
404
405/// Computes node value for leaves at tiers 16, 32, or 48.
406///
407/// Node value is computed as: hash(key || value, domain = depth).
408pub fn hash_upper_leaf(key: RpoDigest, value: Word, depth: u8) -> RpoDigest {
409 const NUM_UPPER_TIERS: usize = TieredSmt::TIER_DEPTHS.len() - 1;
410 debug_assert!(TieredSmt::TIER_DEPTHS[..NUM_UPPER_TIERS].contains(&depth));
411 Rpo256::merge_in_domain(&[key, value.into()], depth.into())
412}
413
414/// Computes node value for leaves at the bottom tier (depth 64).
415///
416/// Node value is computed as: hash([key_0, value_0, ..., key_n, value_n, domain=64]).
417pub fn hash_bottom_leaf(values: &[(RpoDigest, Word)]) -> RpoDigest {
418 let mut elements = Vec::with_capacity(values.len() * 8);
419 for (key, val) in values.iter() {
420 elements.extend_from_slice(key.as_elements());
421 elements.extend_from_slice(val.as_slice());
422 }
423 // TODO: hash in domain
424 Rpo256::hash_elements(&elements)
425}