feat: [3/4] partial storage map queries

CHANGELOG.md

@@ -14,6 +14,7 @@
 - Renamed card's names in the `miden-network-monitor` binary ([#1441](https://github.com/0xMiden/miden-node/pull/1441)).
 - Improved tracing in `miden-network-monitor` binary ([#1366](https://github.com/0xMiden/miden-node/pull/1366)).
 - Integrated RPC stack with Validator component for transaction validation ([#1457](https://github.com/0xMiden/miden-node/pull/1457)).
+- Add partial storage map queries to RPC ([#1428](https://github.com/0xMiden/miden-node/pull/1428)).
 - Added explorer status to the `miden-network-monitor` binary ([#1450](https://github.com/0xMiden/miden-node/pull/1450)).
 
 ### Changes

crates/proto/src/domain/account.rs

@@ -15,7 +15,7 @@ use miden_objects::account::{
 use miden_objects::asset::{Asset, AssetVault};
 use miden_objects::block::BlockNumber;
 use miden_objects::block::account_tree::AccountWitness;
-use miden_objects::crypto::merkle::SparseMerklePath;
+use miden_objects::crypto::merkle::{MerkleError, SmtForest, SmtProof, SparseMerklePath};
 use miden_objects::note::{NoteExecutionMode, NoteTag};
 use miden_objects::utils::{Deserializable, DeserializationError, Serializable};
 use thiserror::Error;
@@ -194,6 +194,7 @@ impl TryFrom<proto::rpc::account_storage_details::AccountStorageMapDetails>
         value: proto::rpc::account_storage_details::AccountStorageMapDetails,
     ) -> Result<Self, Self::Error> {
         use proto::rpc::account_storage_details::account_storage_map_details::map_entries::StorageMapEntry;
+
         let proto::rpc::account_storage_details::AccountStorageMapDetails {
             slot_name,
             too_many_entries,
@@ -257,6 +258,7 @@ impl TryFrom<proto::rpc::account_proof_request::account_detail_request::StorageM
     }
 }
 
+/// Request of slot data values.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum SlotData {
     All,
@@ -466,6 +468,15 @@ pub enum StorageMapEntries {
 }
 
 /// Details about an account storage map slot.
+#[derive(Debug, Clone, PartialEq)]
+pub enum StorageMapData {
+    /// All entries are included used for small storage maps or when `all_entries` is requested.
+    AllEntries(Vec<(Word, Word)>),
+
+    /// Specific entries with their Merkle proofs for partial responses.
+    EntriesWithProofs(Vec<SmtProof>),
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct AccountStorageMapDetails {
     pub slot_name: StorageSlotName,
@@ -476,42 +487,101 @@ impl AccountStorageMapDetails {
     /// Maximum number of storage map entries that can be returned in a single response.
     pub const MAX_RETURN_ENTRIES: usize = 1000;
 
+    /// Creates storage map details with all entries from the storage map.
+    ///
+    /// If the storage map has too many entries (> `MAX_RETURN_ENTRIES`),
+    /// returns `LimitExceeded` variant.
+    pub fn from_all_entries(slot_name: StorageSlotName, storage_map: &StorageMap) -> Self {
+        if storage_map.num_entries() > Self::MAX_RETURN_ENTRIES {
+            Self {
+                slot_name,
+                entries: StorageMapEntries::LimitExceeded,
+            }
+        } else {
+            let map_entries = Vec::from_iter(storage_map.entries().map(|(k, v)| (*k, *v)));
+            Self {
+                slot_name,
+                entries: StorageMapEntries::Entries(map_entries),
+            }
+        }
+    }
+
+    /// Creates storage map details based on the requested slot data.
+    ///
+    /// Handles both "all entries" and "specific keys" requests.
+    /// Returns `LimitExceeded` if too many entries.
     pub fn new(slot_name: StorageSlotName, slot_data: SlotData, storage_map: &StorageMap) -> Self {
         match slot_data {
             SlotData::All => Self::from_all_entries(slot_name, storage_map),
-            SlotData::MapKeys(keys) => Self::from_specific_keys(slot_name, &keys[..], storage_map),
+            SlotData::MapKeys(keys) => {
+                if keys.len() > Self::MAX_RETURN_ENTRIES {
+                    Self {
+                        slot_name,
+                        entries: StorageMapEntries::LimitExceeded,
+                    }
+                } else {
+                    // Query specific keys from the storage map
+                    let mut entries = Vec::with_capacity(keys.len());
+                    for key in keys {
+                        let value = storage_map.get(&key);
+                        entries.push((key, value));
+                    }
+                    Self {
+                        slot_name,
+                        entries: StorageMapEntries::Entries(entries),
+                    }
+                }
+            },
         }
     }
 
-    fn from_all_entries(slot_name: StorageSlotName, storage_map: &StorageMap) -> Self {
-        if storage_map.num_entries() > Self::MAX_RETURN_ENTRIES {
+    /// Creates storage map details from forest-queried entries.
+    ///
+    /// Returns `LimitExceeded` if too many entries.
+    pub fn from_forest_entries(slot_name: StorageSlotName, entries: Vec<(Word, Word)>) -> Self {
+        if entries.len() > Self::MAX_RETURN_ENTRIES {
             Self {
                 slot_name,
                 entries: StorageMapEntries::LimitExceeded,
             }
         } else {
-            let map_entries = Vec::from_iter(storage_map.entries().map(|(k, v)| (*k, *v)));
             Self {
                 slot_name,
-                entries: StorageMapEntries::Entries(map_entries),
+                entries: StorageMapEntries::Entries(entries),
             }
         }
     }
 
-    fn from_specific_keys(
+    /// Creates storage map details with SMT proofs for specific keys.
+    ///
+    /// Returns `LimitExceeded` if too many keys, or `MerkleError` if the forest
+    /// doesn't contain sufficient data.
+    pub fn from_specific_keys(
         slot_name: StorageSlotName,
         keys: &[Word],
-        storage_map: &StorageMap,
-    ) -> Self {
+        storage_forest: &SmtForest,
+        smt_root: Word,
+    ) -> Result<Self, MerkleError> {
         if keys.len() > Self::MAX_RETURN_ENTRIES {
-            Self {
+            return Ok(Self {
                 slot_name,
                 entries: StorageMapEntries::LimitExceeded,
-            }
-        } else {
-            // TODO For now, we return all entries instead of specific keys with proofs
-            Self::from_all_entries(slot_name, storage_map)
+            });
         }
+
+        // Collect key-value pairs by opening proofs for each key
+        let mut entries = Vec::with_capacity(keys.len());
+
+        for key in keys {
+            let proof = storage_forest.open(smt_root, *key)?;
+            let value = proof.get(key).unwrap_or(miden_objects::EMPTY_WORD);
+            entries.push((*key, value));
+        }
+
+        Ok(Self {
+            slot_name,
+            entries: StorageMapEntries::Entries(entries),
+        })
     }
 }
 
@@ -707,7 +777,6 @@ impl From<AccountStorageMapDetails>
         }
     }
 }
-
 // ACCOUNT WITNESS
 // ================================================================================================
 

crates/store/src/errors.rs

@@ -127,6 +127,12 @@ pub enum DatabaseError {
     SqlValueConversion(#[from] DatabaseTypeConversionError),
     #[error("Not implemented: {0}")]
     NotImplemented(String),
+    #[error("storage root not found for account {account_id}, slot {slot_name}, block {block_num}")]
+    StorageRootNotFound {
+        account_id: AccountId,
+        slot_name: String,
+        block_num: BlockNumber,
+    },
 }
 
 impl DatabaseError {

crates/store/src/state.rs

@@ -1191,7 +1191,6 @@ impl State {
             let storage_map = match slot.content() {
                 StorageSlotContent::Map(map) => map,
                 StorageSlotContent::Value(_) => {
-                    // TODO: what to do with value entries? Is it ok to ignore them?
                     return Err(AccountError::StorageSlotNotMap(slot_name).into());
                 },
             };