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#[derive(Clone, Debug, Default, PartialEq, Eq)]
28pub struct AccountStorageDelta {
29 values: BTreeMap<u8, Word>,
30 maps: BTreeMap<u8, StorageMapDelta>,
31}
32
33impl AccountStorageDelta {
34 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 pub fn values(&self) -> &BTreeMap<u8, Word> {
47 &self.values
48 }
49
50 pub fn maps(&self) -> &BTreeMap<u8, StorageMapDelta> {
52 &self.maps
53 }
54
55 pub fn is_empty(&self) -> bool {
57 self.values.is_empty() && self.maps.is_empty()
58 }
59
60 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 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 pub fn merge(&mut self, other: Self) -> Result<(), AccountDeltaError> {
72 self.values.extend(other.values);
73
74 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 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 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 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 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
132impl 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 storage_map_delta_size += slot.get_size_hint() + storage_map_delta.get_size_hint();
177 }
178
179 u8_size * 3 +
181 self.cleared_slots().count() * u8_size +
183 self.updated_slots().count() * (u8_size + word_size) +
185 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#[derive(Clone, Debug, Default, PartialEq, Eq)]
221pub struct StorageMapDelta(BTreeMap<Digest, Word>);
222
223impl StorageMapDelta {
224 pub fn new(map: BTreeMap<Digest, Word>) -> Self {
226 Self(map)
227 }
228
229 pub fn leaves(&self) -> &BTreeMap<Digest, Word> {
231 &self.0
232 }
233
234 pub fn insert(&mut self, key: Digest, value: Word) {
236 self.0.insert(key, value);
237 }
238
239 pub fn is_empty(&self) -> bool {
241 self.0.is_empty()
242 }
243
244 pub fn merge(&mut self, other: Self) {
246 self.0.extend(other.0);
248 }
249
250 fn cleared_keys(&self) -> impl Iterator<Item = &Digest> + '_ {
252 self.0.iter().filter(|&(_, value)| value == &EMPTY_WORD).map(|(key, _)| key)
253 }
254
255 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 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
277impl From<StorageMap> for StorageMapDelta {
279 fn from(map: StorageMap) -> Self {
280 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_count.get_size_hint() +
318 cleared_keys_count * Digest::SERIALIZED_SIZE +
319
320 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#[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 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 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 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}