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 at the specified index.
113 pub fn get_leaf(&self, key: &RpoDigest) -> SmtLeaf {
114 <Self as SparseMerkleTree<SMT_DEPTH>>::get_leaf(self, key)
115 }
116
117 /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle
118 /// path to the leaf, as well as the leaf itself.
119 pub fn open(&self, key: &RpoDigest) -> (MerklePath, SmtLeaf) {
120 <Self as SparseMerkleTree<SMT_DEPTH>>::open(self, key)
121 }
122
123 // ITERATORS
124 // --------------------------------------------------------------------------------------------
125
126 /// Returns an iterator over the leaves of this [Smt].
127 pub fn leaves(&self) -> impl Iterator<Item = (LeafIndex<SMT_DEPTH>, &SmtLeaf)> {
128 self.leaves
129 .iter()
130 .map(|(leaf_index, leaf)| (LeafIndex::new_max_depth(*leaf_index), leaf))
131 }
132
133 /// Returns an iterator over the inner nodes of this [Smt].
134 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
135 self.inner_nodes.values().map(|e| InnerNodeInfo {
136 value: e.hash(),
137 left: e.left,
138 right: e.right,
139 })
140 }
141
142 // STATE MUTATORS
143 // --------------------------------------------------------------------------------------------
144
145 /// Inserts a value at the specified key, returning the previous value associated with that key.
146 /// Recall that by definition, any key that hasn't been updated is associated with
147 /// [`Self::EMPTY_VALUE`].
148 ///
149 /// This also recomputes all hashes between the leaf (associated with the key) and the root,
150 /// updating the root itself.
151 pub fn insert(&mut self, key: RpoDigest, value: Word) -> Word {
152 <Self as SparseMerkleTree<SMT_DEPTH>>::insert(self, key, value)
153 }
154
155 // HELPERS
156 // --------------------------------------------------------------------------------------------
157
158 /// Inserts `value` at leaf index pointed to by `key`. `value` is guaranteed to not be the empty
159 /// value, such that this is indeed an insertion.
160 fn perform_insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {
161 debug_assert_ne!(value, Self::EMPTY_VALUE);
162
163 let leaf_index: LeafIndex<SMT_DEPTH> = Self::key_to_leaf_index(&key);
164
165 match self.leaves.get_mut(&leaf_index.value()) {
166 Some(leaf) => leaf.insert(key, value),
167 None => {
168 self.leaves.insert(leaf_index.value(), SmtLeaf::Single((key, value)));
169
170 None
171 }
172 }
173 }
174
175 /// Removes key-value pair at leaf index pointed to by `key` if it exists.
176 fn perform_remove(&mut self, key: RpoDigest) -> Option<Word> {
177 let leaf_index: LeafIndex<SMT_DEPTH> = Self::key_to_leaf_index(&key);
178
179 if let Some(leaf) = self.leaves.get_mut(&leaf_index.value()) {
180 let (old_value, is_empty) = leaf.remove(key);
181 if is_empty {
182 self.leaves.remove(&leaf_index.value());
183 }
184 old_value
185 } else {
186 // there's nothing stored at the leaf; nothing to update
187 None
188 }
189 }
190}
191
192impl SparseMerkleTree<SMT_DEPTH> for Smt {
193 type Key = RpoDigest;
194 type Value = Word;
195 type Leaf = SmtLeaf;
196 type Opening = (MerklePath, SmtLeaf);
197
198 const EMPTY_VALUE: Self::Value = EMPTY_WORD;
199
200 fn root(&self) -> RpoDigest {
201 self.root
202 }
203
204 fn set_root(&mut self, root: RpoDigest) {
205 self.root = root;
206 }
207
208 fn get_inner_node(&self, index: NodeIndex) -> InnerNode {
209 self.inner_nodes.get(&index).cloned().unwrap_or_else(|| {
210 let node = EmptySubtreeRoots::entry(SMT_DEPTH, index.depth() + 1);
211
212 InnerNode { left: *node, right: *node }
213 })
214 }
215
216 fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode) {
217 self.inner_nodes.insert(index, inner_node);
218 }
219
220 fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value> {
221 // inserting an `EMPTY_VALUE` is equivalent to removing any value associated with `key`
222 if value != Self::EMPTY_VALUE {
223 self.perform_insert(key, value)
224 } else {
225 self.perform_remove(key)
226 }
227 }
228
229 fn get_leaf(&self, key: &RpoDigest) -> Self::Leaf {
230 let leaf_pos = LeafIndex::<SMT_DEPTH>::from(*key).value();
231
232 match self.leaves.get(&leaf_pos) {
233 Some(leaf) => leaf.clone(),
234 None => SmtLeaf::Empty,
235 }
236 }
237
238 fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest {
239 leaf.hash()
240 }
241
242 fn key_to_leaf_index(key: &RpoDigest) -> LeafIndex<SMT_DEPTH> {
243 let most_significant_felt = key[3];
244 LeafIndex::new_max_depth(most_significant_felt.as_int())
245 }
246}
247
248impl Default for Smt {
249 fn default() -> Self {
250 Self::new()
251 }
252}
253
254// LEAF
255// ================================================================================================
256
257#[derive(Clone, Debug, PartialEq, Eq)]
258#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
259pub enum SmtLeaf {
260 Empty,
261 Single((RpoDigest, Word)),
262 Multiple(Vec<(RpoDigest, Word)>),
263}
264
265impl SmtLeaf {
266 /// Converts a leaf to a list of field elements
267 pub fn to_elements(&self) -> Vec<Felt> {
268 self.clone().into_elements()
269 }
270
271 /// Converts a leaf to a list of field elements
272 pub fn into_elements(self) -> Vec<Felt> {
273 self.into_kv_pairs().into_iter().flat_map(kv_to_elements).collect()
274 }
275
276 /// Returns the key-value pairs in the leaf
277 pub fn kv_pairs(&self) -> Vec<&(RpoDigest, Word)> {
278 match self {
279 SmtLeaf::Empty => Vec::new(),
280 SmtLeaf::Single(kv_pair) => vec![kv_pair],
281 SmtLeaf::Multiple(kv_pairs) => kv_pairs.iter().collect(),
282 }
283 }
284
285 /// Converts a leaf the key-value pairs in the leaf
286 pub fn into_kv_pairs(self) -> Vec<(RpoDigest, Word)> {
287 match self {
288 SmtLeaf::Empty => Vec::new(),
289 SmtLeaf::Single(kv_pair) => vec![kv_pair],
290 SmtLeaf::Multiple(kv_pairs) => kv_pairs,
291 }
292 }
293
294 /// Computes the hash of the leaf
295 pub fn hash(&self) -> RpoDigest {
296 match self {
297 SmtLeaf::Empty => EMPTY_WORD.into(),
298 SmtLeaf::Single((key, value)) => Rpo256::merge(&[*key, value.into()]),
299 SmtLeaf::Multiple(kvs) => {
300 let elements: Vec<Felt> = kvs.iter().copied().flat_map(kv_to_elements).collect();
301 Rpo256::hash_elements(&elements)
302 }
303 }
304 }
305
306 // HELPERS
307 // ---------------------------------------------------------------------------------------------
308
309 /// Inserts key-value pair into the leaf; returns the previous value associated with `key`, if
310 /// any.
311 fn insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {
312 match self {
313 SmtLeaf::Empty => {
314 *self = SmtLeaf::Single((key, value));
315 None
316 }
317 SmtLeaf::Single(kv_pair) => {
318 if kv_pair.0 == key {
319 // the key is already in this leaf. Update the value and return the previous
320 // value
321 let old_value = kv_pair.1;
322 kv_pair.1 = value;
323 Some(old_value)
324 } else {
325 // Another entry is present in this leaf. Transform the entry into a list
326 // entry, and make sure the key-value pairs are sorted by key
327 let mut pairs = vec![*kv_pair, (key, value)];
328 pairs.sort_by(|(key_1, _), (key_2, _)| cmp_keys(*key_1, *key_2));
329
330 *self = SmtLeaf::Multiple(pairs);
331
332 None
333 }
334 }
335 SmtLeaf::Multiple(kv_pairs) => {
336 match kv_pairs.binary_search_by(|kv_pair| cmp_keys(kv_pair.0, key)) {
337 Ok(pos) => {
338 let old_value = kv_pairs[pos].1;
339 kv_pairs[pos].1 = value;
340
341 Some(old_value)
342 }
343 Err(pos) => {
344 kv_pairs.insert(pos, (key, value));
345
346 None
347 }
348 }
349 }
350 }
351 }
352
353 /// Removes key-value pair from the leaf stored at key; returns the previous value associated
354 /// with `key`, if any. Also returns an `is_empty` flag, indicating whether the leaf became
355 /// empty, and must be removed from the data structure it is contained in.
356 fn remove(&mut self, key: RpoDigest) -> (Option<Word>, bool) {
357 match self {
358 SmtLeaf::Empty => (None, false),
359 SmtLeaf::Single((key_at_leaf, value_at_leaf)) => {
360 if *key_at_leaf == key {
361 // our key was indeed stored in the leaf, so we return the value that was stored
362 // in it, and indicate that the leaf should be removed
363 let old_value = *value_at_leaf;
364
365 // Note: this is not strictly needed, since the caller is expected to drop this
366 // `SmtLeaf` object.
367 *self = SmtLeaf::Empty;
368
369 (Some(old_value), true)
370 } else {
371 // another key is stored at leaf; nothing to update
372 (None, false)
373 }
374 }
375 SmtLeaf::Multiple(kv_pairs) => {
376 match kv_pairs.binary_search_by(|kv_pair| cmp_keys(kv_pair.0, key)) {
377 Ok(pos) => {
378 let old_value = kv_pairs[pos].1;
379
380 kv_pairs.remove(pos);
381 debug_assert!(!kv_pairs.is_empty());
382
383 if kv_pairs.len() == 1 {
384 // convert the leaf into `Single`
385 *self = SmtLeaf::Single(kv_pairs[0]);
386 }
387
388 (Some(old_value), false)
389 }
390 Err(_) => {
391 // other keys are stored at leaf; nothing to update
392 (None, false)
393 }
394 }
395 }
396 }
397 }
398}
399
400// HELPER FUNCTIONS
401// ================================================================================================
402
403/// Converts a key-value tuple to an iterator of `Felt`s
404fn kv_to_elements((key, value): (RpoDigest, Word)) -> impl Iterator<Item = Felt> {
405 let key_elements = key.into_iter();
406 let value_elements = value.into_iter();
407
408 key_elements.chain(value_elements)
409}
410
411/// Compares two keys, compared element-by-element using their integer representations starting with
412/// the most significant element.
413fn cmp_keys(key_1: RpoDigest, key_2: RpoDigest) -> Ordering {
414 for (v1, v2) in key_1.iter().zip(key_2.iter()).rev() {
415 let v1 = v1.as_int();
416 let v2 = v2.as_int();
417 if v1 != v2 {
418 return v1.cmp(&v2);
419 }
420 }
421
422 Ordering::Equal
423}