1use alloc::string::ToString;
2use alloc::vec::Vec;
3
4use super::{
5 AccountError,
6 AccountStoragePatch,
7 ByteReader,
8 ByteWriter,
9 Deserializable,
10 DeserializationError,
11 Felt,
12 Serializable,
13 Word,
14};
15use crate::account::{
16 AccountComponent,
17 StorageMapPatch,
18 StorageMapPatchEntries,
19 StorageSlotPatch,
20 StorageValuePatch,
21};
22use crate::asset::AssetCallbacks;
23use crate::crypto::SequentialCommit;
24
25pub(crate) mod slot;
26pub use slot::{StorageSlot, StorageSlotContent, StorageSlotId, StorageSlotName, StorageSlotType};
27
28mod map;
29pub use map::{PartialStorageMap, StorageMap, StorageMapKey, StorageMapKeyHash, StorageMapWitness};
30
31mod header;
32pub use header::{AccountStorageHeader, StorageSlotHeader};
33
34mod partial;
35pub use partial::PartialStorage;
36
37#[derive(Debug, Clone, Default, PartialEq, Eq)]
58pub struct AccountStorage {
59 slots: Vec<StorageSlot>,
60}
61
62impl AccountStorage {
63 pub const MAX_NUM_STORAGE_SLOTS: usize = 255;
65
66 pub fn new(mut slots: Vec<StorageSlot>) -> Result<AccountStorage, AccountError> {
79 let num_slots = slots.len();
80
81 if num_slots > Self::MAX_NUM_STORAGE_SLOTS {
82 return Err(AccountError::StorageTooManySlots(num_slots as u64));
83 }
84
85 slots.sort_unstable_by(|a, b| a.name().cmp(b.name()));
87
88 for slots in slots.windows(2) {
91 if slots[0].id() == slots[1].id() {
92 return Err(AccountError::DuplicateStorageSlotName(slots[0].name().clone()));
93 }
94 }
95
96 Ok(Self { slots })
97 }
98
99 pub(super) fn from_components(
107 components: Vec<AccountComponent>,
108 ) -> Result<AccountStorage, AccountError> {
109 let storage_slots = components
110 .into_iter()
111 .flat_map(|component| {
112 let AccountComponent { storage_slots, .. } = component;
113 storage_slots.into_iter()
114 })
115 .collect();
116
117 Self::new(storage_slots)
118 }
119
120 pub fn to_elements(&self) -> Vec<Felt> {
131 <Self as SequentialCommit>::to_elements(self)
132 }
133
134 pub fn to_commitment(&self) -> Word {
136 <Self as SequentialCommit>::to_commitment(self)
137 }
138
139 pub fn num_slots(&self) -> u8 {
141 self.slots.len() as u8
144 }
145
146 pub fn slots(&self) -> &[StorageSlot] {
148 &self.slots
149 }
150
151 pub fn into_slots(self) -> Vec<StorageSlot> {
153 self.slots
154 }
155
156 pub fn to_header(&self) -> AccountStorageHeader {
158 AccountStorageHeader::new(self.slots.iter().map(StorageSlotHeader::from).collect())
159 .expect("slots should be valid as ensured by AccountStorage")
160 }
161
162 pub fn get(&self, slot_name: &StorageSlotName) -> Option<&StorageSlot> {
165 self.slots.iter().find(|slot| slot.name().id() == slot_name.id())
166 }
167
168 pub fn has_callback_slots(&self) -> bool {
177 AssetCallbacks::slot_names()
178 .iter()
179 .any(|slot_name| self.get(slot_name).is_some())
180 }
181
182 fn get_mut(&mut self, slot_name: &StorageSlotName) -> Option<&mut StorageSlot> {
185 self.slots.iter_mut().find(|slot| slot.name().id() == slot_name.id())
186 }
187
188 pub fn get_item(&self, slot_name: &StorageSlotName) -> Result<Word, AccountError> {
195 self.get(slot_name)
196 .map(|slot| slot.content().value())
197 .ok_or_else(|| AccountError::StorageSlotNameNotFound { slot_name: slot_name.clone() })
198 }
199
200 pub fn get_map_item(
208 &self,
209 slot_name: &StorageSlotName,
210 key: StorageMapKey,
211 ) -> Result<Word, AccountError> {
212 self.get(slot_name)
213 .ok_or_else(|| AccountError::StorageSlotNameNotFound { slot_name: slot_name.clone() })
214 .and_then(|slot| match slot.content() {
215 StorageSlotContent::Map(map) => Ok(map.get(&key)),
216 _ => Err(AccountError::StorageSlotNotMap(slot_name.clone())),
217 })
218 }
219
220 pub(super) fn apply_patch(&mut self, patch: &AccountStoragePatch) -> Result<(), AccountError> {
230 for (slot_name, slot_patch) in patch.slots() {
231 match slot_patch {
232 StorageSlotPatch::Value(value_patch) => {
233 self.apply_value_patch(slot_name, value_patch)?
234 },
235 StorageSlotPatch::Map(map_patch) => self.apply_map_patch(slot_name, map_patch)?,
236 }
237 }
238
239 Ok(())
240 }
241
242 fn apply_value_patch(
244 &mut self,
245 slot_name: &StorageSlotName,
246 value_patch: &StorageValuePatch,
247 ) -> Result<(), AccountError> {
248 match value_patch {
249 StorageValuePatch::Create { value } => {
250 self.create_value_slot(slot_name.clone(), *value)?;
251 },
252 StorageValuePatch::Update { value } => {
253 self.set_item(slot_name, *value)?;
254 },
255 StorageValuePatch::Remove => {
256 self.remove_slot(slot_name)?;
257 },
258 }
259
260 Ok(())
261 }
262
263 fn apply_map_patch(
265 &mut self,
266 slot_name: &StorageSlotName,
267 map_patch: &StorageMapPatch,
268 ) -> Result<(), AccountError> {
269 match map_patch {
270 StorageMapPatch::Create { entries } => {
271 self.create_map_slot(slot_name.clone(), entries)?;
272 },
273 StorageMapPatch::Update { entries } => {
274 let slot = self.get_mut(slot_name).ok_or_else(|| {
275 AccountError::StorageSlotNameNotFound { slot_name: slot_name.clone() }
276 })?;
277
278 let storage_map = match slot.content_mut() {
279 StorageSlotContent::Map(map) => map,
280 _ => return Err(AccountError::StorageSlotNotMap(slot_name.clone())),
281 };
282
283 storage_map.apply_patch(entries)?;
284 },
285 StorageMapPatch::Remove => {
286 self.remove_slot(slot_name)?;
287 },
288 }
289
290 Ok(())
291 }
292
293 pub fn set_item(
304 &mut self,
305 slot_name: &StorageSlotName,
306 value: Word,
307 ) -> Result<Word, AccountError> {
308 let slot = self.get_mut(slot_name).ok_or_else(|| {
309 AccountError::StorageSlotNameNotFound { slot_name: slot_name.clone() }
310 })?;
311
312 let StorageSlotContent::Value(old_value) = slot.content() else {
313 return Err(AccountError::StorageSlotNotValue(slot_name.clone()));
314 };
315 let old_value = *old_value;
316
317 let mut new_slot = StorageSlotContent::Value(value);
318 core::mem::swap(slot.content_mut(), &mut new_slot);
319
320 Ok(old_value)
321 }
322
323 pub fn set_map_item(
334 &mut self,
335 slot_name: &StorageSlotName,
336 key: StorageMapKey,
337 value: Word,
338 ) -> Result<(Word, Word), AccountError> {
339 let slot = self.get_mut(slot_name).ok_or_else(|| {
340 AccountError::StorageSlotNameNotFound { slot_name: slot_name.clone() }
341 })?;
342
343 let StorageSlotContent::Map(storage_map) = slot.content_mut() else {
344 return Err(AccountError::StorageSlotNotMap(slot_name.clone()));
345 };
346
347 let old_root = storage_map.root();
348
349 let old_value = storage_map.insert(key, value)?;
350
351 Ok((old_root, old_value))
352 }
353
354 fn create_value_slot(
361 &mut self,
362 slot_name: StorageSlotName,
363 value: Word,
364 ) -> Result<(), AccountError> {
365 self.create_slot(StorageSlot::with_value(slot_name, value))
366 }
367
368 fn create_map_slot(
375 &mut self,
376 slot_name: StorageSlotName,
377 entries: &StorageMapPatchEntries,
378 ) -> Result<(), AccountError> {
379 let storage_map =
380 StorageMap::with_entries(entries.as_map().iter().map(|(key, value)| (*key, *value)))
381 .expect("map should contain only unique entries");
382
383 self.create_slot(StorageSlot::with_map(slot_name, storage_map))
384 }
385
386 fn remove_slot(&mut self, slot_name: &StorageSlotName) -> Result<(), AccountError> {
392 match self.slots.iter().position(|slot| slot.name().id() == slot_name.id()) {
393 Some(index) => {
394 self.slots.remove(index);
395 Ok(())
396 },
397 None => Err(AccountError::StorageSlotNameNotFound { slot_name: slot_name.clone() }),
398 }
399 }
400
401 fn create_slot(&mut self, slot: StorageSlot) -> Result<(), AccountError> {
412 match self.slots.binary_search_by(|existing| existing.name().cmp(slot.name())) {
413 Ok(index) => {
414 self.slots[index] = slot;
415 Ok(())
416 },
417 Err(index) => {
418 if self.slots.len() >= Self::MAX_NUM_STORAGE_SLOTS {
419 return Err(AccountError::StorageTooManySlots(self.slots.len() as u64 + 1));
420 }
421
422 self.slots.insert(index, slot);
423 Ok(())
424 },
425 }
426 }
427}
428
429impl IntoIterator for AccountStorage {
433 type Item = StorageSlot;
434 type IntoIter = alloc::vec::IntoIter<StorageSlot>;
435
436 fn into_iter(self) -> Self::IntoIter {
437 self.slots.into_iter()
438 }
439}
440
441impl SequentialCommit for AccountStorage {
445 type Commitment = Word;
446
447 fn to_elements(&self) -> Vec<Felt> {
448 self.slots()
449 .iter()
450 .flat_map(|slot| {
451 StorageSlotHeader::new(
452 slot.name().clone(),
453 slot.content().slot_type(),
454 slot.content().value(),
455 )
456 .to_elements()
457 })
458 .collect()
459 }
460}
461
462impl Serializable for AccountStorage {
466 fn write_into<W: ByteWriter>(&self, target: &mut W) {
467 target.write_u8(self.slots().len() as u8);
468 target.write_many(self.slots());
469 }
470
471 fn get_size_hint(&self) -> usize {
472 let u8_size = 0u8.get_size_hint();
474 let mut size = u8_size;
475
476 for slot in self.slots() {
477 size += slot.get_size_hint();
478 }
479
480 size
481 }
482}
483
484impl Deserializable for AccountStorage {
485 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
486 let num_slots = source.read_u8()? as usize;
487 let slots = source.read_many_iter::<StorageSlot>(num_slots)?.collect::<Result<_, _>>()?;
488
489 Self::new(slots).map_err(|err| DeserializationError::InvalidValue(err.to_string()))
490 }
491}
492
493#[cfg(test)]
497mod tests {
498 use std::collections::BTreeMap;
499
500 use assert_matches::assert_matches;
501
502 use super::{AccountStorage, Deserializable, Serializable};
503 use crate::Word;
504 use crate::account::{
505 AccountStorageHeader,
506 AccountStoragePatch,
507 StorageSlot,
508 StorageSlotHeader,
509 StorageSlotName,
510 StorageSlotPatch,
511 StorageValuePatch,
512 };
513 use crate::errors::AccountError;
514
515 #[test]
516 fn test_serde_account_storage() -> anyhow::Result<()> {
517 let storage = AccountStorage::new(vec![]).unwrap();
519 let bytes = storage.to_bytes();
520 assert_eq!(storage, AccountStorage::read_from_bytes(&bytes).unwrap());
521
522 let storage = AccountStorage::new(vec![
524 StorageSlot::with_empty_value(StorageSlotName::new("miden::test::value")?),
525 StorageSlot::with_empty_map(StorageSlotName::new("miden::test::map")?),
526 ])
527 .unwrap();
528 let bytes = storage.to_bytes();
529 assert_eq!(storage, AccountStorage::read_from_bytes(&bytes).unwrap());
530
531 Ok(())
532 }
533
534 #[test]
535 fn test_get_slot_by_name() -> anyhow::Result<()> {
536 let counter_slot = StorageSlotName::new("miden::test::counter")?;
537 let map_slot = StorageSlotName::new("miden::test::map")?;
538
539 let slots = vec![
540 StorageSlot::with_empty_value(counter_slot.clone()),
541 StorageSlot::with_empty_map(map_slot.clone()),
542 ];
543 let storage = AccountStorage::new(slots.clone())?;
544
545 assert_eq!(storage.get(&counter_slot).unwrap(), &slots[0]);
546 assert_eq!(storage.get(&map_slot).unwrap(), &slots[1]);
547
548 Ok(())
549 }
550
551 #[test]
552 fn test_account_storage_and_header_fail_on_duplicate_slot_name() -> anyhow::Result<()> {
553 let slot_name0 = StorageSlotName::mock(0);
554 let slot_name1 = StorageSlotName::mock(1);
555 let slot_name2 = StorageSlotName::mock(2);
556
557 let mut slots = vec![
558 StorageSlot::with_empty_value(slot_name0.clone()),
559 StorageSlot::with_empty_value(slot_name1.clone()),
560 StorageSlot::with_empty_map(slot_name0.clone()),
561 StorageSlot::with_empty_value(slot_name2.clone()),
562 ];
563
564 let err = AccountStorage::new(slots.clone()).unwrap_err();
567
568 assert_matches!(err, AccountError::DuplicateStorageSlotName(name) => {
569 assert_eq!(name, slot_name0);
570 });
571
572 slots.sort_unstable_by(|a, b| a.name().cmp(b.name()));
573 let err = AccountStorageHeader::new(slots.iter().map(StorageSlotHeader::from).collect())
574 .unwrap_err();
575
576 assert_matches!(err, AccountError::DuplicateStorageSlotName(name) => {
577 assert_eq!(name, slot_name0);
578 });
579
580 Ok(())
581 }
582
583 #[test]
584 fn create_value_slot_recreates_existing() -> anyhow::Result<()> {
585 let slot_name = StorageSlotName::mock(4);
586 let mut storage = AccountStorage::new(vec![StorageSlot::with_value(
587 slot_name.clone(),
588 Word::from([1u32, 2, 3, 4]),
589 )])?;
590
591 let new_value = Word::from([5u32, 6, 7, 8]);
593 storage.create_value_slot(slot_name.clone(), new_value)?;
594
595 assert_eq!(storage.num_slots(), 1);
596 assert_eq!(storage.get_item(&slot_name)?, new_value);
597
598 Ok(())
599 }
600
601 #[test]
602 fn remove_slot_rejects_absent() -> anyhow::Result<()> {
603 let absent = StorageSlotName::new("miden::test::absent")?;
604 let mut storage = AccountStorage::default();
605
606 let err = storage.remove_slot(&absent).unwrap_err();
607 assert_matches!(err, AccountError::StorageSlotNameNotFound { slot_name } => {
608 assert_eq!(slot_name, absent);
609 });
610
611 Ok(())
612 }
613
614 #[test]
615 fn create_and_remove_value_slot_roundtrip() -> anyhow::Result<()> {
616 let existing0 = StorageSlotName::mock(1);
618 let existing1 = StorageSlotName::mock(7);
619 let created = StorageSlotName::mock(20);
620 assert!(existing0 < created);
621 assert!(created < existing1);
622
623 let value = Word::from([9u32, 8, 7, 6]);
624
625 let mut storage = AccountStorage::new(vec![
626 StorageSlot::with_value(existing0.clone(), value),
627 StorageSlot::with_value(existing1.clone(), value),
628 ])?;
629
630 storage.create_value_slot(created.clone(), value)?;
631 assert_eq!(storage.num_slots(), 3);
632 assert_eq!(storage.get_item(&created)?, value);
633 assert!(
634 storage.slots().is_sorted_by_key(|slot| slot.name()),
635 "slots should remain sorted after insertion"
636 );
637
638 assert_eq!(storage.get_item(&existing0)?, value, "existing slot should remain accessible");
639 assert_eq!(storage.get_item(&existing1)?, value, "existing slot should remain accessible");
640
641 storage.remove_slot(&created)?;
642 assert_eq!(storage.num_slots(), 2);
643 assert_matches!(
644 storage.get_item(&created).unwrap_err(),
645 AccountError::StorageSlotNameNotFound { .. }
646 );
647
648 Ok(())
649 }
650
651 #[test]
652 fn apply_storage_patch() -> anyhow::Result<()> {
653 let updated = StorageSlotName::mock(1);
654 let created = StorageSlotName::mock(2);
655 let removed = StorageSlotName::mock(3);
656
657 let init_value = Word::from([1u32, 2, 3, 4]);
658 let final_value = Word::from([6u32, 7, 8, 9]);
659
660 let mut storage = AccountStorage::new(vec![
661 StorageSlot::with_value(updated.clone(), init_value),
662 StorageSlot::with_value(removed.clone(), init_value),
663 ])?;
664
665 let patches = BTreeMap::from_iter([
666 (
667 created.clone(),
668 StorageSlotPatch::Value(StorageValuePatch::Create { value: final_value }),
669 ),
670 (
671 updated.clone(),
672 StorageSlotPatch::Value(StorageValuePatch::Update { value: final_value }),
673 ),
674 (removed.clone(), StorageSlotPatch::Value(StorageValuePatch::Remove)),
675 ]);
676 let patch = AccountStoragePatch::from_raw(patches)?;
677
678 storage.apply_patch(&patch)?;
679
680 assert_eq!(storage.num_slots(), 2);
681 assert_eq!(storage.get_item(&created)?, final_value);
682 assert_eq!(storage.get_item(&updated)?, final_value);
683 assert_eq!(storage.get(&removed), None);
684
685 Ok(())
686 }
687}