feat(FPI): Make storage maps usable by querying the node

CHANGELOG.md

@@ -10,6 +10,7 @@
 * Added caching for foreign account code (#597).
 * Added support for unauthenticated notes consumption in the CLI (#609).
 * [BREAKING] Added foreign procedure invocation support for private accounts (#619).
+* [BREAKING] Added support for specifying map storage slots for FPI (#645)
 
 ### Fixes
 

Makefile

@@ -13,7 +13,7 @@ WARNINGS=RUSTDOCFLAGS="-D warnings"
 
 NODE_DIR="miden-node"
 NODE_REPO="https://github.com/0xPolygonMiden/miden-node.git"
-NODE_BRANCH="next"
+NODE_BRANCH="igamigo-map-proof"
 NODE_FEATURES_TESTING=--features "testing"
 
 PROVER_DIR="miden-base"

crates/rust-client/Cargo.toml

@@ -57,7 +57,7 @@ uuid = { version = "1.10", features = ["serde", "v4"] }
 web-sys = { version = "0.3", features = ["console"]}
 
 [build-dependencies]
-miden-rpc-proto = { git = "https://github.com/0xPolygonMiden/miden-node", branch = "next" }
+miden-rpc-proto = { git = "https://github.com/0xPolygonMiden/miden-node", branch = "igamigo-map-proof" }
 miden-lib = { workspace = true }
 miette = { version = "7.2", features = ["fancy"] }
 prost = { version = "0.13", default-features = false, features = ["derive"] }

crates/rust-client/src/mock.rs

@@ -42,6 +42,7 @@ use crate::{
         NodeRpcClient, RpcError,
     },
     store::{sqlite_store::SqliteStore, StoreAuthenticator},
+    transactions::ForeignAccount,
     Client,
 };
 
@@ -288,9 +289,8 @@ impl NodeRpcClient for MockRpcApi {
 
     async fn get_account_proofs(
         &mut self,
-        _account_ids: &BTreeSet<AccountId>,
+        _account_ids: &BTreeSet<ForeignAccount>,
         _code_commitments: Vec<AccountCode>,
-        _include_headers: bool,
     ) -> Result<AccountProofs, RpcError> {
         // TODO: Implement fully
         Ok((self.blocks.last().unwrap().header().block_num(), vec![]))

crates/rust-client/src/rpc/domain/accounts.rs

@@ -7,22 +7,20 @@ use core::fmt::{self, Debug, Display, Formatter};
 
 use miden_objects::{
     accounts::{Account, AccountCode, AccountHeader, AccountId, AccountStorageHeader},
-    crypto::merkle::MerklePath,
+    crypto::merkle::{MerklePath, SmtProof},
     Digest, Felt,
 };
 use miden_tx::utils::{Deserializable, Serializable, ToHex};
 use thiserror::Error;
 
-use crate::{
-    rpc::{
-        errors::RpcConversionError,
-        generated::{
-            account::{AccountHeader as ProtoAccountHeader, AccountId as ProtoAccountId},
-            responses::AccountStateHeader as ProtoAccountStateHeader,
-        },
-        RpcError,
+use crate::rpc::{
+    errors::RpcConversionError,
+    generated::{
+        account::{AccountHeader as ProtoAccountHeader, AccountId as ProtoAccountId},
+        requests::get_account_proofs_request,
+        responses::{AccountStateHeader as ProtoAccountStateHeader, StorageSlotMapProof},
     },
-    transactions::ForeignAccountInputs,
+    RpcError,
 };
 
 // ACCOUNT DETAILS
@@ -156,7 +154,12 @@ impl ProtoAccountStateHeader {
         account_id: AccountId,
         known_account_codes: &BTreeMap<Digest, AccountCode>,
     ) -> Result<StateHeaders, RpcError> {
-        let ProtoAccountStateHeader { header, storage_header, account_code } = self;
+        let ProtoAccountStateHeader {
+            header,
+            storage_header,
+            account_code,
+            storage_slots,
+        } = self;
         let account_header = header
             .ok_or(RpcError::ExpectedDataMissing("Account.StateHeader".to_string()))?
             .into_domain(account_id)?;
@@ -181,7 +184,22 @@ impl ProtoAccountStateHeader {
             }
         };
 
-        Ok(StateHeaders { account_header, storage_header, code })
+        // Get map values into slot |-> [proof0, proof1, ...] mapping
+        let mut storage_slot_proofs: BTreeMap<u8, Vec<SmtProof>> = BTreeMap::new();
+        for StorageSlotMapProof { storage_slot, smt_proof } in storage_slots {
+            let proof = SmtProof::read_from_bytes(&smt_proof)?;
+            match storage_slot_proofs.get_mut(&(storage_slot as u8)) {
+                Some(list) => list.push(proof),
+                None => _ = storage_slot_proofs.insert(storage_slot as u8, vec![proof]),
+            }
+        }
+
+        Ok(StateHeaders {
+            account_header,
+            storage_header,
+            code,
+            storage_slots: storage_slot_proofs,
+        })
     }
 }
 
@@ -196,6 +214,7 @@ pub struct StateHeaders {
     pub account_header: AccountHeader,
     pub storage_header: AccountStorageHeader,
     pub code: AccountCode,
+    pub storage_slots: BTreeMap<StorageSlotIndex, Vec<SmtProof>>,
 }
 
 /// Represents a proof of existence of an account's state at a specific block number.
@@ -217,7 +236,10 @@ impl AccountProof {
         account_hash: Digest,
         state_headers: Option<StateHeaders>,
     ) -> Result<Self, AccountProofError> {
-        if let Some(StateHeaders { account_header, storage_header: _, code }) = &state_headers {
+        if let Some(StateHeaders {
+            account_header, code, ..
+        }) = &state_headers
+        {
             if account_header.hash() != account_hash {
                 return Err(AccountProofError::InconsistentAccountHash);
             }
@@ -275,55 +297,66 @@ impl AccountProof {
     }
 }
 
-// FPI ACCOUNT DATA
+// ACCOUNT STORAGE REQUEST
 // ================================================================================================
 
-/// Represents the data needed to perform a Foreign Procedure Invocation (FPI) on an account.
-pub struct FpiAccountData {
-    /// Account ID.
-    account_id: AccountId,
-    /// Authentication path from the `account_root` of the block header to the account.
-    merkle_path: MerklePath,
-    /// Account inputs needed to perform the FPI.
-    inputs: ForeignAccountInputs,
-}
+pub type StorageSlotIndex = u8;
+pub type StorageMapKey = Digest;
 
-impl FpiAccountData {
-    pub fn new(
-        account_id: AccountId,
-        merkle_path: MerklePath,
-        inputs: ForeignAccountInputs,
-    ) -> Self {
-        Self { account_id, merkle_path, inputs }
-    }
+/// Describes storage slots indices to be requested, as well as a list of keys for each of those
+/// slots.
+#[derive(Clone, Debug, Default, Eq, PartialEq)]
+pub struct AccountStorageRequirements(BTreeMap<StorageSlotIndex, Vec<StorageMapKey>>);
 
-    pub fn account_id(&self) -> AccountId {
-        self.account_id
+impl AccountStorageRequirements {
+    pub fn new<'a>(
+        slots_and_keys: impl IntoIterator<
+            Item = (StorageSlotIndex, impl IntoIterator<Item = &'a StorageMapKey>),
+        >,
+    ) -> Self {
+        let map = slots_and_keys
+            .into_iter()
+            .map(|(slot_index, keys_iter)| {
+                let keys_vec: Vec<StorageMapKey> = keys_iter.into_iter().cloned().collect();
+                (slot_index, keys_vec)
+            })
+            .collect();
+
+        AccountStorageRequirements(map)
     }
 
-    pub fn merkle_path(&self) -> &MerklePath {
-        &self.merkle_path
+    pub fn inner(&self) -> &BTreeMap<StorageSlotIndex, Vec<StorageMapKey>> {
+        &self.0
     }
+}
 
-    pub fn inputs(&self) -> &ForeignAccountInputs {
-        &self.inputs
+impl From<AccountStorageRequirements> for Vec<get_account_proofs_request::StorageRequest> {
+    fn from(value: AccountStorageRequirements) -> Vec<get_account_proofs_request::StorageRequest> {
+        let mut requests = Vec::with_capacity(value.0.len());
+        for (slot_index, map_keys) in value.0.into_iter() {
+            requests.push(get_account_proofs_request::StorageRequest {
+                storage_slot_index: slot_index as u32,
+                map_keys: map_keys
+                    .into_iter()
+                    .map(crate::rpc::generated::digest::Digest::from)
+                    .collect(),
+            });
+        }
+        requests
     }
+}
 
-    pub fn into_parts(self) -> (AccountId, MerklePath, ForeignAccountInputs) {
-        (self.account_id, self.merkle_path, self.inputs)
+impl Serializable for AccountStorageRequirements {
+    fn write_into<W: miden_tx::utils::ByteWriter>(&self, target: &mut W) {
+        target.write(&self.0);
     }
 }
 
-impl TryFrom<AccountProof> for FpiAccountData {
-    type Error = RpcError;
-
-    fn try_from(value: AccountProof) -> Result<Self, Self::Error> {
-        let (account_id, merkle_proof, _, state_headers) = value.into_parts();
-        if let Some(StateHeaders { account_header, storage_header, code }) = state_headers {
-            let inputs = ForeignAccountInputs::new(account_header, storage_header, code);
-            return Ok(FpiAccountData::new(account_id, merkle_proof, inputs));
-        }
-        Err(RpcError::ExpectedDataMissing(String::from("AccountProof.StateHeaders")))
+impl Deserializable for AccountStorageRequirements {
+    fn read_from<R: miden_tx::utils::ByteReader>(
+        source: &mut R,
+    ) -> Result<Self, miden_tx::utils::DeserializationError> {
+        Ok(AccountStorageRequirements(source.read()?))
     }
 }
 

crates/rust-client/src/rpc/generated/nostd/requests.rs

@@ -148,12 +148,15 @@ pub struct GetAccountStateDeltaRequest {
     #[prost(fixed32, tag = "3")]
     pub to_block_num: u32,
 }
+/// Request message to get account proofs.
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct GetAccountProofsRequest {
-    /// List of account IDs to get states.
+    /// A list of account requests, including map keys + values.
     #[prost(message, repeated, tag = "1")]
-    pub account_ids: ::prost::alloc::vec::Vec<super::account::AccountId>,
-    /// Optional flag to include header and account code in the response. `false` by default.
+    pub account_requests: ::prost::alloc::vec::Vec<
+        get_account_proofs_request::AccountRequest,
+    >,
+    /// Optional flag to include header and account code in the response. False by default.
     #[prost(bool, optional, tag = "2")]
     pub include_headers: ::core::option::Option<bool>,
     /// Account code commitments corresponding to the last-known `AccountCode` for requested
@@ -163,3 +166,27 @@ pub struct GetAccountProofsRequest {
     #[prost(message, repeated, tag = "3")]
     pub code_commitments: ::prost::alloc::vec::Vec<super::digest::Digest>,
 }
+/// Nested message and enum types in `GetAccountProofsRequest`.
+pub mod get_account_proofs_request {
+    /// Represents per-account requests where each account ID has its own list of
+    /// (storage_slot_index, map_keys) pairs.
+    #[derive(Clone, PartialEq, ::prost::Message)]
+    pub struct AccountRequest {
+        /// The account ID for this request.
+        #[prost(message, optional, tag = "1")]
+        pub account_id: ::core::option::Option<super::super::account::AccountId>,
+        /// List of storage requests for this account.
+        #[prost(message, repeated, tag = "2")]
+        pub storage_requests: ::prost::alloc::vec::Vec<StorageRequest>,
+    }
+    /// Represents a storage slot index and the associated map keys.
+    #[derive(Clone, PartialEq, ::prost::Message)]
+    pub struct StorageRequest {
+        /// Storage slot index (\[0..255\])
+        #[prost(uint32, tag = "1")]
+        pub storage_slot_index: u32,
+        /// A list of map keys (Digests) associated with this storage slot.
+        #[prost(message, repeated, tag = "2")]
+        pub map_keys: ::prost::alloc::vec::Vec<super::super::digest::Digest>,
+    }
+}

crates/rust-client/src/rpc/generated/nostd/responses.rs

@@ -227,8 +227,21 @@ pub struct AccountStateHeader {
     /// Values of all account storage slots (max 255).
     #[prost(bytes = "vec", tag = "2")]
     pub storage_header: ::prost::alloc::vec::Vec<u8>,
-    /// Account code, returned only when none of the request's code commitments match with the
-    /// current one.
+    /// Account code, returned only when none of the request's code commitments match
+    /// the current one.
     #[prost(bytes = "vec", optional, tag = "3")]
     pub account_code: ::core::option::Option<::prost::alloc::vec::Vec<u8>>,
+    /// Storage slots information for this account
+    #[prost(message, repeated, tag = "4")]
+    pub storage_slots: ::prost::alloc::vec::Vec<StorageSlotMapProof>,
+}
+/// Represents a single storage slot with the reuqested keys and their respective values.
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct StorageSlotMapProof {
+    /// The storage slot index (\[0..255\]).
+    #[prost(uint32, tag = "1")]
+    pub storage_slot: u32,
+    /// Merkle proof of the map value
+    #[prost(bytes = "vec", tag = "2")]
+    pub smt_proof: ::prost::alloc::vec::Vec<u8>,
 }

crates/rust-client/src/rpc/generated/std/requests.rs

@@ -148,12 +148,15 @@ pub struct GetAccountStateDeltaRequest {
     #[prost(fixed32, tag = "3")]
     pub to_block_num: u32,
 }
+/// Request message to get account proofs.
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct GetAccountProofsRequest {
-    /// List of account IDs to get states.
+    /// A list of account requests, including map keys + values.
     #[prost(message, repeated, tag = "1")]
-    pub account_ids: ::prost::alloc::vec::Vec<super::account::AccountId>,
-    /// Optional flag to include header and account code in the response. `false` by default.
+    pub account_requests: ::prost::alloc::vec::Vec<
+        get_account_proofs_request::AccountRequest,
+    >,
+    /// Optional flag to include header and account code in the response. False by default.
     #[prost(bool, optional, tag = "2")]
     pub include_headers: ::core::option::Option<bool>,
     /// Account code commitments corresponding to the last-known `AccountCode` for requested
@@ -163,3 +166,27 @@ pub struct GetAccountProofsRequest {
     #[prost(message, repeated, tag = "3")]
     pub code_commitments: ::prost::alloc::vec::Vec<super::digest::Digest>,
 }
+/// Nested message and enum types in `GetAccountProofsRequest`.
+pub mod get_account_proofs_request {
+    /// Represents per-account requests where each account ID has its own list of
+    /// (storage_slot_index, map_keys) pairs.
+    #[derive(Clone, PartialEq, ::prost::Message)]
+    pub struct AccountRequest {
+        /// The account ID for this request.
+        #[prost(message, optional, tag = "1")]
+        pub account_id: ::core::option::Option<super::super::account::AccountId>,
+        /// List of storage requests for this account.
+        #[prost(message, repeated, tag = "2")]
+        pub storage_requests: ::prost::alloc::vec::Vec<StorageRequest>,
+    }
+    /// Represents a storage slot index and the associated map keys.
+    #[derive(Clone, PartialEq, ::prost::Message)]
+    pub struct StorageRequest {
+        /// Storage slot index (\[0..255\])
+        #[prost(uint32, tag = "1")]
+        pub storage_slot_index: u32,
+        /// A list of map keys (Digests) associated with this storage slot.
+        #[prost(message, repeated, tag = "2")]
+        pub map_keys: ::prost::alloc::vec::Vec<super::super::digest::Digest>,
+    }
+}