miden-crypto/src/merkle/smt/full/mod.rs
src/merkle/smt/full/mod.rsRaw
1use core::cmp::Ordering;
2
3use winter_math::StarkField;
4
5use crate::hash::rpo::Rpo256;
6use crate::merkle::{EmptySubtreeRoots, InnerNodeInfo};
7use crate::utils::{
8 collections::{BTreeMap, BTreeSet, Vec},
9 vec,
10};
11use crate::{Felt, EMPTY_WORD};
12
13use super::{
14 InnerNode, LeafIndex, MerkleError, MerklePath, NodeIndex, RpoDigest, SparseMerkleTree, Word,
15};
16
17#[cfg(test)]
18mod tests;
19
20// CONSTANTS
21// ================================================================================================
22
23pub const SMT_DEPTH: u8 = 64;
24
25// SMT
26// ================================================================================================
27
28/// Sparse Merkle tree mapping 256-bit keys to 256-bit values. Both keys and values are represented
29/// by 4 field elements.
30///
31/// All leaves sit at depth 64. The most significant element of the key is used to identify the leaf to
32/// which the key maps.
33///
34/// A leaf is either empty, or holds one or more key-value pairs. An empty leaf hashes to the empty
35/// word. Otherwise, a leaf hashes to the hash of its key-value pairs, ordered by key first, value
36/// second.
37#[derive(Debug, Clone, PartialEq, Eq)]
38#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
39pub struct Smt {
40 root: RpoDigest,
41 leaves: BTreeMap<u64, SmtLeaf>,
42 inner_nodes: BTreeMap<NodeIndex, InnerNode>,
43}
44
45impl Smt {
46 // CONSTANTS
47 // --------------------------------------------------------------------------------------------
48 /// The default value used to compute the hash of empty leaves
49 pub const EMPTY_VALUE: Word = <Self as SparseMerkleTree<SMT_DEPTH>>::EMPTY_VALUE;
50
51 // CONSTRUCTORS
52 // --------------------------------------------------------------------------------------------
53
54 /// Returns a new [Smt].
55 ///
56 /// All leaves in the returned tree are set to [Self::EMPTY_VALUE].
57 pub fn new() -> Self {
58 let root = *EmptySubtreeRoots::entry(SMT_DEPTH, 0);
59
60 Self {
61 root,
62 leaves: BTreeMap::new(),
63 inner_nodes: BTreeMap::new(),
64 }
65 }
66
67 /// Returns a new [Smt] instantiated with leaves set as specified by the provided entries.
68 ///
69 /// All leaves omitted from the entries list are set to [Self::EMPTY_VALUE].
70 ///
71 /// # Errors
72 /// Returns an error if the provided entries contain multiple values for the same key.
73 pub fn with_entries(
74 entries: impl IntoIterator<Item = (RpoDigest, Word)>,
75 ) -> Result<Self, MerkleError> {
76 // create an empty tree
77 let mut tree = Self::new();
78
79 // This being a sparse data structure, the EMPTY_WORD is not assigned to the `BTreeMap`, so
80 // entries with the empty value need additional tracking.
81 let mut key_set_to_zero = BTreeSet::new();
82
83 for (key, value) in entries {
84 let old_value = tree.insert(key, value);
85
86 if old_value != EMPTY_WORD || key_set_to_zero.contains(&key) {
87 return Err(MerkleError::DuplicateValuesForIndex(
88 LeafIndex::<SMT_DEPTH>::from(key).value(),
89 ));
90 }
91
92 if value == EMPTY_WORD {
93 key_set_to_zero.insert(key);
94 };
95 }
96 Ok(tree)
97 }
98
99 // PUBLIC ACCESSORS
100 // --------------------------------------------------------------------------------------------
101
102 /// Returns the depth of the tree
103 pub const fn depth(&self) -> u8 {
104 SMT_DEPTH
105 }
106
107 /// Returns the root of the tree
108 pub fn root(&self) -> RpoDigest {
109 <Self as SparseMerkleTree<SMT_DEPTH>>::root(self)
110 }
111
112 /// Returns the leaf to which `key` maps
113 pub fn get_leaf(&self, key: &RpoDigest) -> SmtLeaf {
114 <Self as SparseMerkleTree<SMT_DEPTH>>::get_leaf(self, key)
115 }
116
117 /// Returns the value associated with `key`
118 pub fn get_value(&self, key: &RpoDigest) -> Word {
119 let leaf_pos = LeafIndex::<SMT_DEPTH>::from(*key).value();
120
121 match self.leaves.get(&leaf_pos) {
122 Some(leaf) => leaf.get_value(key),
123 None => EMPTY_WORD,
124 }
125 }
126
127 /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle
128 /// path to the leaf, as well as the leaf itself.
129 pub fn open(&self, key: &RpoDigest) -> (MerklePath, SmtLeaf) {
130 <Self as SparseMerkleTree<SMT_DEPTH>>::open(self, key)
131 }
132
133 // ITERATORS
134 // --------------------------------------------------------------------------------------------
135
136 /// Returns an iterator over the leaves of this [Smt].
137 pub fn leaves(&self) -> impl Iterator<Item = (LeafIndex<SMT_DEPTH>, &SmtLeaf)> {
138 self.leaves
139 .iter()
140 .map(|(leaf_index, leaf)| (LeafIndex::new_max_depth(*leaf_index), leaf))
141 }
142
143 /// Returns an iterator over the key-value pairs of this [Smt].
144 pub fn entries(&self) -> impl Iterator<Item = &(RpoDigest, Word)> {
145 self.leaves().flat_map(|(_, leaf)| leaf.entries())
146 }
147
148 /// Returns an iterator over the inner nodes of this [Smt].
149 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
150 self.inner_nodes.values().map(|e| InnerNodeInfo {
151 value: e.hash(),
152 left: e.left,
153 right: e.right,
154 })
155 }
156
157 // STATE MUTATORS
158 // --------------------------------------------------------------------------------------------
159
160 /// Inserts a value at the specified key, returning the previous value associated with that key.
161 /// Recall that by definition, any key that hasn't been updated is associated with
162 /// [`Self::EMPTY_VALUE`].
163 ///
164 /// This also recomputes all hashes between the leaf (associated with the key) and the root,
165 /// updating the root itself.
166 pub fn insert(&mut self, key: RpoDigest, value: Word) -> Word {
167 <Self as SparseMerkleTree<SMT_DEPTH>>::insert(self, key, value)
168 }
169
170 // HELPERS
171 // --------------------------------------------------------------------------------------------
172
173 /// Inserts `value` at leaf index pointed to by `key`. `value` is guaranteed to not be the empty
174 /// value, such that this is indeed an insertion.
175 fn perform_insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {
176 debug_assert_ne!(value, Self::EMPTY_VALUE);
177
178 let leaf_index: LeafIndex<SMT_DEPTH> = Self::key_to_leaf_index(&key);
179
180 match self.leaves.get_mut(&leaf_index.value()) {
181 Some(leaf) => leaf.insert(key, value),
182 None => {
183 self.leaves.insert(leaf_index.value(), SmtLeaf::Single((key, value)));
184
185 None
186 }
187 }
188 }
189
190 /// Removes key-value pair at leaf index pointed to by `key` if it exists.
191 fn perform_remove(&mut self, key: RpoDigest) -> Option<Word> {
192 let leaf_index: LeafIndex<SMT_DEPTH> = Self::key_to_leaf_index(&key);
193
194 if let Some(leaf) = self.leaves.get_mut(&leaf_index.value()) {
195 let (old_value, is_empty) = leaf.remove(key);
196 if is_empty {
197 self.leaves.remove(&leaf_index.value());
198 }
199 old_value
200 } else {
201 // there's nothing stored at the leaf; nothing to update
202 None
203 }
204 }
205}
206
207impl SparseMerkleTree<SMT_DEPTH> for Smt {
208 type Key = RpoDigest;
209 type Value = Word;
210 type Leaf = SmtLeaf;
211 type Opening = (MerklePath, SmtLeaf);
212
213 const EMPTY_VALUE: Self::Value = EMPTY_WORD;
214
215 fn root(&self) -> RpoDigest {
216 self.root
217 }
218
219 fn set_root(&mut self, root: RpoDigest) {
220 self.root = root;
221 }
222
223 fn get_inner_node(&self, index: NodeIndex) -> InnerNode {
224 self.inner_nodes.get(&index).cloned().unwrap_or_else(|| {
225 let node = EmptySubtreeRoots::entry(SMT_DEPTH, index.depth() + 1);
226
227 InnerNode { left: *node, right: *node }
228 })
229 }
230
231 fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode) {
232 self.inner_nodes.insert(index, inner_node);
233 }
234
235 fn remove_inner_node(&mut self, index: NodeIndex) {
236 let _ = self.inner_nodes.remove(&index);
237 }
238
239 fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value> {
240 // inserting an `EMPTY_VALUE` is equivalent to removing any value associated with `key`
241 if value != Self::EMPTY_VALUE {
242 self.perform_insert(key, value)
243 } else {
244 self.perform_remove(key)
245 }
246 }
247
248 fn get_leaf(&self, key: &RpoDigest) -> Self::Leaf {
249 let leaf_pos = LeafIndex::<SMT_DEPTH>::from(*key).value();
250
251 match self.leaves.get(&leaf_pos) {
252 Some(leaf) => leaf.clone(),
253 None => SmtLeaf::Empty,
254 }
255 }
256
257 fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest {
258 leaf.hash()
259 }
260
261 fn key_to_leaf_index(key: &RpoDigest) -> LeafIndex<SMT_DEPTH> {
262 let most_significant_felt = key[3];
263 LeafIndex::new_max_depth(most_significant_felt.as_int())
264 }
265}
266
267impl Default for Smt {
268 fn default() -> Self {
269 Self::new()
270 }
271}
272
273// LEAF
274// ================================================================================================
275
276#[derive(Clone, Debug, PartialEq, Eq)]
277#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
278pub enum SmtLeaf {
279 Empty,
280 Single((RpoDigest, Word)),
281 Multiple(Vec<(RpoDigest, Word)>),
282}
283
284impl SmtLeaf {
285 /// Converts a leaf to a list of field elements
286 pub fn to_elements(&self) -> Vec<Felt> {
287 self.clone().into_elements()
288 }
289
290 /// Converts a leaf to a list of field elements
291 pub fn into_elements(self) -> Vec<Felt> {
292 self.into_entries().into_iter().flat_map(kv_to_elements).collect()
293 }
294
295 /// Returns the key-value pairs in the leaf
296 pub fn entries(&self) -> Vec<&(RpoDigest, Word)> {
297 match self {
298 SmtLeaf::Empty => Vec::new(),
299 SmtLeaf::Single(kv_pair) => vec![kv_pair],
300 SmtLeaf::Multiple(kv_pairs) => kv_pairs.iter().collect(),
301 }
302 }
303
304 /// Converts a leaf the key-value pairs in the leaf
305 pub fn into_entries(self) -> Vec<(RpoDigest, Word)> {
306 match self {
307 SmtLeaf::Empty => Vec::new(),
308 SmtLeaf::Single(kv_pair) => vec![kv_pair],
309 SmtLeaf::Multiple(kv_pairs) => kv_pairs,
310 }
311 }
312
313 /// Computes the hash of the leaf
314 pub fn hash(&self) -> RpoDigest {
315 match self {
316 SmtLeaf::Empty => EMPTY_WORD.into(),
317 SmtLeaf::Single((key, value)) => Rpo256::merge(&[*key, value.into()]),
318 SmtLeaf::Multiple(kvs) => {
319 let elements: Vec<Felt> = kvs.iter().copied().flat_map(kv_to_elements).collect();
320 Rpo256::hash_elements(&elements)
321 }
322 }
323 }
324
325 // HELPERS
326 // ---------------------------------------------------------------------------------------------
327
328 /// Returns the value associated with `key` in the leaf
329 fn get_value(&self, key: &RpoDigest) -> Word {
330 match self {
331 SmtLeaf::Empty => EMPTY_WORD,
332 SmtLeaf::Single((key_in_leaf, value_in_leaf)) => {
333 if key == key_in_leaf {
334 *value_in_leaf
335 } else {
336 EMPTY_WORD
337 }
338 }
339 SmtLeaf::Multiple(kv_pairs) => {
340 for (key_in_leaf, value_in_leaf) in kv_pairs {
341 if key == key_in_leaf {
342 return *value_in_leaf;
343 }
344 }
345
346 EMPTY_WORD
347 }
348 }
349 }
350
351 /// Inserts key-value pair into the leaf; returns the previous value associated with `key`, if
352 /// any.
353 fn insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {
354 match self {
355 SmtLeaf::Empty => {
356 *self = SmtLeaf::Single((key, value));
357 None
358 }
359 SmtLeaf::Single(kv_pair) => {
360 if kv_pair.0 == key {
361 // the key is already in this leaf. Update the value and return the previous
362 // value
363 let old_value = kv_pair.1;
364 kv_pair.1 = value;
365 Some(old_value)
366 } else {
367 // Another entry is present in this leaf. Transform the entry into a list
368 // entry, and make sure the key-value pairs are sorted by key
369 let mut pairs = vec![*kv_pair, (key, value)];
370 pairs.sort_by(|(key_1, _), (key_2, _)| cmp_keys(*key_1, *key_2));
371
372 *self = SmtLeaf::Multiple(pairs);
373
374 None
375 }
376 }
377 SmtLeaf::Multiple(kv_pairs) => {
378 match kv_pairs.binary_search_by(|kv_pair| cmp_keys(kv_pair.0, key)) {
379 Ok(pos) => {
380 let old_value = kv_pairs[pos].1;
381 kv_pairs[pos].1 = value;
382
383 Some(old_value)
384 }
385 Err(pos) => {
386 kv_pairs.insert(pos, (key, value));
387
388 None
389 }
390 }
391 }
392 }
393 }
394
395 /// Removes key-value pair from the leaf stored at key; returns the previous value associated
396 /// with `key`, if any. Also returns an `is_empty` flag, indicating whether the leaf became
397 /// empty, and must be removed from the data structure it is contained in.
398 fn remove(&mut self, key: RpoDigest) -> (Option<Word>, bool) {
399 match self {
400 SmtLeaf::Empty => (None, false),
401 SmtLeaf::Single((key_at_leaf, value_at_leaf)) => {
402 if *key_at_leaf == key {
403 // our key was indeed stored in the leaf, so we return the value that was stored
404 // in it, and indicate that the leaf should be removed
405 let old_value = *value_at_leaf;
406
407 // Note: this is not strictly needed, since the caller is expected to drop this
408 // `SmtLeaf` object.
409 *self = SmtLeaf::Empty;
410
411 (Some(old_value), true)
412 } else {
413 // another key is stored at leaf; nothing to update
414 (None, false)
415 }
416 }
417 SmtLeaf::Multiple(kv_pairs) => {
418 match kv_pairs.binary_search_by(|kv_pair| cmp_keys(kv_pair.0, key)) {
419 Ok(pos) => {
420 let old_value = kv_pairs[pos].1;
421
422 kv_pairs.remove(pos);
423 debug_assert!(!kv_pairs.is_empty());
424
425 if kv_pairs.len() == 1 {
426 // convert the leaf into `Single`
427 *self = SmtLeaf::Single(kv_pairs[0]);
428 }
429
430 (Some(old_value), false)
431 }
432 Err(_) => {
433 // other keys are stored at leaf; nothing to update
434 (None, false)
435 }
436 }
437 }
438 }
439 }
440}
441
442// HELPER FUNCTIONS
443// ================================================================================================
444
445/// Converts a key-value tuple to an iterator of `Felt`s
446fn kv_to_elements((key, value): (RpoDigest, Word)) -> impl Iterator<Item = Felt> {
447 let key_elements = key.into_iter();
448 let value_elements = value.into_iter();
449
450 key_elements.chain(value_elements)
451}
452
453/// Compares two keys, compared element-by-element using their integer representations starting with
454/// the most significant element.
455fn cmp_keys(key_1: RpoDigest, key_2: RpoDigest) -> Ordering {
456 for (v1, v2) in key_1.iter().zip(key_2.iter()).rev() {
457 let v1 = v1.as_int();
458 let v2 = v2.as_int();
459 if v1 != v2 {
460 return v1.cmp(&v2);
461 }
462 }
463
464 Ordering::Equal
465}