1use alloc::collections::BTreeMap;
2use alloc::collections::btree_map::Entry;
3use alloc::string::ToString;
4use alloc::vec::Vec;
5
6use super::{
7    AccountDeltaError,
8    ByteReader,
9    ByteWriter,
10    Deserializable,
11    DeserializationError,
12    Serializable,
13    Word,
14};
15use crate::account::StorageMap;
16use crate::{EMPTY_WORD, Felt, LexicographicWord, ZERO};
17
18#[derive(Clone, Debug, PartialEq, Eq)]
29pub struct AccountStorageDelta {
30    values: BTreeMap<u8, Word>,
32    maps: BTreeMap<u8, StorageMapDelta>,
34}
35
36impl AccountStorageDelta {
37    pub fn new() -> Self {
39        Self {
40            values: BTreeMap::new(),
41            maps: BTreeMap::new(),
42        }
43    }
44
45    pub fn from_parts(
53        values: BTreeMap<u8, Word>,
54        maps: BTreeMap<u8, StorageMapDelta>,
55    ) -> Result<Self, AccountDeltaError> {
56        let delta = Self { values, maps };
57        delta.validate()?;
58
59        Ok(delta)
60    }
61
62    pub fn values(&self) -> &BTreeMap<u8, Word> {
64        &self.values
65    }
66
67    pub fn maps(&self) -> &BTreeMap<u8, StorageMapDelta> {
69        &self.maps
70    }
71
72    pub fn is_empty(&self) -> bool {
74        self.values.is_empty() && self.maps.is_empty()
75    }
76
77    pub fn set_item(&mut self, slot_index: u8, new_slot_value: Word) {
79        self.values.insert(slot_index, new_slot_value);
80    }
81
82    pub fn set_map_item(&mut self, slot_index: u8, key: Word, new_value: Word) {
84        self.maps.entry(slot_index).or_default().insert(key, new_value);
85    }
86
87    pub fn merge(&mut self, other: Self) -> Result<(), AccountDeltaError> {
89        self.values.extend(other.values);
90
91        for (slot, update) in other.maps.into_iter() {
93            match self.maps.entry(slot) {
94                Entry::Vacant(entry) => {
95                    entry.insert(update);
96                },
97                Entry::Occupied(mut entry) => entry.get_mut().merge(update),
98            }
99        }
100
101        self.validate()
102    }
103
104    fn validate(&self) -> Result<(), AccountDeltaError> {
112        for slot in self.maps.keys() {
113            if self.values.contains_key(slot) {
114                return Err(AccountDeltaError::StorageSlotUsedAsDifferentTypes(*slot));
115            }
116        }
117
118        Ok(())
119    }
120
121    fn cleared_slots(&self) -> impl Iterator<Item = u8> + '_ {
123        self.values.iter().filter(|&(_, value)| value.is_empty()).map(|(slot, _)| *slot)
124    }
125
126    fn updated_slots(&self) -> impl Iterator<Item = (&u8, &Word)> + '_ {
128        self.values.iter().filter(|&(_, value)| !value.is_empty())
129    }
130
131    pub(super) fn append_delta_elements(&self, elements: &mut Vec<Felt>) {
134        const DOMAIN_VALUE: Felt = Felt::new(2);
135        const DOMAIN_MAP: Felt = Felt::new(3);
136
137        let highest_value_slot_idx = self.values.last_key_value().map(|(slot_idx, _)| slot_idx);
138        let highest_map_slot_idx = self.maps.last_key_value().map(|(slot_idx, _)| slot_idx);
139        let highest_slot_idx =
140            highest_value_slot_idx.max(highest_map_slot_idx).copied().unwrap_or(0);
141
142        for slot_idx in 0..=highest_slot_idx {
143            let slot_idx_felt = Felt::from(slot_idx);
144
145            match self.values.get(&slot_idx) {
148                Some(new_value) => {
149                    elements.extend_from_slice(&[DOMAIN_VALUE, slot_idx_felt, ZERO, ZERO]);
150                    elements.extend_from_slice(new_value.as_elements());
151                },
152                None => {
153                    if let Some(map_delta) = self.maps().get(&slot_idx) {
154                        if map_delta.is_empty() {
155                            continue;
156                        }
157
158                        for (key, value) in map_delta.entries() {
159                            elements.extend_from_slice(key.inner().as_elements());
160                            elements.extend_from_slice(value.as_elements());
161                        }
162
163                        let num_changed_entries = Felt::try_from(map_delta.num_entries()).expect(
164                            "number of changed entries should not exceed max representable felt",
165                        );
166
167                        elements.extend_from_slice(&[
168                            DOMAIN_MAP,
169                            slot_idx_felt,
170                            num_changed_entries,
171                            ZERO,
172                        ]);
173                        elements.extend_from_slice(EMPTY_WORD.as_elements());
174                    }
175                },
176            }
177        }
178    }
179
180    pub fn into_parts(self) -> (BTreeMap<u8, Word>, BTreeMap<u8, StorageMapDelta>) {
182        (self.values, self.maps)
183    }
184}
185
186impl Default for AccountStorageDelta {
187    fn default() -> Self {
188        Self::new()
189    }
190}
191
192#[cfg(any(feature = "testing", test))]
193impl AccountStorageDelta {
194    pub fn from_iters(
196        cleared_values: impl IntoIterator<Item = u8>,
197        updated_values: impl IntoIterator<Item = (u8, Word)>,
198        updated_maps: impl IntoIterator<Item = (u8, StorageMapDelta)>,
199    ) -> Self {
200        Self {
201            values: BTreeMap::from_iter(
202                cleared_values.into_iter().map(|key| (key, EMPTY_WORD)).chain(updated_values),
203            ),
204            maps: BTreeMap::from_iter(updated_maps),
205        }
206    }
207}
208
209impl Serializable for AccountStorageDelta {
210    fn write_into<W: ByteWriter>(&self, target: &mut W) {
211        let cleared: Vec<u8> = self.cleared_slots().collect();
212        let updated: Vec<(&u8, &Word)> = self.updated_slots().collect();
213
214        target.write_u8(cleared.len() as u8);
215        target.write_many(cleared.iter());
216
217        target.write_u8(updated.len() as u8);
218        target.write_many(updated.iter());
219
220        target.write_u8(self.maps.len() as u8);
221        target.write_many(self.maps.iter());
222    }
223
224    fn get_size_hint(&self) -> usize {
225        let u8_size = 0u8.get_size_hint();
226        let word_size = EMPTY_WORD.get_size_hint();
227
228        let mut storage_map_delta_size = 0;
229        for (slot, storage_map_delta) in self.maps.iter() {
230            storage_map_delta_size += slot.get_size_hint() + storage_map_delta.get_size_hint();
233        }
234
235        u8_size * 3 +
237        self.cleared_slots().count() * u8_size +
239        self.updated_slots().count() * (u8_size + word_size) +
241        storage_map_delta_size
243    }
244}
245
246impl Deserializable for AccountStorageDelta {
247    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
248        let mut values = BTreeMap::new();
249
250        let num_cleared_items = source.read_u8()? as usize;
251        for _ in 0..num_cleared_items {
252            let cleared_slot = source.read_u8()?;
253            values.insert(cleared_slot, EMPTY_WORD);
254        }
255
256        let num_updated_items = source.read_u8()? as usize;
257        for _ in 0..num_updated_items {
258            let (updated_slot, updated_value) = source.read()?;
259            values.insert(updated_slot, updated_value);
260        }
261
262        let num_maps = source.read_u8()? as usize;
263        let maps = source.read_many::<(u8, StorageMapDelta)>(num_maps)?.into_iter().collect();
264
265        Self::from_parts(values, maps)
266            .map_err(|err| DeserializationError::InvalidValue(err.to_string()))
267    }
268}
269
270#[derive(Clone, Debug, Default, PartialEq, Eq)]
281pub struct StorageMapDelta(BTreeMap<LexicographicWord, Word>);
282
283impl StorageMapDelta {
284    pub fn new(map: BTreeMap<LexicographicWord, Word>) -> Self {
286        Self(map)
287    }
288
289    pub fn num_entries(&self) -> usize {
291        self.0.len()
292    }
293
294    pub fn entries(&self) -> &BTreeMap<LexicographicWord, Word> {
296        &self.0
297    }
298
299    pub fn insert(&mut self, key: Word, value: Word) {
301        self.0.insert(LexicographicWord::new(key), value);
302    }
303
304    pub fn is_empty(&self) -> bool {
306        self.0.is_empty()
307    }
308
309    pub fn merge(&mut self, other: Self) {
311        self.0.extend(other.0);
313    }
314
315    pub fn as_map_mut(&mut self) -> &mut BTreeMap<LexicographicWord, Word> {
317        &mut self.0
318    }
319
320    fn cleared_keys(&self) -> impl Iterator<Item = &Word> + '_ {
322        self.0.iter().filter(|&(_, value)| value.is_empty()).map(|(key, _)| key.inner())
323    }
324
325    fn updated_entries(&self) -> impl Iterator<Item = (&Word, &Word)> + '_ {
327        self.0.iter().filter_map(|(key, value)| {
328            if !value.is_empty() {
329                Some((key.inner(), value))
330            } else {
331                None
332            }
333        })
334    }
335}
336
337#[cfg(any(feature = "testing", test))]
338impl StorageMapDelta {
339    pub fn from_iters(
341        cleared_leaves: impl IntoIterator<Item = Word>,
342        updated_leaves: impl IntoIterator<Item = (Word, Word)>,
343    ) -> Self {
344        Self(BTreeMap::from_iter(
345            cleared_leaves
346                .into_iter()
347                .map(|key| (LexicographicWord::new(key), EMPTY_WORD))
348                .chain(
349                    updated_leaves
350                        .into_iter()
351                        .map(|(key, value)| (LexicographicWord::new(key), value)),
352                ),
353        ))
354    }
355
356    pub fn into_map(self) -> BTreeMap<LexicographicWord, Word> {
358        self.0
359    }
360}
361
362impl From<StorageMap> for StorageMapDelta {
364    fn from(map: StorageMap) -> Self {
365        StorageMapDelta::new(
366            map.into_entries()
367                .into_iter()
368                .map(|(key, value)| (LexicographicWord::new(key), value))
369                .collect(),
370        )
371    }
372}
373
374impl Serializable for StorageMapDelta {
375    fn write_into<W: ByteWriter>(&self, target: &mut W) {
376        let cleared: Vec<&Word> = self.cleared_keys().collect();
377        let updated: Vec<(&Word, &Word)> = self.updated_entries().collect();
378
379        target.write_usize(cleared.len());
380        target.write_many(cleared.iter());
381
382        target.write_usize(updated.len());
383        target.write_many(updated.iter());
384    }
385
386    fn get_size_hint(&self) -> usize {
387        let word_size = EMPTY_WORD.get_size_hint();
388
389        let cleared_keys_count = self.cleared_keys().count();
390        let updated_entries_count = self.updated_entries().count();
391
392        cleared_keys_count.get_size_hint() +
394        cleared_keys_count * Word::SERIALIZED_SIZE +
395
396        updated_entries_count.get_size_hint() +
398        updated_entries_count * (Word::SERIALIZED_SIZE + word_size)
399    }
400}
401
402impl Deserializable for StorageMapDelta {
403    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
404        let mut map = BTreeMap::new();
405
406        let cleared_count = source.read_usize()?;
407        for _ in 0..cleared_count {
408            let cleared_key = source.read()?;
409            map.insert(LexicographicWord::new(cleared_key), EMPTY_WORD);
410        }
411
412        let updated_count = source.read_usize()?;
413        for _ in 0..updated_count {
414            let (updated_key, updated_value) = source.read()?;
415            map.insert(LexicographicWord::new(updated_key), updated_value);
416        }
417
418        Ok(Self::new(map))
419    }
420}
421
422#[cfg(test)]
426mod tests {
427    use anyhow::Context;
428
429    use super::{AccountStorageDelta, Deserializable, Serializable};
430    use crate::account::StorageMapDelta;
431    use crate::testing::storage::AccountStorageDeltaBuilder;
432    use crate::{ONE, Word, ZERO};
433
434    #[test]
435    fn account_storage_delta_validation() {
436        let delta = AccountStorageDelta::from_iters(
437            [1, 2, 3],
438            [(4, Word::from([ONE, ONE, ONE, ONE])), (5, Word::from([ONE, ONE, ONE, ZERO]))],
439            [],
440        );
441        assert!(delta.validate().is_ok());
442
443        let bytes = delta.to_bytes();
444        assert_eq!(AccountStorageDelta::read_from_bytes(&bytes), Ok(delta));
445
446        let delta = AccountStorageDelta::from_iters(
448            [1, 2, 3],
449            [(2, Word::from([ONE, ONE, ONE, ONE])), (5, Word::from([ONE, ONE, ONE, ZERO]))],
450            [(1, StorageMapDelta::default())],
451        );
452        assert!(delta.validate().is_err());
453
454        let bytes = delta.to_bytes();
455        assert!(AccountStorageDelta::read_from_bytes(&bytes).is_err());
456
457        let delta = AccountStorageDelta::from_iters(
459            [1, 3],
460            [(2, Word::from([ONE, ONE, ONE, ONE])), (5, Word::from([ONE, ONE, ONE, ZERO]))],
461            [(2, StorageMapDelta::default())],
462        );
463        assert!(delta.validate().is_err());
464
465        let bytes = delta.to_bytes();
466        assert!(AccountStorageDelta::read_from_bytes(&bytes).is_err());
467    }
468
469    #[test]
470    fn test_is_empty() {
471        let storage_delta = AccountStorageDelta::new();
472        assert!(storage_delta.is_empty());
473
474        let storage_delta = AccountStorageDelta::from_iters([1], [], []);
475        assert!(!storage_delta.is_empty());
476
477        let storage_delta =
478            AccountStorageDelta::from_iters([], [(2, Word::from([ONE, ONE, ONE, ONE]))], []);
479        assert!(!storage_delta.is_empty());
480
481        let storage_delta =
482            AccountStorageDelta::from_iters([], [], [(3, StorageMapDelta::default())]);
483        assert!(!storage_delta.is_empty());
484    }
485
486    #[test]
487    fn test_serde_account_storage_delta() {
488        let storage_delta = AccountStorageDelta::new();
489        let serialized = storage_delta.to_bytes();
490        let deserialized = AccountStorageDelta::read_from_bytes(&serialized).unwrap();
491        assert_eq!(deserialized, storage_delta);
492
493        let storage_delta = AccountStorageDelta::from_iters([1], [], []);
494        let serialized = storage_delta.to_bytes();
495        let deserialized = AccountStorageDelta::read_from_bytes(&serialized).unwrap();
496        assert_eq!(deserialized, storage_delta);
497
498        let storage_delta =
499            AccountStorageDelta::from_iters([], [(2, Word::from([ONE, ONE, ONE, ONE]))], []);
500        let serialized = storage_delta.to_bytes();
501        let deserialized = AccountStorageDelta::read_from_bytes(&serialized).unwrap();
502        assert_eq!(deserialized, storage_delta);
503
504        let storage_delta =
505            AccountStorageDelta::from_iters([], [], [(3, StorageMapDelta::default())]);
506        let serialized = storage_delta.to_bytes();
507        let deserialized = AccountStorageDelta::read_from_bytes(&serialized).unwrap();
508        assert_eq!(deserialized, storage_delta);
509    }
510
511    #[test]
512    fn test_serde_storage_map_delta() {
513        let storage_map_delta = StorageMapDelta::default();
514        let serialized = storage_map_delta.to_bytes();
515        let deserialized = StorageMapDelta::read_from_bytes(&serialized).unwrap();
516        assert_eq!(deserialized, storage_map_delta);
517
518        let storage_map_delta = StorageMapDelta::from_iters([Word::from([ONE, ONE, ONE, ONE])], []);
519        let serialized = storage_map_delta.to_bytes();
520        let deserialized = StorageMapDelta::read_from_bytes(&serialized).unwrap();
521        assert_eq!(deserialized, storage_map_delta);
522
523        let storage_map_delta =
524            StorageMapDelta::from_iters([], [(Word::empty(), Word::from([ONE, ONE, ONE, ONE]))]);
525        let serialized = storage_map_delta.to_bytes();
526        let deserialized = StorageMapDelta::read_from_bytes(&serialized).unwrap();
527        assert_eq!(deserialized, storage_map_delta);
528    }
529
530    #[rstest::rstest]
531    #[case::some_some(Some(1), Some(2), Some(2))]
532    #[case::none_some(None, Some(2), Some(2))]
533    #[case::some_none(Some(1), None, None)]
534    #[test]
535    fn merge_items(
536        #[case] x: Option<u32>,
537        #[case] y: Option<u32>,
538        #[case] expected: Option<u32>,
539    ) -> anyhow::Result<()> {
540        fn create_delta(item: Option<u32>) -> anyhow::Result<AccountStorageDelta> {
542            const SLOT: u8 = 123;
543            let item = item.map(|x| (SLOT, Word::from([x, 0, 0, 0])));
544
545            AccountStorageDeltaBuilder::new()
546                .add_cleared_items(item.is_none().then_some(SLOT))
547                .add_updated_values(item)
548                .build()
549                .context("failed to build storage delta")
550        }
551
552        let mut delta_x = create_delta(x)?;
553        let delta_y = create_delta(y)?;
554        let expected = create_delta(expected)?;
555
556        delta_x.merge(delta_y).context("failed to merge deltas")?;
557
558        assert_eq!(delta_x, expected);
559
560        Ok(())
561    }
562
563    #[rstest::rstest]
564    #[case::some_some(Some(1), Some(2), Some(2))]
565    #[case::none_some(None, Some(2), Some(2))]
566    #[case::some_none(Some(1), None, None)]
567    #[test]
568    fn merge_maps(#[case] x: Option<u32>, #[case] y: Option<u32>, #[case] expected: Option<u32>) {
569        fn create_delta(value: Option<u32>) -> StorageMapDelta {
570            let key = Word::from([10u32, 0, 0, 0]);
571            match value {
572                Some(value) => {
573                    StorageMapDelta::from_iters([], [(key, Word::from([value, 0, 0, 0]))])
574                },
575                None => StorageMapDelta::from_iters([key], []),
576            }
577        }
578
579        let mut delta_x = create_delta(x);
580        let delta_y = create_delta(y);
581        let expected = create_delta(expected);
582
583        delta_x.merge(delta_y);
584
585        assert_eq!(delta_x, expected);
586    }
587}