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}