Add testing accessors for AccountStoragePatch

crates/miden-protocol/src/account/patch/storage.rs

@@ -633,6 +633,35 @@ mod tests {
         });
     }
 
+    #[test]
+    fn account_storage_patch_accessors() {
+        let value_slot = StorageSlotName::mock(1);
+        let map_slot = StorageSlotName::mock(2);
+        let absent_slot = StorageSlotName::mock(3);
+
+        let value = Word::from([1u32, 2, 3, 4]);
+        let map_key = StorageMapKey::from_array([10, 11, 12, 13]);
+        let map_value = Word::from([5u32, 6, 7, 8]);
+        let absent_key = StorageMapKey::from_array([99, 99, 99, 99]);
+
+        let patch = AccountStoragePatch::from_iters(
+            [],
+            [(value_slot.clone(), value)],
+            [(map_slot.clone(), StorageMapPatch::from_iters([], [(map_key, map_value)]))],
+        );
+
+        assert_eq!(patch.get_value(&value_slot), Some(value));
+        assert_eq!(patch.get_value(&absent_slot), None);
+
+        let map_patch = patch.get_map(&map_slot).unwrap();
+        assert_eq!(map_patch.entries().get(&map_key), Some(&map_value));
+        assert_eq!(patch.get_map(&absent_slot), None);
+
+        assert_eq!(patch.get_map_value(&map_slot, &map_key), Some(map_value));
+        assert_eq!(patch.get_map_value(&map_slot, &absent_key), None);
+        assert_eq!(patch.get_map_value(&absent_slot, &map_key), None);
+    }
+
     #[test]
     fn test_is_empty() {
         let storage_patch = AccountStoragePatch::new();

crates/miden-protocol/src/testing/storage.rs

@@ -44,6 +44,37 @@ impl AccountStoragePatch {
         Self::from_raw(deltas)
     }
 
+    // ACCESSORS
+    // -------------------------------------------------------------------------------------------
+
+    /// Returns the updated value for the given slot, or `None` if the slot was not updated.
+    ///
+    /// # Panics
+    /// Panics if the slot patch is a map.
+    pub fn get_value(&self, slot_name: &StorageSlotName) -> Option<Word> {
+        self.get(slot_name).cloned().map(StorageSlotPatch::unwrap_value)
+    }
+
+    /// Returns the map patch for the given slot, or `None` if the slot was not updated.
+    ///
+    /// # Panics
+    /// Panics if the slot patch is a value.
+    pub fn get_map(&self, slot_name: &StorageSlotName) -> Option<&StorageMapPatch> {
+        self.get(slot_name).map(|patch| match patch {
+            StorageSlotPatch::Map(map_patch) => map_patch,
+            StorageSlotPatch::Value(_) => panic!("called get_map on a value slot patch"),
+        })
+    }
+
+    /// Returns the updated value for the given map entry, or `None` if the slot or key was not
+    /// updated.
+    ///
+    /// # Panics
+    /// Panics if the slot patch is a value.
+    pub fn get_map_value(&self, slot_name: &StorageSlotName, key: &StorageMapKey) -> Option<Word> {
+        self.get_map(slot_name)?.entries().get(key).copied()
+    }
+
     // MUTATORS
     // -------------------------------------------------------------------------------------------
 

crates/miden-testing/src/kernel_tests/tx/test_account_delta.rs

@@ -16,7 +16,6 @@ use miden_protocol::account::{
     StorageMapKey,
     StorageSlot,
     StorageSlotName,
-    StorageSlotPatch,
 };
 use miden_protocol::asset::{
     Asset,
@@ -771,24 +770,18 @@ async fn asset_and_storage_patch() -> anyhow::Result<()> {
     // We expect one updated item and one updated map
     assert_eq!(executed_transaction.account_delta().storage().values().count(), 1);
     assert_eq!(
-        executed_transaction
-            .account_delta()
-            .storage()
-            .get(&MOCK_VALUE_SLOT0)
-            .cloned()
-            .map(StorageSlotPatch::unwrap_value),
+        executed_transaction.account_delta().storage().get_value(&MOCK_VALUE_SLOT0),
         Some(updated_slot_value)
     );
 
     assert_eq!(executed_transaction.account_delta().storage().maps().count(), 1);
-    let map_patch = executed_transaction
-        .account_delta()
-        .storage()
-        .get(&MOCK_MAP_SLOT)
-        .cloned()
-        .map(StorageSlotPatch::unwrap_map)
-        .unwrap();
-    assert_eq!(*map_patch.entries().get(&updated_map_key).unwrap(), updated_map_value);
+    assert_eq!(
+        executed_transaction
+            .account_delta()
+            .storage()
+            .get_map_value(&MOCK_MAP_SLOT, &updated_map_key),
+        Some(updated_map_value)
+    );
 
     // vault delta
     // --------------------------------------------------------------------------------------------
@@ -897,18 +890,9 @@ async fn proven_tx_storage_maps_matches_executed_tx_for_new_account() -> anyhow:
     for (slot_name, expected_map) in
         [(map0_slot_name, map0), (map1_slot_name, map1), (map2_slot_name, map2)]
     {
-        let map_patch = tx
-            .account_delta()
-            .storage()
-            .get(&slot_name)
-            .cloned()
-            .map(StorageSlotPatch::unwrap_map)
-            .unwrap();
-        assert_eq!(
-            map_patch.entries().iter().collect::<BTreeMap<_, _>>(),
-            expected_map.entries().collect(),
-            "map delta does not match for slot {slot_name}",
-        );
+        let map_patch_entries = tx.account_delta().storage().get_map(&slot_name).unwrap().entries();
+        let expected: BTreeMap<_, _> = expected_map.entries().map(|(k, v)| (*k, *v)).collect();
+        assert_eq!(map_patch_entries, &expected, "map delta does not match for slot {slot_name}",);
     }
 
     let proven_tx = LocalTransactionProver::default().prove_dummy(tx.clone())?;
@@ -966,24 +950,8 @@ async fn delta_for_new_account_retains_empty_value_storage_slots() -> anyhow::Re
     };
 
     assert_eq!(delta.storage().values().count(), 2);
-    assert_eq!(
-        delta
-            .storage()
-            .get(&slot_name0)
-            .cloned()
-            .map(StorageSlotPatch::unwrap_value)
-            .unwrap(),
-        Word::empty()
-    );
-    assert_eq!(
-        delta
-            .storage()
-            .get(&slot_name1)
-            .cloned()
-            .map(StorageSlotPatch::unwrap_value)
-            .unwrap(),
-        slot_value2
-    );
+    assert_eq!(delta.storage().get_value(&slot_name0), Some(Word::empty()));
+    assert_eq!(delta.storage().get_value(&slot_name1), Some(slot_value2));
 
     let recreated_account = Account::try_from(delta)?;
     // The recreated account should match the original account with the nonce incremented (and the
@@ -1017,15 +985,7 @@ async fn delta_for_new_account_retains_empty_map_storage_slots() -> anyhow::Resu
     };
 
     assert_eq!(delta.storage().maps().count(), 1);
-    assert!(
-        delta
-            .storage()
-            .get(&slot_name0)
-            .cloned()
-            .map(StorageSlotPatch::unwrap_map)
-            .unwrap()
-            .is_empty()
-    );
+    assert!(delta.storage().get_map(&slot_name0).unwrap().is_empty());
 
     let recreated_account = Account::try_from(delta)?;
     // The recreated account should match the original account with the nonce incremented (and the

crates/miden-testing/src/kernel_tests/tx/test_lazy_loading.rs

@@ -2,7 +2,7 @@
 //!
 //! Once lazy loading is enabled generally, it can be removed and/or integrated into other tests.
 
-use miden_protocol::account::{AccountId, AccountStorage, StorageMapKey, StorageSlotPatch};
+use miden_protocol::account::{AccountId, AccountStorage, StorageMapKey};
 use miden_protocol::asset::{Asset, FungibleAsset};
 use miden_protocol::testing::account_id::{
     ACCOUNT_ID_FEE_FAUCET,
@@ -234,13 +234,7 @@ async fn setting_map_item_with_lazy_loading_succeeds() -> anyhow::Result<()> {
         .execute()
         .await?;
 
-    let map_patch = tx
-        .account_delta()
-        .storage()
-        .get(mock_map_slot)
-        .cloned()
-        .map(StorageSlotPatch::unwrap_map)
-        .unwrap();
+    let map_patch = tx.account_delta().storage().get_map(mock_map_slot).unwrap();
     assert_eq!(map_patch.entries().get(&existing_key).unwrap(), &value0);
     assert_eq!(map_patch.entries().get(&non_existent_key).unwrap(), &value1);