feat: allow updates to all item slots (#737)

bench-tx/src/utils.rs

@@ -3,7 +3,7 @@ pub use alloc::collections::BTreeMap;
 
 use miden_lib::transaction::TransactionKernel;
 use miden_objects::{
-    accounts::{Account, AccountCode, AccountId, AccountStorage, SlotItem, StorageSlot},
+    accounts::{Account, AccountCode, AccountId, AccountStorage, SlotItem},
     assembly::ModuleAst,
     assets::{Asset, AssetVault},
     Felt, Word,
@@ -88,14 +88,8 @@ pub fn get_account_with_default_account_code(
     let account_assembler = TransactionKernel::assembler();
 
     let account_code = AccountCode::new(account_code_ast.clone(), &account_assembler).unwrap();
-    let account_storage = AccountStorage::new(
-        vec![SlotItem {
-            index: 0,
-            slot: StorageSlot::new_value(public_key),
-        }],
-        vec![],
-    )
-    .unwrap();
+    let account_storage =
+        AccountStorage::new(vec![SlotItem::new_value(0, 0, public_key)], BTreeMap::new()).unwrap();
 
     let account_vault = match assets {
         Some(asset) => AssetVault::new(&[asset]).unwrap(),

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

@@ -403,7 +403,7 @@ 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
     # => [KEY, NEW_VALUE, index, ...]

miden-lib/src/accounts/faucets/mod.rs

@@ -1,9 +1,8 @@
-use alloc::string::ToString;
+use alloc::{collections::BTreeMap, string::ToString};
 
 use miden_objects::{
     accounts::{
         Account, AccountCode, AccountId, AccountStorage, AccountStorageType, AccountType, SlotItem,
-        StorageSlot,
     },
     assembly::LibraryPath,
     assets::TokenSymbol,
@@ -72,17 +71,8 @@ pub fn create_basic_fungible_faucet(
     // - slot 0: authentication data
     // - slot 1: token metadata as [max_supply, decimals, token_symbol, 0]
     let account_storage = AccountStorage::new(
-        vec![
-            SlotItem {
-                index: 0,
-                slot: StorageSlot::new_value(auth_data),
-            },
-            SlotItem {
-                index: 1,
-                slot: StorageSlot::new_value(metadata),
-            },
-        ],
-        vec![],
+        vec![SlotItem::new_value(0, 0, auth_data), SlotItem::new_value(1, 0, metadata)],
+        BTreeMap::new(),
     )?;
 
     let account_seed = AccountId::get_account_seed(

miden-lib/src/accounts/wallets/mod.rs

@@ -1,9 +1,11 @@
-use alloc::string::{String, ToString};
+use alloc::{
+    collections::BTreeMap,
+    string::{String, ToString},
+};
 
 use miden_objects::{
     accounts::{
-        Account, AccountCode, AccountId, AccountStorage, AccountStorageType, AccountType,
-        StorageSlot,
+        Account, AccountCode, AccountId, AccountStorage, AccountStorageType, AccountType, SlotItem,
     },
     assembly::ModuleAst,
     AccountError, Word,
@@ -58,13 +60,8 @@ pub fn create_basic_wallet(
     let account_assembler = TransactionKernel::assembler();
     let account_code = AccountCode::new(account_code_ast.clone(), &account_assembler)?;
 
-    let account_storage = AccountStorage::new(
-        vec![miden_objects::accounts::SlotItem {
-            index: 0,
-            slot: StorageSlot::new_value(storage_slot_0_data),
-        }],
-        vec![],
-    )?;
+    let account_storage =
+        AccountStorage::new(vec![SlotItem::new_value(0, 0, storage_slot_0_data)], BTreeMap::new())?;
 
     let account_seed = AccountId::get_account_seed(
         init_seed,

miden-lib/src/transaction/inputs.rs

@@ -223,7 +223,7 @@ fn add_account_to_advice_inputs(
 
     // If there are storage maps, we populate the merkle store and advice map
     if !account.storage().maps().is_empty() {
-        for map in account.storage().maps() {
+        for map in account.storage().maps().values() {
             // extend the merkle store and map with the storage maps
             inputs.extend_merkle_store(map.inner_nodes());
 

miden-tx/src/kernel_tests/test_account.rs

@@ -10,15 +10,16 @@ use miden_objects::{
             ACCOUNT_ID_REGULAR_ACCOUNT_IMMUTABLE_CODE_ON_CHAIN,
             ACCOUNT_ID_REGULAR_ACCOUNT_UPDATABLE_CODE_OFF_CHAIN,
         },
-        AccountId, AccountType, StorageSlotType,
+        AccountId, AccountType, SlotItem, StorageSlotType,
     },
     crypto::{hash::rpo::RpoDigest, merkle::LeafIndex},
     testing::{
         account::MockAccountType,
         notes::AssetPreservationStatus,
         prepare_word,
         storage::{
-            storage_item_0, storage_item_1, storage_item_2, storage_map_2, STORAGE_LEAVES_2,
+            storage_map_2, STORAGE_INDEX_0, STORAGE_INDEX_1, STORAGE_INDEX_2, STORAGE_LEAVES_2,
+            STORAGE_VALUE_0, STORAGE_VALUE_1,
         },
     },
 };
@@ -246,7 +247,10 @@ fn test_is_faucet_procedure() {
 
 #[test]
 fn test_get_item() {
-    for storage_item in [storage_item_0(), storage_item_1()] {
+    for storage_item in [
+        SlotItem::new_value(STORAGE_INDEX_0, 0, STORAGE_VALUE_0),
+        SlotItem::new_value(STORAGE_INDEX_1, 0, STORAGE_VALUE_1),
+    ] {
         let (tx_inputs, tx_args) =
             mock_inputs(MockAccountType::StandardExisting, AssetPreservationStatus::Preserved);
 
@@ -335,7 +339,11 @@ fn test_set_item() {
 // Test different account storage types
 #[test]
 fn test_get_storage_data_type() {
-    for storage_item in [storage_item_0(), storage_item_1(), storage_item_2()] {
+    for storage_item in [
+        SlotItem::new_value(STORAGE_INDEX_0, 0, STORAGE_VALUE_0),
+        SlotItem::new_value(STORAGE_INDEX_1, 0, STORAGE_VALUE_1),
+        SlotItem::new_map(STORAGE_INDEX_2, 0, storage_map_2().root().into()),
+    ] {
         let (tx_inputs, tx_args) =
             mock_inputs(MockAccountType::StandardExisting, AssetPreservationStatus::Preserved);
 
@@ -385,7 +393,7 @@ fn test_get_map_item() {
     let (tx_inputs, tx_args) =
         mock_inputs(MockAccountType::StandardExisting, AssetPreservationStatus::Preserved);
 
-    let storage_item = storage_item_2();
+    let storage_item = SlotItem::new_map(STORAGE_INDEX_2, 0, storage_map_2().root().into());
     for (key, value) in STORAGE_LEAVES_2 {
         let code = format!(
             "
@@ -433,7 +441,7 @@ fn test_set_map_item() {
     let (tx_inputs, tx_args) =
         mock_inputs(MockAccountType::StandardExisting, AssetPreservationStatus::Preserved);
 
-    let storage_item = storage_item_2();
+    let storage_item = SlotItem::new_map(STORAGE_INDEX_2, 0, storage_map_2().root().into());
 
     let code = format!(
         "

miden-tx/tests/integration/main.rs

@@ -5,7 +5,7 @@ use miden_lib::transaction::TransactionKernel;
 use miden_objects::{
     accounts::{
         account_id::testing::ACCOUNT_ID_SENDER, Account, AccountCode, AccountId, AccountStorage,
-        SlotItem, StorageSlot,
+        SlotItem,
     },
     assembly::{ModuleAst, ProgramAst},
     assets::{Asset, AssetVault, FungibleAsset},
@@ -73,20 +73,16 @@ pub fn get_account_with_default_account_code(
     public_key: Word,
     assets: Option<Asset>,
 ) -> Account {
+    use std::collections::BTreeMap;
+
     use miden_objects::testing::account_code::DEFAULT_ACCOUNT_CODE;
     let account_code_src = DEFAULT_ACCOUNT_CODE;
     let account_code_ast = ModuleAst::parse(account_code_src).unwrap();
     let account_assembler = TransactionKernel::assembler();
 
     let account_code = AccountCode::new(account_code_ast.clone(), &account_assembler).unwrap();
-    let account_storage = AccountStorage::new(
-        vec![SlotItem {
-            index: 0,
-            slot: StorageSlot::new_value(public_key),
-        }],
-        vec![],
-    )
-    .unwrap();
+    let account_storage =
+        AccountStorage::new(vec![SlotItem::new_value(0, 0, public_key)], BTreeMap::new()).unwrap();
 
     let account_vault = match assets {
         Some(asset) => AssetVault::new(&[asset]).unwrap(),

miden-tx/tests/integration/scripts/faucet.rs

@@ -1,5 +1,7 @@
 extern crate alloc;
 
+use std::collections::BTreeMap;
+
 use miden_lib::{
     accounts::faucets::create_basic_fungible_faucet,
     transaction::{memory::FAUCET_STORAGE_DATA_SLOT, TransactionKernel},
@@ -8,7 +10,7 @@ use miden_lib::{
 use miden_objects::{
     accounts::{
         account_id::testing::ACCOUNT_ID_FUNGIBLE_FAUCET_OFF_CHAIN, Account, AccountCode, AccountId,
-        AccountStorage, AccountStorageType, SlotItem, StorageSlot,
+        AccountStorage, AccountStorageType, SlotItem,
     },
     assembly::{ModuleAst, ProgramAst},
     assets::{Asset, AssetVault, FungibleAsset, TokenSymbol},
@@ -294,16 +296,10 @@ fn get_faucet_account_with_max_supply_and_total_issuance(
     let faucet_storage_slot_1 = [Felt::new(max_supply), Felt::new(0), Felt::new(0), Felt::new(0)];
     let mut faucet_account_storage = AccountStorage::new(
         vec![
-            SlotItem {
-                index: 0,
-                slot: StorageSlot::new_value(public_key),
-            },
-            SlotItem {
-                index: 1,
-                slot: StorageSlot::new_value(faucet_storage_slot_1),
-            },
+            SlotItem::new_value(0, 0, public_key),
+            SlotItem::new_value(1, 0, faucet_storage_slot_1),
         ],
-        vec![],
+        BTreeMap::new(),
     )
     .unwrap();
 

miden-tx/tests/integration/wallet/mod.rs

@@ -1,11 +1,13 @@
+use std::collections::BTreeMap;
+
 use miden_lib::{accounts::wallets::create_basic_wallet, AuthScheme};
 use miden_objects::{
     accounts::{
         account_id::testing::{
             ACCOUNT_ID_FUNGIBLE_FAUCET_ON_CHAIN, ACCOUNT_ID_OFF_CHAIN_SENDER,
             ACCOUNT_ID_REGULAR_ACCOUNT_UPDATABLE_CODE_OFF_CHAIN,
         },
-        Account, AccountId, AccountStorage, SlotItem, StorageSlot,
+        Account, AccountId, AccountStorage, SlotItem,
     },
     assembly::ProgramAst,
     assets::{Asset, AssetVault, FungibleAsset},
@@ -85,14 +87,9 @@ fn prove_receive_asset_via_wallet() {
     assert_eq!(executed_transaction.account_delta().nonce(), Some(Felt::new(2)));
 
     // clone account info
-    let account_storage = AccountStorage::new(
-        vec![SlotItem {
-            index: 0,
-            slot: StorageSlot::new_value(target_pub_key),
-        }],
-        vec![],
-    )
-    .unwrap();
+    let account_storage =
+        AccountStorage::new(vec![SlotItem::new_value(0, 0, target_pub_key)], BTreeMap::new())
+            .unwrap();
     let account_code = target_account.code().clone();
     // vault delta
     let target_account_after: Account = Account::from_parts(
@@ -167,14 +164,9 @@ fn prove_send_asset_via_wallet() {
     assert!(prove_and_verify_transaction(executed_transaction.clone()).is_ok());
 
     // clones account info
-    let sender_account_storage = AccountStorage::new(
-        vec![SlotItem {
-            index: 0,
-            slot: StorageSlot::new_value(sender_pub_key),
-        }],
-        vec![],
-    )
-    .unwrap();
+    let sender_account_storage =
+        AccountStorage::new(vec![SlotItem::new_value(0, 0, sender_pub_key)], BTreeMap::new())
+            .unwrap();
     let sender_account_code = sender_account.code().clone();
 
     // vault delta

objects/src/accounts/data.rs

@@ -94,6 +94,8 @@ impl Deserializable for AccountData {
 
 #[cfg(test)]
 mod tests {
+    use alloc::collections::BTreeMap;
+
     use miden_crypto::{
         dsa::rpo_falcon512::SecretKey,
         utils::{Deserializable, Serializable},
@@ -118,7 +120,7 @@ mod tests {
 
         // create account and auth
         let vault = AssetVault::new(&[]).unwrap();
-        let storage = AccountStorage::new(vec![], vec![]).unwrap();
+        let storage = AccountStorage::new(vec![], BTreeMap::new()).unwrap();
         let nonce = Felt::new(0);
         let account = Account::from_parts(id, vault, storage, code, nonce);
         let account_seed = Some(Word::default());