miden-crypto/src/merkle/smt/full/leaf.rs
src/merkle/smt/full/leaf.rsRaw
1use alloc::{string::ToString, vec::Vec};
2use core::cmp::Ordering;
3
4use super::{Felt, LeafIndex, Rpo256, RpoDigest, SmtLeafError, Word, EMPTY_WORD, SMT_DEPTH};
5use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
6
7#[derive(Clone, Debug, PartialEq, Eq)]
8#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
9pub enum SmtLeaf {
10 Empty(LeafIndex<SMT_DEPTH>),
11 Single((RpoDigest, Word)),
12 Multiple(Vec<(RpoDigest, Word)>),
13}
14
15impl SmtLeaf {
16 // CONSTRUCTORS
17 // ---------------------------------------------------------------------------------------------
18
19 /// Returns a new leaf with the specified entries
20 ///
21 /// # Errors
22 /// - Returns an error if 2 keys in `entries` map to a different leaf index
23 /// - Returns an error if 1 or more keys in `entries` map to a leaf index
24 /// different from `leaf_index`
25 pub fn new(
26 entries: Vec<(RpoDigest, Word)>,
27 leaf_index: LeafIndex<SMT_DEPTH>,
28 ) -> Result<Self, SmtLeafError> {
29 match entries.len() {
30 0 => Ok(Self::new_empty(leaf_index)),
31 1 => {
32 let (key, value) = entries[0];
33
34 if LeafIndex::<SMT_DEPTH>::from(key) != leaf_index {
35 return Err(SmtLeafError::SingleKeyInconsistentWithLeafIndex {
36 key,
37 leaf_index,
38 });
39 }
40
41 Ok(Self::new_single(key, value))
42 }
43 _ => {
44 let leaf = Self::new_multiple(entries)?;
45
46 // `new_multiple()` checked that all keys map to the same leaf index. We still need
47 // to ensure that that leaf index is `leaf_index`.
48 if leaf.index() != leaf_index {
49 Err(SmtLeafError::MultipleKeysInconsistentWithLeafIndex {
50 leaf_index_from_keys: leaf.index(),
51 leaf_index_supplied: leaf_index,
52 })
53 } else {
54 Ok(leaf)
55 }
56 }
57 }
58 }
59
60 /// Returns a new empty leaf with the specified leaf index
61 pub fn new_empty(leaf_index: LeafIndex<SMT_DEPTH>) -> Self {
62 Self::Empty(leaf_index)
63 }
64
65 /// Returns a new single leaf with the specified entry. The leaf index is derived from the
66 /// entry's key.
67 pub fn new_single(key: RpoDigest, value: Word) -> Self {
68 Self::Single((key, value))
69 }
70
71 /// Returns a new single leaf with the specified entry. The leaf index is derived from the
72 /// entries' keys.
73 ///
74 /// # Errors
75 /// - Returns an error if 2 keys in `entries` map to a different leaf index
76 pub fn new_multiple(entries: Vec<(RpoDigest, Word)>) -> Result<Self, SmtLeafError> {
77 if entries.len() < 2 {
78 return Err(SmtLeafError::InvalidNumEntriesForMultiple(entries.len()));
79 }
80
81 // Check that all keys map to the same leaf index
82 {
83 let mut keys = entries.iter().map(|(key, _)| key);
84
85 let first_key = *keys.next().expect("ensured at least 2 entries");
86 let first_leaf_index: LeafIndex<SMT_DEPTH> = first_key.into();
87
88 for &next_key in keys {
89 let next_leaf_index: LeafIndex<SMT_DEPTH> = next_key.into();
90
91 if next_leaf_index != first_leaf_index {
92 return Err(SmtLeafError::InconsistentKeys {
93 entries,
94 key_1: first_key,
95 key_2: next_key,
96 });
97 }
98 }
99 }
100
101 Ok(Self::Multiple(entries))
102 }
103
104 // PUBLIC ACCESSORS
105 // ---------------------------------------------------------------------------------------------
106
107 /// Returns true if the leaf is empty
108 pub fn is_empty(&self) -> bool {
109 matches!(self, Self::Empty(_))
110 }
111
112 /// Returns the leaf's index in the [`super::Smt`]
113 pub fn index(&self) -> LeafIndex<SMT_DEPTH> {
114 match self {
115 SmtLeaf::Empty(leaf_index) => *leaf_index,
116 SmtLeaf::Single((key, _)) => key.into(),
117 SmtLeaf::Multiple(entries) => {
118 // Note: All keys are guaranteed to have the same leaf index
119 let (first_key, _) = entries[0];
120 first_key.into()
121 }
122 }
123 }
124
125 /// Returns the number of entries stored in the leaf
126 pub fn num_entries(&self) -> u64 {
127 match self {
128 SmtLeaf::Empty(_) => 0,
129 SmtLeaf::Single(_) => 1,
130 SmtLeaf::Multiple(entries) => {
131 entries.len().try_into().expect("shouldn't have more than 2^64 entries")
132 }
133 }
134 }
135
136 /// Computes the hash of the leaf
137 pub fn hash(&self) -> RpoDigest {
138 match self {
139 SmtLeaf::Empty(_) => EMPTY_WORD.into(),
140 SmtLeaf::Single((key, value)) => Rpo256::merge(&[*key, value.into()]),
141 SmtLeaf::Multiple(kvs) => {
142 let elements: Vec<Felt> = kvs.iter().copied().flat_map(kv_to_elements).collect();
143 Rpo256::hash_elements(&elements)
144 }
145 }
146 }
147
148 // ITERATORS
149 // ---------------------------------------------------------------------------------------------
150
151 /// Returns the key-value pairs in the leaf
152 pub fn entries(&self) -> Vec<&(RpoDigest, Word)> {
153 match self {
154 SmtLeaf::Empty(_) => Vec::new(),
155 SmtLeaf::Single(kv_pair) => vec![kv_pair],
156 SmtLeaf::Multiple(kv_pairs) => kv_pairs.iter().collect(),
157 }
158 }
159
160 // CONVERSIONS
161 // ---------------------------------------------------------------------------------------------
162
163 /// Converts a leaf to a list of field elements
164 pub fn to_elements(&self) -> Vec<Felt> {
165 self.clone().into_elements()
166 }
167
168 /// Converts a leaf to a list of field elements
169 pub fn into_elements(self) -> Vec<Felt> {
170 self.into_entries().into_iter().flat_map(kv_to_elements).collect()
171 }
172
173 /// Converts a leaf the key-value pairs in the leaf
174 pub fn into_entries(self) -> Vec<(RpoDigest, Word)> {
175 match self {
176 SmtLeaf::Empty(_) => Vec::new(),
177 SmtLeaf::Single(kv_pair) => vec![kv_pair],
178 SmtLeaf::Multiple(kv_pairs) => kv_pairs,
179 }
180 }
181
182 // HELPERS
183 // ---------------------------------------------------------------------------------------------
184
185 /// Returns the value associated with `key` in the leaf, or `None` if `key` maps to another leaf.
186 pub(super) fn get_value(&self, key: &RpoDigest) -> Option<Word> {
187 // Ensure that `key` maps to this leaf
188 if self.index() != key.into() {
189 return None;
190 }
191
192 match self {
193 SmtLeaf::Empty(_) => Some(EMPTY_WORD),
194 SmtLeaf::Single((key_in_leaf, value_in_leaf)) => {
195 if key == key_in_leaf {
196 Some(*value_in_leaf)
197 } else {
198 Some(EMPTY_WORD)
199 }
200 }
201 SmtLeaf::Multiple(kv_pairs) => {
202 for (key_in_leaf, value_in_leaf) in kv_pairs {
203 if key == key_in_leaf {
204 return Some(*value_in_leaf);
205 }
206 }
207
208 Some(EMPTY_WORD)
209 }
210 }
211 }
212
213 /// Inserts key-value pair into the leaf; returns the previous value associated with `key`, if
214 /// any.
215 ///
216 /// The caller needs to ensure that `key` has the same leaf index as all other keys in the leaf
217 pub(super) fn insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {
218 match self {
219 SmtLeaf::Empty(_) => {
220 *self = SmtLeaf::new_single(key, value);
221 None
222 }
223 SmtLeaf::Single(kv_pair) => {
224 if kv_pair.0 == key {
225 // the key is already in this leaf. Update the value and return the previous
226 // value
227 let old_value = kv_pair.1;
228 kv_pair.1 = value;
229 Some(old_value)
230 } else {
231 // Another entry is present in this leaf. Transform the entry into a list
232 // entry, and make sure the key-value pairs are sorted by key
233 let mut pairs = vec![*kv_pair, (key, value)];
234 pairs.sort_by(|(key_1, _), (key_2, _)| cmp_keys(*key_1, *key_2));
235
236 *self = SmtLeaf::Multiple(pairs);
237
238 None
239 }
240 }
241 SmtLeaf::Multiple(kv_pairs) => {
242 match kv_pairs.binary_search_by(|kv_pair| cmp_keys(kv_pair.0, key)) {
243 Ok(pos) => {
244 let old_value = kv_pairs[pos].1;
245 kv_pairs[pos].1 = value;
246
247 Some(old_value)
248 }
249 Err(pos) => {
250 kv_pairs.insert(pos, (key, value));
251
252 None
253 }
254 }
255 }
256 }
257 }
258
259 /// Removes key-value pair from the leaf stored at key; returns the previous value associated
260 /// with `key`, if any. Also returns an `is_empty` flag, indicating whether the leaf became
261 /// empty, and must be removed from the data structure it is contained in.
262 pub(super) fn remove(&mut self, key: RpoDigest) -> (Option<Word>, bool) {
263 match self {
264 SmtLeaf::Empty(_) => (None, false),
265 SmtLeaf::Single((key_at_leaf, value_at_leaf)) => {
266 if *key_at_leaf == key {
267 // our key was indeed stored in the leaf, so we return the value that was stored
268 // in it, and indicate that the leaf should be removed
269 let old_value = *value_at_leaf;
270
271 // Note: this is not strictly needed, since the caller is expected to drop this
272 // `SmtLeaf` object.
273 *self = SmtLeaf::new_empty(key.into());
274
275 (Some(old_value), true)
276 } else {
277 // another key is stored at leaf; nothing to update
278 (None, false)
279 }
280 }
281 SmtLeaf::Multiple(kv_pairs) => {
282 match kv_pairs.binary_search_by(|kv_pair| cmp_keys(kv_pair.0, key)) {
283 Ok(pos) => {
284 let old_value = kv_pairs[pos].1;
285
286 kv_pairs.remove(pos);
287 debug_assert!(!kv_pairs.is_empty());
288
289 if kv_pairs.len() == 1 {
290 // convert the leaf into `Single`
291 *self = SmtLeaf::Single(kv_pairs[0]);
292 }
293
294 (Some(old_value), false)
295 }
296 Err(_) => {
297 // other keys are stored at leaf; nothing to update
298 (None, false)
299 }
300 }
301 }
302 }
303 }
304}
305
306impl Serializable for SmtLeaf {
307 fn write_into<W: ByteWriter>(&self, target: &mut W) {
308 // Write: num entries
309 self.num_entries().write_into(target);
310
311 // Write: leaf index
312 let leaf_index: u64 = self.index().value();
313 leaf_index.write_into(target);
314
315 // Write: entries
316 for (key, value) in self.entries() {
317 key.write_into(target);
318 value.write_into(target);
319 }
320 }
321}
322
323impl Deserializable for SmtLeaf {
324 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
325 // Read: num entries
326 let num_entries = source.read_u64()?;
327
328 // Read: leaf index
329 let leaf_index: LeafIndex<SMT_DEPTH> = {
330 let value = source.read_u64()?;
331 LeafIndex::new_max_depth(value)
332 };
333
334 // Read: entries
335 let mut entries: Vec<(RpoDigest, Word)> = Vec::new();
336 for _ in 0..num_entries {
337 let key: RpoDigest = source.read()?;
338 let value: Word = source.read()?;
339
340 entries.push((key, value));
341 }
342
343 Self::new(entries, leaf_index)
344 .map_err(|err| DeserializationError::InvalidValue(err.to_string()))
345 }
346}
347
348// HELPER FUNCTIONS
349// ================================================================================================
350
351/// Converts a key-value tuple to an iterator of `Felt`s
352fn kv_to_elements((key, value): (RpoDigest, Word)) -> impl Iterator<Item = Felt> {
353 let key_elements = key.into_iter();
354 let value_elements = value.into_iter();
355
356 key_elements.chain(value_elements)
357}
358
359/// Compares two keys, compared element-by-element using their integer representations starting with
360/// the most significant element.
361fn cmp_keys(key_1: RpoDigest, key_2: RpoDigest) -> Ordering {
362 for (v1, v2) in key_1.iter().zip(key_2.iter()).rev() {
363 let v1 = v1.as_int();
364 let v2 = v2.as_int();
365 if v1 != v2 {
366 return v1.cmp(&v2);
367 }
368 }
369
370 Ordering::Equal
371}