Refactor apply_delta for Storage Maps (#758)

* refactor: decouple AccountStorageDelta.updated_items and AccountStorageDelta.updated_maps

* refactor set_map_item

* after review

* refactoring masm code

* after last feedback

miden-lib/asm/kernels/transaction/api.masm

@@ -2,6 +2,7 @@ use.std::collections::smt
 
 use.miden::kernels::tx::account
 use.miden::kernels::tx::asset_vault
+use.miden::kernels::tx::constants
 use.miden::kernels::tx::faucet
 use.miden::kernels::tx::memory
 use.miden::kernels::tx::note
@@ -16,6 +17,9 @@ const.ERR_FAUCET_RESERVED_DATA_SLOT=0x00020000
 # Procedure can only be called from faucet accounts
 const.ERR_ACCT_MUST_BE_A_FAUCET=0x00020001
 
+# Getting a map item on a non-map slot
+const.ERR_READING_MAP_VALUE_FROM_NON_MAP_SLOT=0x00020049
+
 # EVENTS
 # =================================================================================================
 
@@ -29,14 +33,6 @@ const.ACCOUNT_VAULT_BEFORE_REMOVE_ASSET_EVENT=131074
 # Event emitted after an asset is removed from the account vault.
 const.ACCOUNT_VAULT_AFTER_REMOVE_ASSET_EVENT=131075
 
-# TYPES
-# =================================================================================================
-
-# Type of storage slot item in the account storage
-const.STORAGE_VALUE=0
-const.STORAGE_MAP=1
-const.STORAGE_ARRAY=2
-
 # AUTHENTICATION
 # =================================================================================================
 
@@ -208,7 +204,7 @@ export.get_account_map_item
     # => [slot_type, index, KEY, ...]
 
     # fails if slot_type is not 1 = map
-    push.STORAGE_MAP eq assert
+    exec.constants::get_storage_slot_type_map eq assert.err=ERR_READING_MAP_VALUE_FROM_NON_MAP_SLOT
     # => [index, KEY, ...]
 
     # fetch the account storage item, which is ROOT of the map
@@ -244,15 +240,7 @@ export.set_account_map_item.1
     exec.authenticate_account_origin
     # => [index, KEY, NEW_VALUE, ...]
 
-    # check if storage type is map
-    dup exec.account::get_storage_slot_type_info drop
-    # => [slot_type, index, KEY, NEW_VALUE, ...]
-
-    # fails if slot_type is not 1 = map
-    push.STORAGE_MAP eq assert
-    # => [index, KEY, NEW_VALUE, ...]
-
-    # stores index for later
+    # store index for later
     dup loc_store.0
     # => [index, KEY, NEW_VALUE, ...]
 
@@ -261,12 +249,8 @@ export.set_account_map_item.1
     # => [KEY, NEW_VALUE, OLD_MAP_ROOT, ...]
 
     # set the new map item
-    loc_load.0 exec.account::set_map_item swapw
-    # => [NEW_MAP_ROOT, OLD_MAP_VALUE, ...]
-
-    # set the root of the map in the respective account storage slot
-    loc_load.0 exec.account::set_item
-    # => [OLD_MAP_ROOT, OLD_MAP_VALUE, ...]
+    loc_load.0 exec.account::set_map_item
+    # => [OLD_MAP_ROOT, OLD_VALUE, ...]
 
     # organize the stack for return (16 elements)
     movupw.2 dropw

miden-lib/asm/miden/kernels/tx/account.masm

@@ -22,6 +22,12 @@ const.ERR_ACCOUNT_SEED_DIGEST_MISMATCH=0x0002003E
 # Account pow is insufficient
 const.ERR_ACCOUNT_INVALID_POW=0x0002003F
 
+# Setting non-value item on a value slot
+const.ERR_SETTING_NON_VALUE_ITEM_ON_VALUE_SLOT=0x00020047
+
+# Setting a map item on a non-map slot
+const.ERR_SETTING_MAP_ITEM_ON_NON_MAP_SLOT=0x00020048
+
 # CONSTANTS
 # =================================================================================================
 
@@ -344,7 +350,33 @@ export.get_item
     # => [VALUE]
 end
 
-#! Sets an item in the account storage. Panics if the index is out of bounds.
+#! Sets an item in the account storage. Doesn't emit any events.
+#!
+#! Stack: [index, NEW_VALUE]
+#! Output: [OLD_VALUE]
+#!
+#! - index is the index of the item to set.
+#! - NEW_VALUE is the value to set.
+#! - OLD_VALUE is the previous value of the item.
+#! - OLD_ROOT is the old storage root.
+#! - NEW_ROOT is the new storage root.
+proc.set_item_raw
+    # get the storage root
+    exec.memory::get_acct_storage_root
+    # => [OLD_ROOT, index, NEW_VALUE]
+
+    # set the item in storage
+    movup.4 push.STORAGE_TREE_DEPTH mtree_set
+    # => [OLD_VALUE, NEW_ROOT]
+
+    # set the new storage root
+    swapw exec.memory::set_acct_storage_root
+    # => [OLD_VALUE]
+end
+
+#! Sets an item in the account storage. Panics if 
+#! - the index is out of bounds
+#! - the slot type is not value
 #!
 #! Stack: [index, V']
 #! Output: [V]
@@ -357,29 +389,29 @@ export.set_item
     emit.ACCOUNT_STORAGE_BEFORE_SET_ITEM_EVENT
     # => [index, V']
 
+    # check if storage type is slot
+    dup exec.memory::get_acct_storage_slot_type_data u32split drop
+    # => [slot_type, index, V', 0, 0, 0]
+
+    # fails if slot_type is not 0 = value
+    exec.constants::get_storage_slot_type_value eq assert.err=ERR_SETTING_NON_VALUE_ITEM_ON_VALUE_SLOT
+    # => [index, V', 0, 0, 0]
+
     # duplicate the index and the V' to be able to emit an event after an account storage item is
     # being updated
     movdn.4 dupw dup.8
     # => [index, V', V', index]
 
-    # get the storage root
-    exec.memory::get_acct_storage_root
-    # => [R, index, V', V', index]
-
-    # set the item in storage
-    movup.4 push.STORAGE_TREE_DEPTH mtree_set
-    # => [V, R', V', index]
-
-    # set the new storage root
-    swapw exec.memory::set_acct_storage_root
+    # set V' in the storage slot
+    exec.set_item_raw
     # => [V, V', index]
 
     # emit event to signal that an account storage item is being updated
     swapw movup.8 emit.ACCOUNT_STORAGE_AFTER_SET_ITEM_EVENT drop dropw
     # => [V]
 end
 
-#! Sets an item in the specified account storage map. Panics if no map is found for ROOT.
+#! Sets an item in the specified account storage map.
 #!
 #! Stack: [index, KEY, NEW_VALUE, OLD_ROOT, ...]
 #! Output: [OLD_VALUE, NEW_ROOT, ...]
@@ -389,7 +421,23 @@ end
 #! - KEY is the key to set.
 #! - OLD_VALUE is the previous value of the item.
 #! - NEW_ROOT is the new root of the map.
-export.set_map_item.2
+#!
+#! Panics if
+#! - the slot type is not map
+#! - no map is found for ROOT
+export.set_map_item.3
+    # store index for later
+    dup loc_store.0
+    # => [index, KEY, NEW_VALUE, ...]
+
+    # check if storage type is map
+    dup exec.memory::get_acct_storage_slot_type_data u32split drop
+    # => [slot_type, index, KEY, NEW_VALUE, OLD_ROOT]
+
+    # fails if slot_type is not 1 = map
+    exec.constants::get_storage_slot_type_map eq assert.err=ERR_SETTING_MAP_ITEM_ON_NON_MAP_SLOT
+    # => [index, KEY, NEW_VALUE, OLD_ROOT]
+
     push.0 drop                                             # TODO: remove line, see miden-vm/#1122
     emit.ACCOUNT_STORAGE_BEFORE_SET_MAP_ITEM_EVENT
     # => [index, KEY, NEW_VALUE, OLD_ROOT]
@@ -403,17 +451,22 @@ export.set_map_item.2
     # note smt::set expects the stack to be [NEW_VALUE, KEY, OLD_ROOT, ...]
     swapw exec.smt::set
     # => [OLD_VALUE, NEW_ROOT, KEY, NEW_VALUE, index, ...]
+
     # store OLD_VALUE and NEW_ROOT until the end of the procedure
-    loc_storew.0 dropw loc_storew.1 dropw
+    loc_storew.1 dropw loc_storew.2 dropw
     # => [KEY, NEW_VALUE, index, ...]
 
     # emit event to signal that an account storage item is being updated
     movup.8 emit.ACCOUNT_STORAGE_AFTER_SET_MAP_ITEM_EVENT drop
     # => [KEY, NEW_VALUE, ...]
 
     # load OLD_VALUE and NEW_ROOT on the top of the stack
-    loc_loadw.1 swapw loc_loadw.0
-    # => [OLD_VALUE, NEW_ROOT]
+    loc_loadw.2 swapw loc_loadw.1 swapw
+    # => [NEW_ROOT, OLD_VALUE, ...]
+
+    # set the root of the map in the respective account storage slot
+    loc_load.0 exec.set_item_raw
+    # => [OLD_MAP_ROOT, OLD_VALUE, ...]
 end
 
 #! Verifies that the procedure root is part of the account code Merkle tree. Panics if the

miden-lib/asm/miden/kernels/tx/constants.masm

@@ -30,6 +30,14 @@ const.REGULAR_ACCOUNT_SEED_DIGEST_MODULUS=8388608
 # zeros for a faucet account (2^31).
 const.FAUCET_ACCOUNT_SEED_DIGEST_MODULUS=2147483648
 
+# TYPES
+# =================================================================================================
+
+# Type of storage slot item in the account storage
+const.STORAGE_SLOT_TYPE_VALUE=0
+const.STORAGE_SLOT_TYPE_MAP=1
+const.STORAGE_SLOT_TYPE_ARRAY=2
+
 # PROCEDURES
 # =================================================================================================
 
@@ -136,3 +144,33 @@ end
 export.get_empty_smt_root
     push.15321474589252129342.17373224439259377994.15071539326562317628.3312677166725950353
 end
+
+#! Returns the type of storage slot value in the account storage.
+#!
+#! Stack: []
+#! Output: [type_storage_value]
+#!
+#! - type_storage_value is the type of storage slot item in the account storage.
+export.get_storage_slot_type_value
+    push.STORAGE_SLOT_TYPE_VALUE
+end
+
+#! Returns the type of storage slot map in the account storage.
+#!
+#! Stack: []
+#! Output: [type_storage_map]
+#!
+#! - type_storage_map is the type of storage slot item in the account storage.
+export.get_storage_slot_type_map
+    push.STORAGE_SLOT_TYPE_MAP
+end
+
+#! Returns the type of storage slot array in the account storage.
+#!
+#! Stack: []
+#! Output: [type_storage_array]
+#!
+#! - type_storage_array is the type of storage slot item in the account storage.
+export.get_storage_slot_type_array
+    push.STORAGE_SLOT_TYPE_ARRAY
+end

miden-tx/src/error.rs

@@ -241,8 +241,11 @@ const ERR_INVALID_NOTE_TYPE: u32 = 131140;
 const ERR_INVALID_NOTE_IDX: u32 = 131154;
 const ERR_NOTE_INVALID_TAG_PREFIX_FOR_TYPE: u32 = 131141;
 const ERR_NOTE_TAG_MUST_BE_U32: u32 = 131142;
+const ERR_SETTING_NON_VALUE_ITEM_ON_VALUE_SLOT: u32 = 131143;
+const ERR_SETTING_MAP_ITEM_ON_NON_MAP_SLOT: u32 = 131144;
+const ERR_READING_MAP_VALUE_FROM_NON_MAP_SLOT: u32 = 131145;
 
-pub const KERNEL_ERRORS: [(u32, &str); 72] = [
+pub const KERNEL_ERRORS: [(u32, &str); 75] = [
     (ERR_FAUCET_RESERVED_DATA_SLOT, "For faucets, storage slot 254 is reserved and can not be used with set_account_item procedure"),
     (ERR_ACCT_MUST_BE_A_FAUCET, "Procedure can only be called from faucet accounts"),
     (ERR_P2ID_WRONG_NUMBER_OF_INPUTS, "P2ID scripts expect exactly 1 note input"),
@@ -315,4 +318,7 @@ pub const KERNEL_ERRORS: [(u32, &str); 72] = [
     (ERR_INVALID_NOTE_IDX, "Invalid note index"),
     (ERR_NOTE_INVALID_TAG_PREFIX_FOR_TYPE, "The note's tag failed the most significant validation"),
     (ERR_NOTE_TAG_MUST_BE_U32, "The note's tag high bits must be set to 0"),
+    (ERR_SETTING_NON_VALUE_ITEM_ON_VALUE_SLOT, "Setting a non-value item on a value slot"),
+    (ERR_SETTING_MAP_ITEM_ON_NON_MAP_SLOT, "Setting a map item on a non-map slot"),
+    (ERR_READING_MAP_VALUE_FROM_NON_MAP_SLOT, "Slot type is not a map"),
 ];

miden-tx/src/tests/mod.rs

@@ -303,7 +303,8 @@ fn executed_transaction_account_delta() {
 
     // storage delta
     // --------------------------------------------------------------------------------------------
-    assert_eq!(executed_transaction.account_delta().storage().updated_items.len(), 2);
+    // We expect one updated item and one updated map
+    assert_eq!(executed_transaction.account_delta().storage().updated_items.len(), 1);
     assert_eq!(
         executed_transaction.account_delta().storage().updated_items[0].0,
         STORAGE_INDEX_0

objects/src/accounts/delta/storage.rs

@@ -95,12 +95,6 @@ impl AccountStorageDelta {
             if index > &MAX_MUTABLE_STORAGE_SLOT_IDX {
                 return Err(AccountDeltaError::ImmutableStorageSlot(*index as usize));
             }
-            // for every storage map delta there should be one corresponding storage item update
-            if !self.updated_items.iter().any(|(idx, _)| idx == index) {
-                return Err(AccountDeltaError::StorageMapDeltaWithoutStorageItemChange(
-                    *index as usize,
-                ));
-            }
             if !storage_map_delta.is_empty() {
                 storage_map_delta.validate()?;
             }

objects/src/accounts/mod.rs

@@ -404,10 +404,7 @@ mod tests {
         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)]),
-                (2_u8, storage_map.root().into()),
-            ])
+            .add_updated_items([(1_u8, [Felt::new(1), Felt::new(2), Felt::new(3), Felt::new(4)])])
             .add_updated_maps([(2_u8, updated_map)])
             .build()
             .unwrap();

objects/src/accounts/storage/map.rs

@@ -1,15 +1,14 @@
-use vm_core::EMPTY_WORD;
-
 use super::{
-    AccountError, ByteReader, ByteWriter, Deserializable, DeserializationError, Felt, Serializable,
-    Word,
+    AccountError, ByteReader, ByteWriter, Deserializable, DeserializationError, Digest, Felt,
+    Serializable, Word,
 };
 use crate::{
     accounts::StorageMapDelta,
     crypto::{
         hash::rpo::RpoDigest,
         merkle::{InnerNodeInfo, LeafIndex, Smt, SmtLeaf, SmtProof, SMT_DEPTH},
     },
+    EMPTY_WORD,
 };
 
 // ACCOUNT STORAGE MAP
@@ -111,35 +110,19 @@ impl StorageMap {
     ///
     /// This method assumes that the delta has been validated by the calling method and so, no
     /// additional validation of delta is performed.
-    pub fn apply_delta(&mut self, delta: &StorageMapDelta) -> Result<(), AccountError> {
+    pub fn apply_delta(&mut self, delta: &StorageMapDelta) -> Result<Digest, AccountError> {
         // apply the updated leaves to the storage map
         for &(key, value) in delta.updated_leaves.iter() {
-            self.set_map_item(key, value)?;
+            self.insert(key.into(), value);
         }
 
         // apply the cleared leaves to the storage map
         // currently we cannot remove leaves from the storage map, so we just set them to empty
         for &key in delta.cleared_leaves.iter() {
-            self.set_map_item(key, EMPTY_WORD)?;
-        }
-
-        Ok(())
-    }
-
-    /// Sets a map item from the storage at the specified index.
-    pub fn set_map_item(&mut self, key: Word, value: Word) -> Result<(Word, Word), AccountError> {
-        let old_map_root = self.root();
-        let old_value = self.get_value(&RpoDigest::from(key));
-
-        if value == EMPTY_WORD {
-            // if the value is empty, remove the leaf from the storage map
-            self.map.insert(key.into(), value);
-        } else {
-            // insert the value into the storage map
-            self.map.insert(key.into(), value);
+            self.insert(key.into(), EMPTY_WORD);
         }
 
-        Ok((old_map_root.into(), old_value))
+        Ok(self.root())
     }
 }