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