miden-crypto/src/merkle/tiered_smt/mod.rs
src/merkle/tiered_smt/mod.rsRaw
1use super::{
2 BTreeMap, BTreeSet, EmptySubtreeRoots, Felt, MerkleError, MerklePath, NodeIndex, Rpo256,
3 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 // HELPER METHODS
193 // --------------------------------------------------------------------------------------------
194
195 /// Checks if the specified index is valid in the context of this Merkle tree.
196 ///
197 /// # Errors
198 /// Returns an error if:
199 /// - The specified index depth is 0 or greater than 64.
200 /// - The node for the specified index does not exists in the Merkle tree. This is possible
201 /// when an ancestors of the specified index is a leaf node.
202 fn validate_node_access(&self, index: NodeIndex) -> Result<(), MerkleError> {
203 if index.is_root() {
204 return Err(MerkleError::DepthTooSmall(index.depth()));
205 } else if index.depth() > Self::MAX_DEPTH {
206 return Err(MerkleError::DepthTooBig(index.depth() as u64));
207 } else {
208 // make sure that there are no leaf nodes in the ancestors of the index; since leaf
209 // nodes can live at specific depth, we just need to check these depths.
210 let tier = get_index_tier(&index);
211 let mut tier_index = index;
212 for &depth in Self::TIER_DEPTHS[..tier].iter().rev() {
213 tier_index.move_up_to(depth);
214 if self.upper_leaves.contains_key(&tier_index) {
215 return Err(MerkleError::NodeNotInSet(index));
216 }
217 }
218 }
219
220 Ok(())
221 }
222
223 /// Returns a node at the specified index. If the node does not exist at this index, a root
224 /// for an empty subtree at the index's depth is returned.
225 ///
226 /// Unlike [TieredSmt::get_node()] this does not perform any checks to verify that the returned
227 /// node is valid in the context of this tree.
228 fn get_node_unchecked(&self, index: &NodeIndex) -> RpoDigest {
229 match self.nodes.get(index) {
230 Some(node) => *node,
231 None => EmptySubtreeRoots::empty_hashes(Self::MAX_DEPTH)[index.depth() as usize],
232 }
233 }
234
235 /// Returns an index at which a node for the specified key should be inserted. If a leaf node
236 /// already exists at that index, returns the key associated with that leaf node.
237 ///
238 /// In case the index falls into the bottom tier (depth = 64), leaf node key is not returned
239 /// as the bottom tier may contain multiple key-value pairs in the same leaf.
240 fn get_insert_location(&self, key: &RpoDigest) -> (NodeIndex, Option<RpoDigest>) {
241 // traverse the tree from the root down checking nodes at tiers 16, 32, and 48. Return if
242 // a node at any of the tiers is either a leaf or a root of an empty subtree.
243 let mse = Word::from(key)[3].as_int();
244 for depth in (Self::TIER_DEPTHS[0]..Self::MAX_DEPTH).step_by(Self::TIER_SIZE as usize) {
245 let index = NodeIndex::new_unchecked(depth, mse >> (Self::MAX_DEPTH - depth));
246 if let Some(leaf_key) = self.upper_leaves.get(&index) {
247 return (index, Some(*leaf_key));
248 } else if !self.nodes.contains_key(&index) {
249 return (index, None);
250 }
251 }
252
253 // if we got here, that means all of the nodes checked so far are internal nodes, and
254 // the new node would need to be inserted in the bottom tier.
255 let index = NodeIndex::new_unchecked(Self::MAX_DEPTH, mse);
256 (index, None)
257 }
258
259 /// Inserts the provided key-value pair at the specified index and updates the root of this
260 /// Merkle tree by recomputing the path to the root.
261 fn insert_node(&mut self, mut index: NodeIndex, key: RpoDigest, value: Word) {
262 let depth = index.depth();
263
264 // insert the key into index-key map and compute the new value of the node
265 let mut node = if index.depth() == Self::MAX_DEPTH {
266 // for the bottom tier, we add the key-value pair to the existing leaf, or create a
267 // new leaf with this key-value pair
268 self.bottom_leaves
269 .entry(index.value())
270 .and_modify(|leaves| leaves.add_value(key, value))
271 .or_insert(BottomLeaf::new(key, value))
272 .hash()
273 } else {
274 // for the upper tiers, we just update the index-key map and compute the value of the
275 // node
276 self.upper_leaves.insert(index, key);
277 // the node value is computed as: hash(remaining_key || value, domain = depth)
278 let remaining_path = get_remaining_path(key, depth.into());
279 Rpo256::merge_in_domain(&[remaining_path, value.into()], depth.into())
280 };
281
282 // insert the node and update the path from the node to the root
283 for _ in 0..index.depth() {
284 self.nodes.insert(index, node);
285 let sibling = self.get_node_unchecked(&index.sibling());
286 node = Rpo256::merge(&index.build_node(node, sibling));
287 index.move_up();
288 }
289
290 // update the root
291 self.root = node;
292 }
293}
294
295impl Default for TieredSmt {
296 fn default() -> Self {
297 Self {
298 root: EmptySubtreeRoots::empty_hashes(Self::MAX_DEPTH)[0],
299 nodes: BTreeMap::new(),
300 upper_leaves: BTreeMap::new(),
301 bottom_leaves: BTreeMap::new(),
302 values: BTreeMap::new(),
303 }
304 }
305}
306
307// HELPER FUNCTIONS
308// ================================================================================================
309
310/// Returns the remaining path for the specified key at the specified depth.
311///
312/// Remaining path is computed by setting n most significant bits of the key to zeros, where n is
313/// the specified depth.
314fn get_remaining_path(key: RpoDigest, depth: u32) -> RpoDigest {
315 let mut key = Word::from(key);
316 key[3] = if depth == 64 {
317 ZERO
318 } else {
319 // remove `depth` bits from the most significant key element
320 ((key[3].as_int() << depth) >> depth).into()
321 };
322 key.into()
323}
324
325/// Returns index for the specified key inserted at the specified depth.
326///
327/// The value for the key is computed by taking n most significant bits from the most significant
328/// element of the key, where n is the specified depth.
329fn key_to_index(key: &RpoDigest, depth: u8) -> NodeIndex {
330 let mse = Word::from(key)[3].as_int();
331 let value = match depth {
332 16 | 32 | 48 | 64 => mse >> ((TieredSmt::MAX_DEPTH - depth) as u32),
333 _ => unreachable!("invalid depth: {depth}"),
334 };
335 NodeIndex::new_unchecked(depth, value)
336}
337
338/// Returns tiered common prefix length between the most significant elements of the provided keys.
339///
340/// Specifically:
341/// - returns 64 if the most significant elements are equal.
342/// - returns 48 if the common prefix is between 48 and 63 bits.
343/// - returns 32 if the common prefix is between 32 and 47 bits.
344/// - returns 16 if the common prefix is between 16 and 31 bits.
345/// - returns 0 if the common prefix is fewer than 16 bits.
346fn get_common_prefix_tier(key1: &RpoDigest, key2: &RpoDigest) -> u8 {
347 let e1 = Word::from(key1)[3].as_int();
348 let e2 = Word::from(key2)[3].as_int();
349 let ex = (e1 ^ e2).leading_zeros() as u8;
350 (ex / 16) * 16
351}
352
353/// Returns a tier for the specified index.
354///
355/// The tiers are defined as follows:
356/// - Tier 0: depth 0 through 16 (inclusive).
357/// - Tier 1: depth 17 through 32 (inclusive).
358/// - Tier 2: depth 33 through 48 (inclusive).
359/// - Tier 3: depth 49 through 64 (inclusive).
360const fn get_index_tier(index: &NodeIndex) -> usize {
361 debug_assert!(index.depth() <= TieredSmt::MAX_DEPTH, "invalid depth");
362 match index.depth() {
363 0..=16 => 0,
364 17..=32 => 1,
365 33..=48 => 2,
366 _ => 3,
367 }
368}
369
370// BOTTOM LEAF
371// ================================================================================================
372
373/// Stores contents of the bottom leaf (i.e., leaf at depth = 64) in a [TieredSmt].
374///
375/// Bottom leaf can contain one or more key-value pairs all sharing the same 64-bit key prefix.
376/// The values are sorted by key to make sure the structure of the leaf is independent of the
377/// insertion order. This guarantees that a leaf with the same set of key-value pairs always has
378/// the same hash value.
379#[derive(Debug, Clone, PartialEq, Eq)]
380struct BottomLeaf {
381 values: BTreeMap<[u64; 4], Word>,
382}
383
384impl BottomLeaf {
385 /// Returns a new [BottomLeaf] with a single key-value pair added.
386 pub fn new(key: RpoDigest, value: Word) -> Self {
387 let mut values = BTreeMap::new();
388 let key = get_remaining_path(key, TieredSmt::MAX_DEPTH as u32);
389 values.insert(key.into(), value);
390 Self { values }
391 }
392
393 /// Adds a new key-value pair to this leaf.
394 pub fn add_value(&mut self, key: RpoDigest, value: Word) {
395 let key = get_remaining_path(key, TieredSmt::MAX_DEPTH as u32);
396 self.values.insert(key.into(), value);
397 }
398
399 /// Computes a hash of this leaf.
400 pub fn hash(&self) -> RpoDigest {
401 let mut elements = Vec::with_capacity(self.values.len() * 2);
402 for (key, val) in self.values.iter() {
403 key.iter().for_each(|&v| elements.push(Felt::new(v)));
404 elements.extend_from_slice(val);
405 }
406 // TODO: hash in domain
407 Rpo256::hash_elements(&elements)
408 }
409}