fix: StorageMap bug for patch release (#745)

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## 0.3.1 (2024-06-12)
 * Replaced `cargo-make` with just `make` for running tasks (#696).
 * Made `DataStore` conditionally async using `winter-maybe-async` (#725)
+* Fixed `StorageMap`s implementation and included into apply_delta (#745)
 
 ## 0.3.0 (2024-05-14)
 

objects/src/accounts/mod.rs

@@ -302,30 +302,52 @@ pub mod testing {
         },
         code::testing::make_account_code,
         Account, AccountDelta, AccountId, AccountStorage, AccountStorageDelta, AccountVaultDelta,
-        Felt, SlotItem, StorageSlot, StorageSlotType, Word,
+        Felt, SlotItem, StorageMap, StorageSlot, StorageSlotType, Word,
     };
     use crate::assets::{Asset, AssetVault, FungibleAsset};
 
     // UTILITIES
     // --------------------------------------------------------------------------------------------
 
-    pub fn build_account(assets: Vec<Asset>, nonce: Felt, storage_items: Vec<Word>) -> Account {
+    pub fn build_account(
+        assets: Vec<Asset>,
+        nonce: Felt,
+        storage_items: Vec<Word>,
+        storage_map: Option<StorageMap>,
+    ) -> Account {
         let id = AccountId::try_from(ACCOUNT_ID_REGULAR_ACCOUNT_IMMUTABLE_CODE_ON_CHAIN).unwrap();
         let code = make_account_code();
 
         // build account data
         let vault = AssetVault::new(&assets).unwrap();
 
         let slot_type = StorageSlotType::Value { value_arity: 0 };
-        let slot_items: Vec<SlotItem> = storage_items
+        let mut slot_items: Vec<SlotItem> = storage_items
             .into_iter()
             .enumerate()
             .map(|(index, item)| SlotItem {
                 index: index as u8,
                 slot: StorageSlot { slot_type, value: item },
             })
             .collect();
-        let storage = AccountStorage::new(slot_items, vec![]).unwrap();
+
+        let mut maps = Vec::new();
+        let slot_type_map = StorageSlotType::Map { value_arity: 0 };
+        if let Some(storage_map) = storage_map {
+            let map_item = SlotItem {
+                // slot index * 10 to avoid conflicts with value slots
+                // this should only stay until we merge next with the proper solution
+                index: 100_u8,
+                slot: StorageSlot {
+                    slot_type: slot_type_map,
+                    value: storage_map.root().into(),
+                },
+            };
+            slot_items.push(map_item);
+            maps.push(storage_map);
+        }
+
+        let storage = AccountStorage::new(slot_items, maps).unwrap();
 
         Account::new(id, vault, storage, code, nonce)
     }
@@ -361,15 +383,18 @@ mod tests {
         Felt, Word,
     };
 
-    use super::{testing::*, AccountDelta, AccountStorageDelta, AccountVaultDelta};
-    use crate::accounts::{delta::AccountStorageDeltaBuilder, Account};
+    use super::{testing::*, AccountDelta, AccountStorageDelta, AccountVaultDelta, Digest};
+    use crate::accounts::{
+        delta::{AccountStorageDeltaBuilder, StorageMapDelta},
+        Account, StorageMap,
+    };
 
     #[test]
     fn test_serde_account() {
         let init_nonce = Felt::new(1);
         let (asset_0, _) = build_assets();
         let word = [Felt::new(1), Felt::new(2), Felt::new(3), Felt::new(4)];
-        let account = build_account(vec![asset_0], init_nonce, vec![word]);
+        let account = build_account(vec![asset_0], init_nonce, vec![word], None);
 
         let serialized = account.to_bytes();
         let deserialized = Account::read_from_bytes(&serialized).unwrap();
@@ -399,23 +424,53 @@ mod tests {
         let init_nonce = Felt::new(1);
         let (asset_0, asset_1) = build_assets();
         let word = [Felt::new(1), Felt::new(2), Felt::new(3), Felt::new(4)];
-        let mut account = build_account(vec![asset_0], init_nonce, vec![word]);
+        // StorageMap with values
+        let storage_map_leaves_2: [(Digest, Word); 2] = [
+            (
+                Digest::new([Felt::new(101), Felt::new(102), Felt::new(103), Felt::new(104)]),
+                [Felt::new(1_u64), Felt::new(2_u64), Felt::new(3_u64), Felt::new(4_u64)],
+            ),
+            (
+                Digest::new([Felt::new(105), Felt::new(106), Felt::new(107), Felt::new(108)]),
+                [Felt::new(5_u64), Felt::new(6_u64), Felt::new(7_u64), Felt::new(8_u64)],
+            ),
+        ];
+        let mut storage_map = StorageMap::with_entries(storage_map_leaves_2).unwrap();
+        let mut account =
+            build_account(vec![asset_0], init_nonce, vec![word], Some(storage_map.clone()));
+
+        let new_map_entry = (
+            Digest::new([Felt::new(101), Felt::new(102), Felt::new(103), Felt::new(104)]),
+            [Felt::new(9_u64), Felt::new(10_u64), Felt::new(11_u64), Felt::new(12_u64)],
+        );
+        let updated_map =
+            StorageMapDelta::from(vec![], vec![(new_map_entry.0.into(), new_map_entry.1)]);
+        storage_map.insert(new_map_entry.0, new_map_entry.1);
 
         // build account delta
         let final_nonce = Felt::new(2);
         let storage_delta = AccountStorageDeltaBuilder::new()
             .add_cleared_items([0])
-            .add_updated_items([(1_u8, [Felt::new(1), Felt::new(2), Felt::new(3), Felt::new(4)])])
+            .add_updated_items([
+                (1_u8, [Felt::new(1), Felt::new(2), Felt::new(3), Felt::new(4)]),
+                (100_u8, storage_map.root().into()),
+            ])
+            .add_updated_maps([(100_u8, updated_map)])
             .build()
             .unwrap();
         let account_delta =
             build_account_delta(vec![asset_1], vec![asset_0], final_nonce, storage_delta);
 
         // apply delta and create final_account
         account.apply_delta(&account_delta).unwrap();
-        let final_account = build_account(vec![asset_1], final_nonce, vec![Word::default(), word]);
 
-        // assert account is what it should be
+        let final_account = build_account(
+            vec![asset_1],
+            final_nonce,
+            vec![Word::default(), word],
+            Some(storage_map),
+        );
+
         assert_eq!(account, final_account);
     }
 
@@ -425,7 +480,7 @@ mod tests {
         // build account
         let init_nonce = Felt::new(1);
         let (asset, _) = build_assets();
-        let mut account = build_account(vec![asset], init_nonce, vec![Word::default()]);
+        let mut account = build_account(vec![asset], init_nonce, vec![Word::default()], None);
 
         // build account delta
         let storage_delta = AccountStorageDeltaBuilder::new()
@@ -445,7 +500,7 @@ mod tests {
         // build account
         let init_nonce = Felt::new(2);
         let (asset, _) = build_assets();
-        let mut account = build_account(vec![asset], init_nonce, vec![Word::default()]);
+        let mut account = build_account(vec![asset], init_nonce, vec![Word::default()], None);
 
         // build account delta
         let final_nonce = Felt::new(1);
@@ -466,7 +521,7 @@ mod tests {
         // build account
         let init_nonce = Felt::new(1);
         let word = [Felt::new(1), Felt::new(2), Felt::new(3), Felt::new(4)];
-        let mut account = build_account(vec![], init_nonce, vec![word]);
+        let mut account = build_account(vec![], init_nonce, vec![word], None);
 
         // build account delta
         let final_nonce = Felt::new(2);

objects/src/accounts/storage/mod.rs

@@ -2,9 +2,12 @@ use alloc::{collections::BTreeMap, string::ToString, vec::Vec};
 
 use super::{
     AccountError, AccountStorageDelta, ByteReader, ByteWriter, Deserializable,
-    DeserializationError, Digest, Felt, Hasher, Serializable, Word,
+    DeserializationError, Digest, Felt, Hasher, Serializable, Word, ZERO,
+};
+use crate::{
+    accounts::StorageMapDelta,
+    crypto::merkle::{LeafIndex, NodeIndex, SimpleSmt},
 };
-use crate::crypto::merkle::{LeafIndex, NodeIndex, SimpleSmt};
 
 mod slot;
 pub use slot::StorageSlotType;
@@ -206,6 +209,11 @@ impl AccountStorage {
         &self.maps
     }
 
+    // Returns the storage map with a given root.
+    pub fn find_storage_map_by_root(&mut self, target_root: Digest) -> Option<&mut StorageMap> {
+        self.maps.iter_mut().find(|map| map.root() == target_root)
+    }
+
     // DATA MUTATORS
     // --------------------------------------------------------------------------------------------
 
@@ -219,6 +227,12 @@ impl AccountStorage {
     /// - The delta implies an update to a reserved account slot.
     /// - The updates violate storage layout constraints.
     pub(super) fn apply_delta(&mut self, delta: &AccountStorageDelta) -> Result<(), AccountError> {
+        // Map updates are applied first as we need to find the storage map by its old root
+        // and every map updates always involves updating the root in the Storage slots as well.
+        for &(slot_idx, ref map_delta) in delta.updated_maps.iter() {
+            self.set_map_item(slot_idx, map_delta.clone())?;
+        }
+
         for &slot_idx in delta.cleared_items.iter() {
             self.set_item(slot_idx, Word::default())?;
         }
@@ -253,14 +267,61 @@ impl AccountStorage {
                     });
                 }
             },
-            slot_type => Err(AccountError::StorageSlotNotValueSlot(index, slot_type))?,
+            StorageSlotType::Map { value_arity } => {
+                if value_arity > 0 {
+                    return Err(AccountError::StorageSlotInvalidValueArity {
+                        slot: index,
+                        expected: 0,
+                        actual: value_arity,
+                    });
+                }
+            },
+            slot_type => Err(AccountError::StorageSlotArrayNotSupportedYet(index, slot_type))?,
         }
 
         // update the slot and return
         let index = LeafIndex::new(index as u64).expect("index is u8 - index within range");
         let slot_value = self.slots.insert(index, value);
         Ok(slot_value)
     }
+
+    /// Updates a storage map at the specified index.
+    ///
+    /// # Errors
+    /// Returns an error if:
+    /// - The index specifies a reserved storage slot.
+    /// - The index is not u8.
+    /// - The map does not exist at the specified index.
+    pub fn set_map_item(
+        &mut self,
+        index: u8,
+        map_delta: StorageMapDelta,
+    ) -> Result<(), AccountError> {
+        // layout commitment slot cannot be updated
+        if index == Self::SLOT_LAYOUT_COMMITMENT_INDEX {
+            return Err(AccountError::StorageSlotIsReserved(index));
+        }
+
+        // load the storage map
+        let index = LeafIndex::new(index as u64).expect("index is u8 - index within range");
+        let old_map_root: Digest = self.slots.get_leaf(&index).into();
+
+        let storage_map = self
+            .find_storage_map_by_root(old_map_root)
+            .ok_or(AccountError::StorageMapNotFound { index: index.value() })?;
+
+        // apply the updated leaves to the storage map
+        for (key, value) in map_delta.updated_leaves.iter() {
+            storage_map.insert(key.into(), *value);
+        }
+
+        // apply the cleared leaves to the storage map
+        for key in map_delta.cleared_leaves.iter() {
+            storage_map.insert(key.into(), [ZERO; 4]);
+        }
+
+        Ok(())
+    }
 }
 
 // SERIALIZATION

objects/src/errors.rs

@@ -33,8 +33,9 @@ pub enum AccountError {
     SeedDigestTooFewTrailingZeros { expected: u32, actual: u32 },
     StorageSlotInvalidValueArity { slot: u8, expected: u8, actual: u8 },
     StorageSlotIsReserved(u8),
-    StorageSlotNotValueSlot(u8, StorageSlotType),
+    StorageSlotArrayNotSupportedYet(u8, StorageSlotType),
     StorageMapToManyMaps { expected: usize, actual: usize },
+    StorageMapNotFound { index: u64 },
     StubDataIncorrectLength(usize, usize),
 }