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#[derive(Clone, Debug, Default, PartialEq, Eq)]
27pub struct AccountStorageDelta {
28 values: BTreeMap<u8, Word>,
29 maps: BTreeMap<u8, StorageMapDelta>,
30}
31
32impl AccountStorageDelta {
33 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 pub fn values(&self) -> &BTreeMap<u8, Word> {
46 &self.values
47 }
48
49 pub fn maps(&self) -> &BTreeMap<u8, StorageMapDelta> {
51 &self.maps
52 }
53
54 pub fn is_empty(&self) -> bool {
56 self.values.is_empty() && self.maps.is_empty()
57 }
58
59 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 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 pub fn merge(&mut self, other: Self) -> Result<(), AccountDeltaError> {
71 self.values.extend(other.values);
72
73 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 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 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 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 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
131impl 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 storage_map_delta_size += slot.get_size_hint() + storage_map_delta.get_size_hint();
176 }
177
178 u8_size * 3 +
180 self.cleared_slots().count() * u8_size +
182 self.updated_slots().count() * (u8_size + word_size) +
184 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#[derive(Clone, Debug, Default, PartialEq, Eq)]
220pub struct StorageMapDelta(BTreeMap<Digest, Word>);
221
222impl StorageMapDelta {
223 pub fn new(map: BTreeMap<Digest, Word>) -> Self {
225 Self(map)
226 }
227
228 pub fn leaves(&self) -> &BTreeMap<Digest, Word> {
230 &self.0
231 }
232
233 pub fn insert(&mut self, key: Digest, value: Word) {
235 self.0.insert(key, value);
236 }
237
238 pub fn is_empty(&self) -> bool {
240 self.0.is_empty()
241 }
242
243 pub fn merge(&mut self, other: Self) {
245 self.0.extend(other.0);
247 }
248
249 fn cleared_keys(&self) -> impl Iterator<Item = &Digest> + '_ {
251 self.0.iter().filter(|&(_, value)| value == &EMPTY_WORD).map(|(key, _)| key)
252 }
253
254 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 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
276impl From<StorageMap> for StorageMapDelta {
278 fn from(map: StorageMap) -> Self {
279 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_count.get_size_hint() +
317 cleared_keys_count * Digest::SERIALIZED_SIZE +
318
319 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#[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 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 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 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}