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 insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value> {
236 // inserting an `EMPTY_VALUE` is equivalent to removing any value associated with `key`
237 if value != Self::EMPTY_VALUE {
238 self.perform_insert(key, value)
239 } else {
240 self.perform_remove(key)
241 }
242 }
243
244 fn get_leaf(&self, key: &RpoDigest) -> Self::Leaf {
245 let leaf_pos = LeafIndex::<SMT_DEPTH>::from(*key).value();
246
247 match self.leaves.get(&leaf_pos) {
248 Some(leaf) => leaf.clone(),
249 None => SmtLeaf::Empty,
250 }
251 }
252
253 fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest {
254 leaf.hash()
255 }
256
257 fn key_to_leaf_index(key: &RpoDigest) -> LeafIndex<SMT_DEPTH> {
258 let most_significant_felt = key[3];
259 LeafIndex::new_max_depth(most_significant_felt.as_int())
260 }
261}
262
263impl Default for Smt {
264 fn default() -> Self {
265 Self::new()
266 }
267}
268
269// LEAF
270// ================================================================================================
271
272#[derive(Clone, Debug, PartialEq, Eq)]
273#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
274pub enum SmtLeaf {
275 Empty,
276 Single((RpoDigest, Word)),
277 Multiple(Vec<(RpoDigest, Word)>),
278}
279
280impl SmtLeaf {
281 /// Converts a leaf to a list of field elements
282 pub fn to_elements(&self) -> Vec<Felt> {
283 self.clone().into_elements()
284 }
285
286 /// Converts a leaf to a list of field elements
287 pub fn into_elements(self) -> Vec<Felt> {
288 self.into_entries().into_iter().flat_map(kv_to_elements).collect()
289 }
290
291 /// Returns the key-value pairs in the leaf
292 pub fn entries(&self) -> Vec<&(RpoDigest, Word)> {
293 match self {
294 SmtLeaf::Empty => Vec::new(),
295 SmtLeaf::Single(kv_pair) => vec![kv_pair],
296 SmtLeaf::Multiple(kv_pairs) => kv_pairs.iter().collect(),
297 }
298 }
299
300 /// Converts a leaf the key-value pairs in the leaf
301 pub fn into_entries(self) -> Vec<(RpoDigest, Word)> {
302 match self {
303 SmtLeaf::Empty => Vec::new(),
304 SmtLeaf::Single(kv_pair) => vec![kv_pair],
305 SmtLeaf::Multiple(kv_pairs) => kv_pairs,
306 }
307 }
308
309 /// Computes the hash of the leaf
310 pub fn hash(&self) -> RpoDigest {
311 match self {
312 SmtLeaf::Empty => EMPTY_WORD.into(),
313 SmtLeaf::Single((key, value)) => Rpo256::merge(&[*key, value.into()]),
314 SmtLeaf::Multiple(kvs) => {
315 let elements: Vec<Felt> = kvs.iter().copied().flat_map(kv_to_elements).collect();
316 Rpo256::hash_elements(&elements)
317 }
318 }
319 }
320
321 // HELPERS
322 // ---------------------------------------------------------------------------------------------
323
324 /// Returns the value associated with `key` in the leaf
325 fn get_value(&self, key: &RpoDigest) -> Word {
326 match self {
327 SmtLeaf::Empty => EMPTY_WORD,
328 SmtLeaf::Single((key_in_leaf, value_in_leaf)) => {
329 if key == key_in_leaf {
330 *value_in_leaf
331 } else {
332 EMPTY_WORD
333 }
334 }
335 SmtLeaf::Multiple(kv_pairs) => {
336 for (key_in_leaf, value_in_leaf) in kv_pairs {
337 if key == key_in_leaf {
338 return *value_in_leaf;
339 }
340 }
341
342 EMPTY_WORD
343 }
344 }
345 }
346
347 /// Inserts key-value pair into the leaf; returns the previous value associated with `key`, if
348 /// any.
349 fn insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {
350 match self {
351 SmtLeaf::Empty => {
352 *self = SmtLeaf::Single((key, value));
353 None
354 }
355 SmtLeaf::Single(kv_pair) => {
356 if kv_pair.0 == key {
357 // the key is already in this leaf. Update the value and return the previous
358 // value
359 let old_value = kv_pair.1;
360 kv_pair.1 = value;
361 Some(old_value)
362 } else {
363 // Another entry is present in this leaf. Transform the entry into a list
364 // entry, and make sure the key-value pairs are sorted by key
365 let mut pairs = vec![*kv_pair, (key, value)];
366 pairs.sort_by(|(key_1, _), (key_2, _)| cmp_keys(*key_1, *key_2));
367
368 *self = SmtLeaf::Multiple(pairs);
369
370 None
371 }
372 }
373 SmtLeaf::Multiple(kv_pairs) => {
374 match kv_pairs.binary_search_by(|kv_pair| cmp_keys(kv_pair.0, key)) {
375 Ok(pos) => {
376 let old_value = kv_pairs[pos].1;
377 kv_pairs[pos].1 = value;
378
379 Some(old_value)
380 }
381 Err(pos) => {
382 kv_pairs.insert(pos, (key, value));
383
384 None
385 }
386 }
387 }
388 }
389 }
390
391 /// Removes key-value pair from the leaf stored at key; returns the previous value associated
392 /// with `key`, if any. Also returns an `is_empty` flag, indicating whether the leaf became
393 /// empty, and must be removed from the data structure it is contained in.
394 fn remove(&mut self, key: RpoDigest) -> (Option<Word>, bool) {
395 match self {
396 SmtLeaf::Empty => (None, false),
397 SmtLeaf::Single((key_at_leaf, value_at_leaf)) => {
398 if *key_at_leaf == key {
399 // our key was indeed stored in the leaf, so we return the value that was stored
400 // in it, and indicate that the leaf should be removed
401 let old_value = *value_at_leaf;
402
403 // Note: this is not strictly needed, since the caller is expected to drop this
404 // `SmtLeaf` object.
405 *self = SmtLeaf::Empty;
406
407 (Some(old_value), true)
408 } else {
409 // another key is stored at leaf; nothing to update
410 (None, false)
411 }
412 }
413 SmtLeaf::Multiple(kv_pairs) => {
414 match kv_pairs.binary_search_by(|kv_pair| cmp_keys(kv_pair.0, key)) {
415 Ok(pos) => {
416 let old_value = kv_pairs[pos].1;
417
418 kv_pairs.remove(pos);
419 debug_assert!(!kv_pairs.is_empty());
420
421 if kv_pairs.len() == 1 {
422 // convert the leaf into `Single`
423 *self = SmtLeaf::Single(kv_pairs[0]);
424 }
425
426 (Some(old_value), false)
427 }
428 Err(_) => {
429 // other keys are stored at leaf; nothing to update
430 (None, false)
431 }
432 }
433 }
434 }
435 }
436}
437
438// HELPER FUNCTIONS
439// ================================================================================================
440
441/// Converts a key-value tuple to an iterator of `Felt`s
442fn kv_to_elements((key, value): (RpoDigest, Word)) -> impl Iterator<Item = Felt> {
443 let key_elements = key.into_iter();
444 let value_elements = value.into_iter();
445
446 key_elements.chain(value_elements)
447}
448
449/// Compares two keys, compared element-by-element using their integer representations starting with
450/// the most significant element.
451fn cmp_keys(key_1: RpoDigest, key_2: RpoDigest) -> Ordering {
452 for (v1, v2) in key_1.iter().zip(key_2.iter()).rev() {
453 let v1 = v1.as_int();
454 let v2 = v2.as_int();
455 if v1 != v2 {
456 return v1.cmp(&v2);
457 }
458 }
459
460 Ordering::Equal
461}