miden_objects/account/delta/
storage.rs

1use alloc::{
2    collections::{BTreeMap, btree_map::Entry},
3    string::ToString,
4    vec::Vec,
5};
6
7use super::{
8    AccountDeltaError, ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
9    Word,
10};
11use crate::{
12    Digest, EMPTY_WORD,
13    account::{AccountStorage, StorageMap, StorageSlot},
14    crypto::merkle::SmtLeaf,
15};
16// ACCOUNT STORAGE DELTA
17// ================================================================================================
18
19/// [AccountStorageDelta] stores the differences between two states of account storage.
20///
21/// The delta consists of two maps:
22/// - A map containing the updates to value storage slots. The keys in this map are indexes of the
23///   updated storage slots and the values are the new values for these slots.
24/// - A map containing updates to storage maps. The keys in this map are indexes of the updated
25///   storage slots and the values are corresponding storage map delta objects.
26#[derive(Clone, Debug, Default, PartialEq, Eq)]
27pub struct AccountStorageDelta {
28    values: BTreeMap<u8, Word>,
29    maps: BTreeMap<u8, StorageMapDelta>,
30}
31
32impl AccountStorageDelta {
33    /// Creates a new storage delta from the provided fields.
34    pub fn new(
35        values: BTreeMap<u8, Word>,
36        maps: BTreeMap<u8, StorageMapDelta>,
37    ) -> Result<Self, AccountDeltaError> {
38        let result = Self { values, maps };
39        result.validate()?;
40
41        Ok(result)
42    }
43
44    /// Returns a reference to the updated values in this storage delta.
45    pub fn values(&self) -> &BTreeMap<u8, Word> {
46        &self.values
47    }
48
49    /// Returns a reference to the updated maps in this storage delta.
50    pub fn maps(&self) -> &BTreeMap<u8, StorageMapDelta> {
51        &self.maps
52    }
53
54    /// Returns true if storage delta contains no updates.
55    pub fn is_empty(&self) -> bool {
56        self.values.is_empty() && self.maps.is_empty()
57    }
58
59    /// Tracks a slot change
60    pub fn set_item(&mut self, slot_index: u8, new_slot_value: Word) {
61        self.values.insert(slot_index, new_slot_value);
62    }
63
64    /// Tracks a map item change
65    pub fn set_map_item(&mut self, slot_index: u8, key: Digest, new_value: Word) {
66        self.maps.entry(slot_index).or_default().insert(key, new_value);
67    }
68
69    /// Merges another delta into this one, overwriting any existing values.
70    pub fn merge(&mut self, other: Self) -> Result<(), AccountDeltaError> {
71        self.values.extend(other.values);
72
73        // merge maps
74        for (slot, update) in other.maps.into_iter() {
75            match self.maps.entry(slot) {
76                Entry::Vacant(entry) => {
77                    entry.insert(update);
78                },
79                Entry::Occupied(mut entry) => entry.get_mut().merge(update),
80            }
81        }
82
83        self.validate()
84    }
85
86    /// Checks whether this storage delta is valid.
87    ///
88    /// # Errors:
89    /// - Any of the updated slot is referenced from both maps (e.g., updated twice).
90    fn validate(&self) -> Result<(), AccountDeltaError> {
91        for slot in self.maps.keys() {
92            if self.values.contains_key(slot) {
93                return Err(AccountDeltaError::StorageSlotUsedAsDifferentTypes(*slot));
94            }
95        }
96
97        Ok(())
98    }
99
100    /// Returns an iterator of all the cleared storage slots.
101    fn cleared_slots(&self) -> impl Iterator<Item = u8> + '_ {
102        self.values
103            .iter()
104            .filter(|&(_, value)| (value == &EMPTY_WORD))
105            .map(|(slot, _)| *slot)
106    }
107
108    /// Returns an iterator of all the updated storage slots.
109    fn updated_slots(&self) -> impl Iterator<Item = (&u8, &Word)> + '_ {
110        self.values.iter().filter(|&(_, value)| value != &EMPTY_WORD)
111    }
112}
113
114#[cfg(any(feature = "testing", test))]
115impl AccountStorageDelta {
116    /// Creates an [AccountStorageDelta] from the given iterators.
117    pub fn from_iters(
118        cleared_items: impl IntoIterator<Item = u8>,
119        updated_values: impl IntoIterator<Item = (u8, Word)>,
120        updated_maps: impl IntoIterator<Item = (u8, StorageMapDelta)>,
121    ) -> Self {
122        Self {
123            values: BTreeMap::from_iter(
124                cleared_items.into_iter().map(|key| (key, EMPTY_WORD)).chain(updated_values),
125            ),
126            maps: BTreeMap::from_iter(updated_maps),
127        }
128    }
129}
130
131/// Converts an [AccountStorage] into an [AccountStorageDelta] for initial delta construction.
132impl From<AccountStorage> for AccountStorageDelta {
133    fn from(storage: AccountStorage) -> Self {
134        let mut values = BTreeMap::new();
135        let mut maps = BTreeMap::new();
136        for (slot_idx, slot) in storage.into_iter().enumerate() {
137            let slot_idx: u8 = slot_idx.try_into().expect("slot index must fit into `u8`");
138            match slot {
139                StorageSlot::Value(value) => {
140                    values.insert(slot_idx, value);
141                },
142                StorageSlot::Map(map) => {
143                    maps.insert(slot_idx, map.into());
144                },
145            }
146        }
147
148        Self { values, maps }
149    }
150}
151
152impl Serializable for AccountStorageDelta {
153    fn write_into<W: ByteWriter>(&self, target: &mut W) {
154        let cleared: Vec<u8> = self.cleared_slots().collect();
155        let updated: Vec<(&u8, &Word)> = self.updated_slots().collect();
156
157        target.write_u8(cleared.len() as u8);
158        target.write_many(cleared.iter());
159
160        target.write_u8(updated.len() as u8);
161        target.write_many(updated.iter());
162
163        target.write_u8(self.maps.len() as u8);
164        target.write_many(self.maps.iter());
165    }
166
167    fn get_size_hint(&self) -> usize {
168        let u8_size = 0u8.get_size_hint();
169        let word_size = EMPTY_WORD.get_size_hint();
170
171        let mut storage_map_delta_size = 0;
172        for (slot, storage_map_delta) in self.maps.iter() {
173            // The serialized size of each entry is the combination of slot (key) and the delta
174            // (value).
175            storage_map_delta_size += slot.get_size_hint() + storage_map_delta.get_size_hint();
176        }
177
178        // Length Prefixes
179        u8_size * 3 +
180        // Cleared Slots
181        self.cleared_slots().count() * u8_size +
182        // Updated Slots
183        self.updated_slots().count() * (u8_size + word_size) +
184        // Storage Map Delta
185        storage_map_delta_size
186    }
187}
188
189impl Deserializable for AccountStorageDelta {
190    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
191        let mut values = BTreeMap::new();
192
193        let num_cleared_items = source.read_u8()? as usize;
194        for _ in 0..num_cleared_items {
195            let cleared_slot = source.read_u8()?;
196            values.insert(cleared_slot, EMPTY_WORD);
197        }
198
199        let num_updated_items = source.read_u8()? as usize;
200        for _ in 0..num_updated_items {
201            let (updated_slot, updated_value) = source.read()?;
202            values.insert(updated_slot, updated_value);
203        }
204
205        let num_maps = source.read_u8()? as usize;
206        let maps = source.read_many::<(u8, StorageMapDelta)>(num_maps)?.into_iter().collect();
207
208        Self::new(values, maps).map_err(|err| DeserializationError::InvalidValue(err.to_string()))
209    }
210}
211
212// STORAGE MAP DELTA
213// ================================================================================================
214
215/// [StorageMapDelta] stores the differences between two states of account storage maps.
216///
217/// The differences are represented as leaf updates: a map of updated item key ([Digest]) to
218/// value ([Word]). For cleared items the value is [EMPTY_WORD].
219#[derive(Clone, Debug, Default, PartialEq, Eq)]
220pub struct StorageMapDelta(BTreeMap<Digest, Word>);
221
222impl StorageMapDelta {
223    /// Creates a new storage map delta from the provided leaves.
224    pub fn new(map: BTreeMap<Digest, Word>) -> Self {
225        Self(map)
226    }
227
228    /// Returns a reference to the updated leaves in this storage map delta.
229    pub fn leaves(&self) -> &BTreeMap<Digest, Word> {
230        &self.0
231    }
232
233    /// Inserts an item into the storage map delta.
234    pub fn insert(&mut self, key: Digest, value: Word) {
235        self.0.insert(key, value);
236    }
237
238    /// Returns true if storage map delta contains no updates.
239    pub fn is_empty(&self) -> bool {
240        self.0.is_empty()
241    }
242
243    /// Merge `other` into this delta, giving precedence to `other`.
244    pub fn merge(&mut self, other: Self) {
245        // Aggregate the changes into a map such that `other` overwrites self.
246        self.0.extend(other.0);
247    }
248
249    /// Returns an iterator of all the cleared keys in the storage map.
250    fn cleared_keys(&self) -> impl Iterator<Item = &Digest> + '_ {
251        self.0.iter().filter(|&(_, value)| value == &EMPTY_WORD).map(|(key, _)| key)
252    }
253
254    /// Returns an iterator of all the updated entries in the storage map.
255    fn updated_entries(&self) -> impl Iterator<Item = (&Digest, &Word)> + '_ {
256        self.0.iter().filter(|&(_, value)| value != &EMPTY_WORD)
257    }
258}
259
260#[cfg(any(feature = "testing", test))]
261impl StorageMapDelta {
262    /// Creates a new [StorageMapDelta] from the provided iterators.
263    pub fn from_iters(
264        cleared_leaves: impl IntoIterator<Item = Word>,
265        updated_leaves: impl IntoIterator<Item = (Word, Word)>,
266    ) -> Self {
267        Self(BTreeMap::from_iter(
268            cleared_leaves
269                .into_iter()
270                .map(|key| (key.into(), EMPTY_WORD))
271                .chain(updated_leaves.into_iter().map(|(key, value)| (key.into(), value))),
272        ))
273    }
274}
275
276/// Converts a [StorageMap] into a [StorageMapDelta] for initial delta construction.
277impl From<StorageMap> for StorageMapDelta {
278    fn from(map: StorageMap) -> Self {
279        // TODO: implement `IntoIterator` for `Smt` and get rid of copying keys/values:
280        let mut delta = StorageMapDelta::new(BTreeMap::default());
281        for (_, leaf) in map.leaves() {
282            match leaf {
283                SmtLeaf::Empty(_) => continue,
284                SmtLeaf::Single((key, value)) => {
285                    delta.insert(*key, *value);
286                },
287                SmtLeaf::Multiple(vec) => {
288                    vec.iter().for_each(|(key, value)| delta.insert(*key, *value));
289                },
290            }
291        }
292
293        delta
294    }
295}
296
297impl Serializable for StorageMapDelta {
298    fn write_into<W: ByteWriter>(&self, target: &mut W) {
299        let cleared: Vec<&Digest> = self.cleared_keys().collect();
300        let updated: Vec<(&Digest, &Word)> = self.updated_entries().collect();
301
302        target.write_usize(cleared.len());
303        target.write_many(cleared.iter());
304
305        target.write_usize(updated.len());
306        target.write_many(updated.iter());
307    }
308
309    fn get_size_hint(&self) -> usize {
310        let word_size = EMPTY_WORD.get_size_hint();
311
312        let cleared_keys_count = self.cleared_keys().count();
313        let updated_entries_count = self.updated_entries().count();
314
315        // Cleared Keys
316        cleared_keys_count.get_size_hint() +
317        cleared_keys_count * Digest::SERIALIZED_SIZE +
318
319        // Updated Entries
320        updated_entries_count.get_size_hint() +
321        updated_entries_count * (Digest::SERIALIZED_SIZE + word_size)
322    }
323}
324
325impl Deserializable for StorageMapDelta {
326    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
327        let mut map = BTreeMap::new();
328
329        let cleared_count = source.read_usize()?;
330        for _ in 0..cleared_count {
331            let cleared_key = source.read()?;
332            map.insert(cleared_key, EMPTY_WORD);
333        }
334
335        let updated_count = source.read_usize()?;
336        for _ in 0..updated_count {
337            let (updated_key, updated_value) = source.read()?;
338            map.insert(updated_key, updated_value);
339        }
340
341        Ok(Self::new(map))
342    }
343}
344
345// TESTS
346// ================================================================================================
347
348#[cfg(test)]
349mod tests {
350    use super::{AccountStorageDelta, Deserializable, Serializable};
351    use crate::{
352        ONE, ZERO, account::StorageMapDelta, testing::storage::AccountStorageDeltaBuilder,
353    };
354
355    #[test]
356    fn account_storage_delta_validation() {
357        let delta = AccountStorageDelta::from_iters(
358            [1, 2, 3],
359            [(4, [ONE, ONE, ONE, ONE]), (5, [ONE, ONE, ONE, ZERO])],
360            [],
361        );
362        assert!(delta.validate().is_ok());
363
364        let bytes = delta.to_bytes();
365        assert_eq!(AccountStorageDelta::read_from_bytes(&bytes), Ok(delta));
366
367        // duplicate across cleared items and maps
368        let delta = AccountStorageDelta::from_iters(
369            [1, 2, 3],
370            [(2, [ONE, ONE, ONE, ONE]), (5, [ONE, ONE, ONE, ZERO])],
371            [(1, StorageMapDelta::default())],
372        );
373        assert!(delta.validate().is_err());
374
375        let bytes = delta.to_bytes();
376        assert!(AccountStorageDelta::read_from_bytes(&bytes).is_err());
377
378        // duplicate across updated items and maps
379        let delta = AccountStorageDelta::from_iters(
380            [1, 3],
381            [(2, [ONE, ONE, ONE, ONE]), (5, [ONE, ONE, ONE, ZERO])],
382            [(2, StorageMapDelta::default())],
383        );
384        assert!(delta.validate().is_err());
385
386        let bytes = delta.to_bytes();
387        assert!(AccountStorageDelta::read_from_bytes(&bytes).is_err());
388    }
389
390    #[test]
391    fn test_is_empty() {
392        let storage_delta = AccountStorageDelta::default();
393        assert!(storage_delta.is_empty());
394
395        let storage_delta = AccountStorageDelta::from_iters([1], [], []);
396        assert!(!storage_delta.is_empty());
397
398        let storage_delta = AccountStorageDelta::from_iters([], [(2, [ONE, ONE, ONE, ONE])], []);
399        assert!(!storage_delta.is_empty());
400
401        let storage_delta =
402            AccountStorageDelta::from_iters([], [], [(3, StorageMapDelta::default())]);
403        assert!(!storage_delta.is_empty());
404    }
405
406    #[test]
407    fn test_serde_account_storage_delta() {
408        let storage_delta = AccountStorageDelta::default();
409        let serialized = storage_delta.to_bytes();
410        let deserialized = AccountStorageDelta::read_from_bytes(&serialized).unwrap();
411        assert_eq!(deserialized, storage_delta);
412
413        let storage_delta = AccountStorageDelta::from_iters([1], [], []);
414        let serialized = storage_delta.to_bytes();
415        let deserialized = AccountStorageDelta::read_from_bytes(&serialized).unwrap();
416        assert_eq!(deserialized, storage_delta);
417
418        let storage_delta = AccountStorageDelta::from_iters([], [(2, [ONE, ONE, ONE, ONE])], []);
419        let serialized = storage_delta.to_bytes();
420        let deserialized = AccountStorageDelta::read_from_bytes(&serialized).unwrap();
421        assert_eq!(deserialized, storage_delta);
422
423        let storage_delta =
424            AccountStorageDelta::from_iters([], [], [(3, StorageMapDelta::default())]);
425        let serialized = storage_delta.to_bytes();
426        let deserialized = AccountStorageDelta::read_from_bytes(&serialized).unwrap();
427        assert_eq!(deserialized, storage_delta);
428    }
429
430    #[test]
431    fn test_serde_storage_map_delta() {
432        let storage_map_delta = StorageMapDelta::default();
433        let serialized = storage_map_delta.to_bytes();
434        let deserialized = StorageMapDelta::read_from_bytes(&serialized).unwrap();
435        assert_eq!(deserialized, storage_map_delta);
436
437        let storage_map_delta = StorageMapDelta::from_iters([[ONE, ONE, ONE, ONE]], []);
438        let serialized = storage_map_delta.to_bytes();
439        let deserialized = StorageMapDelta::read_from_bytes(&serialized).unwrap();
440        assert_eq!(deserialized, storage_map_delta);
441
442        let storage_map_delta =
443            StorageMapDelta::from_iters([], [([ZERO, ZERO, ZERO, ZERO], [ONE, ONE, ONE, ONE])]);
444        let serialized = storage_map_delta.to_bytes();
445        let deserialized = StorageMapDelta::read_from_bytes(&serialized).unwrap();
446        assert_eq!(deserialized, storage_map_delta);
447    }
448
449    #[rstest::rstest]
450    #[case::some_some(Some(1), Some(2), Some(2))]
451    #[case::none_some(None, Some(2), Some(2))]
452    #[case::some_none(Some(1), None, None)]
453    #[test]
454    fn merge_items(#[case] x: Option<u64>, #[case] y: Option<u64>, #[case] expected: Option<u64>) {
455        /// Creates a delta containing the item as an update if Some, else with the item cleared.
456        fn create_delta(item: Option<u64>) -> AccountStorageDelta {
457            const SLOT: u8 = 123;
458            let item = item.map(|x| (SLOT, [vm_core::Felt::new(x), ZERO, ZERO, ZERO]));
459
460            AccountStorageDeltaBuilder::default()
461                .add_cleared_items(item.is_none().then_some(SLOT))
462                .add_updated_values(item)
463                .build()
464                .unwrap()
465        }
466
467        let mut delta_x = create_delta(x);
468        let delta_y = create_delta(y);
469        let expected = create_delta(expected);
470
471        delta_x.merge(delta_y).unwrap();
472
473        assert_eq!(delta_x, expected);
474    }
475
476    #[rstest::rstest]
477    #[case::some_some(Some(1), Some(2), Some(2))]
478    #[case::none_some(None, Some(2), Some(2))]
479    #[case::some_none(Some(1), None, None)]
480    #[test]
481    fn merge_maps(#[case] x: Option<u64>, #[case] y: Option<u64>, #[case] expected: Option<u64>) {
482        fn create_delta(value: Option<u64>) -> StorageMapDelta {
483            let key = [vm_core::Felt::new(10), ZERO, ZERO, ZERO];
484            match value {
485                Some(value) => StorageMapDelta::from_iters(
486                    [],
487                    [(key, [vm_core::Felt::new(value), ZERO, ZERO, ZERO])],
488                ),
489                None => StorageMapDelta::from_iters([key], []),
490            }
491        }
492
493        let mut delta_x = create_delta(x);
494        let delta_y = create_delta(y);
495        let expected = create_delta(expected);
496
497        delta_x.merge(delta_y);
498
499        assert_eq!(delta_x, expected);
500    }
501}