miden_objects/account/delta/
storage.rs

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