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