From 1a99dc6e5828628409e79b958cff5be779c36f94 Mon Sep 17 00:00:00 2001 From: KOVACS Krisztian Date: Tue, 4 Aug 2026 20:26:34 +0200 Subject: [PATCH 1/2] refactor(rpc/get_account): batch storage map proofs GetAccount now returns a PartialSmt with all requested key-value pairs. - A request for explicit map keys always returns one partial SMT covering those keys. - Partial SMTs are scoped to one storage map at one block. Proofs from different maps or roots are never merged. - The response carries the original, unhashed `StorageMapKey`s. The SMT contains hashed keys, so the raw keys cannot be recovered from the tree. - Values are not duplicated outside the SMT. A client obtains a value by hashing the raw key and calling `PartialSmt::get_value()`. - `too_many_entries`, `all_entries`, and `partial_map` become mutually exclusive results. - The removed `entries_with_proofs` field number and name are reserved. Even though compatibility is intentionally broken, reusing the field could make an old client try to decode the new length-delimited message as the old one. - The compact protobuf representation mirrors `miden_crypto::merkle::smt::UniqueNodes`. Do not encode `PartialSmt` as an opaque byte string. --- bin/ntx-builder/src/clients/rpc.rs | 18 +- bin/stress-test/src/store/mod.rs | 8 +- crates/proto/src/domain/account.rs | 228 +++++++---- crates/proto/src/domain/account/tests.rs | 127 ++++++ crates/proto/src/domain/merkle.rs | 369 +++++++++++++++++- crates/store/src/account_state_forest/mod.rs | 10 +- .../store/src/account_state_forest/tests.rs | 13 +- crates/store/src/state/account.rs | 2 +- proto/proto/rpc.proto | 38 +- proto/proto/types/primitives.proto | 51 +++ 10 files changed, 733 insertions(+), 131 deletions(-) diff --git a/bin/ntx-builder/src/clients/rpc.rs b/bin/ntx-builder/src/clients/rpc.rs index 44cd4605bc..f31d79de16 100644 --- a/bin/ntx-builder/src/clients/rpc.rs +++ b/bin/ntx-builder/src/clients/rpc.rs @@ -498,16 +498,22 @@ impl RpcClient { )) })?; - let StorageMapEntries::EntriesWithProofs(proofs) = &map_details.entries else { + let StorageMapEntries::PartialMap { map_keys, partial_smt } = &map_details.entries else { return Err(RpcError::InvalidResponse( - "response did not include storage map entry proofs".into(), + "response did not include a partial storage map".into(), )); }; - let proof = proofs.first().cloned().ok_or_else(|| { - RpcError::InvalidResponse( - "response did not include a proof for the requested key".into(), - ) + if !map_keys.contains(&map_key) { + return Err(RpcError::InvalidResponse( + "response partial storage map did not include the requested key".into(), + )); + } + + let proof = partial_smt.open(&map_key.hash().as_word()).map_err(|err| { + RpcError::InvalidResponse(format!( + "response did not track the requested storage map key: {err}" + )) })?; StorageMapWitness::new(proof, [map_key]) diff --git a/bin/stress-test/src/store/mod.rs b/bin/stress-test/src/store/mod.rs index 6f9440148d..854af06a8f 100644 --- a/bin/stress-test/src/store/mod.rs +++ b/bin/stress-test/src/store/mod.rs @@ -121,7 +121,7 @@ async fn get_account( account_id: AccountId, storage_map_slot: String, ) -> GetAccountRun { - use proto::rpc::account_storage_details::account_storage_map_details::Entries; + use proto::rpc::account_storage_details::account_storage_map_details::Result; let request = get_account_request(account_id, storage_map_slot); @@ -136,9 +136,9 @@ async fn get_account( .and_then(|details| details.storage_details.as_ref()) .and_then(|storage_details| storage_details.map_details.first()); let (storage_map_entries, storage_map_limit_exceeded) = match map_details { - Some(details) if details.too_many_entries => (0, true), - Some(details) => match &details.entries { - Some(Entries::AllEntries(entries)) => (entries.entries.len(), false), + Some(details) => match &details.result { + Some(Result::TooManyEntries(true)) => (0, true), + Some(Result::AllEntries(entries)) => (entries.entries.len(), false), _ => (0, false), }, None => (0, false), diff --git a/crates/proto/src/domain/account.rs b/crates/proto/src/domain/account.rs index e361d7a103..31f2faf99c 100644 --- a/crates/proto/src/domain/account.rs +++ b/crates/proto/src/domain/account.rs @@ -16,8 +16,8 @@ use miden_protocol::account::{ use miden_protocol::asset::Asset; use miden_protocol::block::BlockNumber; use miden_protocol::block::account_tree::AccountWitness; -use miden_protocol::crypto::merkle::SparseMerklePath; -use miden_protocol::crypto::merkle::smt::SmtProof; +use miden_protocol::crypto::merkle::smt::{PartialSmt, SmtProof}; +use miden_protocol::crypto::merkle::{MerkleError, SparseMerklePath}; use miden_protocol::utils::serde::{Deserializable, DeserializationError, Serializable}; use super::try_convert; @@ -276,12 +276,21 @@ impl }, ProtoSlotData::MapKeys(keys) => { let keys = try_convert(keys.map_keys).collect::, _>>()?; + if has_duplicate_storage_map_keys(&keys) { + return Err(ConversionError::message( + "storage map key request contains duplicate keys", + )); + } SlotData::MapKeys(keys) }, }) } } +fn has_duplicate_storage_map_keys(keys: &[StorageMapKey]) -> bool { + keys.iter().enumerate().any(|(index, key)| keys[..index].contains(key)) +} + // ACCOUNT HEADER CONVERSIONS //================================================================================================ @@ -440,9 +449,12 @@ pub enum StorageMapEntries { /// requested for small maps. AllEntries(Vec<(StorageMapKey, Word)>), - /// Specific entries with their SMT proofs for client-side verification. Used when specific keys - /// are requested from the storage map. - EntriesWithProofs(Vec), + /// Specific raw map keys covered by a single partial SMT. Used when specific keys are requested + /// from the storage map. + PartialMap { + map_keys: Vec, + partial_smt: PartialSmt, + }, } impl AccountStorageMapDetails { @@ -501,18 +513,53 @@ impl AccountStorageMapDetails { /// /// Use this when the caller has already obtained the proofs from an `SmtForest`. /// Returns `LimitExceeded` if too many proofs are provided. - pub fn from_proofs(slot_name: StorageSlotName, proofs: Vec) -> Self { - if proofs.len() > Self::MAX_SMT_PROOF_ENTRIES { - Self { + pub fn from_proofs( + slot_name: StorageSlotName, + map_root: Word, + map_keys: Vec, + proofs: Vec, + ) -> Result { + if map_keys.len() != proofs.len() { + return Err(MerkleError::InternalError(format!( + "storage map key count {} does not match proof count {}", + map_keys.len(), + proofs.len() + ))); + } + if has_duplicate_storage_map_keys(&map_keys) { + return Err(MerkleError::InternalError( + "storage map key list contains duplicate keys".into(), + )); + } + + if map_keys.len() > Self::MAX_SMT_PROOF_ENTRIES { + return Ok(Self { slot_name, entries: StorageMapEntries::LimitExceeded, - } + }); + } + + let partial_smt = if proofs.is_empty() { + PartialSmt::new(map_root) } else { - Self { - slot_name, - entries: StorageMapEntries::EntriesWithProofs(proofs), - } + PartialSmt::from_proofs(proofs)? + }; + + if partial_smt.root() != map_root { + return Err(MerkleError::ConflictingRoots { + expected_root: map_root, + actual_root: partial_smt.root(), + }); + } + + for map_key in &map_keys { + partial_smt.get_value(&map_key.hash().as_word())?; } + + Ok(Self { + slot_name, + entries: StorageMapEntries::PartialMap { map_keys, partial_smt }, + }) } /// Creates storage map details indicating the limit was exceeded. @@ -534,48 +581,59 @@ impl TryFrom ) -> Result { use proto::rpc::account_storage_details::account_storage_map_details::{ AllMapEntries, - Entries as ProtoEntries, - MapEntriesWithProofs, + PartialStorageMap, + Result as ProtoResult, }; let decoder = value.decoder(); - let proto::rpc::account_storage_details::AccountStorageMapDetails { - slot_name, - too_many_entries, - entries, - } = value; + let proto::rpc::account_storage_details::AccountStorageMapDetails { slot_name, result } = + value; let slot_name = StorageSlotName::new(slot_name).context("slot_name")?; - let entries = if too_many_entries { - StorageMapEntries::LimitExceeded - } else { - match decode!(decoder, entries)? { - ProtoEntries::AllEntries(AllMapEntries { entries }) => { - let entries = entries - .into_iter() - .map(|entry| { - let decoder = entry.decoder(); - let key = StorageMapKey::new(decode!(decoder, entry.key)?); - let value = decode!(decoder, entry.value)?; - Ok((key, value)) - }) - .collect::, ConversionError>>() - .context("entries")?; - StorageMapEntries::AllEntries(entries) - }, - ProtoEntries::EntriesWithProofs(MapEntriesWithProofs { entries }) => { - let proofs = entries - .into_iter() - .map(|entry| { - let decoder = entry.decoder(); - decode!(decoder, entry.proof) - }) - .collect::, ConversionError>>() - .context("entries")?; - StorageMapEntries::EntriesWithProofs(proofs) - }, - } + let entries = match decode!(decoder, result)? { + ProtoResult::TooManyEntries(true) => StorageMapEntries::LimitExceeded, + ProtoResult::TooManyEntries(false) => { + return Err(ConversionError::message("too_many_entries must be true when set")); + }, + ProtoResult::AllEntries(AllMapEntries { entries }) => { + let entries = entries + .into_iter() + .map(|entry| { + let decoder = entry.decoder(); + let key = StorageMapKey::new(decode!(decoder, entry.key)?); + let value = decode!(decoder, entry.value)?; + Ok((key, value)) + }) + .collect::, ConversionError>>() + .context("entries")?; + StorageMapEntries::AllEntries(entries) + }, + ProtoResult::PartialMap(PartialStorageMap { map_keys, partial_smt }) => { + if map_keys.len() > Self::MAX_SMT_PROOF_ENTRIES { + return Err(ConversionError::message(format!( + "partial storage map contains {} keys, exceeding the limit of {}", + map_keys.len(), + Self::MAX_SMT_PROOF_ENTRIES + ))); + } + let map_keys = map_keys + .into_iter() + .map(|key| Word::try_from(key).map(StorageMapKey::new)) + .collect::, _>>() + .context("map_keys")?; + if has_duplicate_storage_map_keys(&map_keys) { + return Err(ConversionError::message( + "partial storage map contains duplicate keys", + )); + } + let partial_smt: PartialSmt = + decode!(decoder, partial_smt).context("partial_smt")?; + for map_key in &map_keys { + partial_smt.get_value(&map_key.hash().as_word()).context("map_keys")?; + } + StorageMapEntries::PartialMap { map_keys, partial_smt } + }, }; Ok(Self { slot_name, entries }) @@ -588,14 +646,14 @@ impl From fn from(value: AccountStorageMapDetails) -> Self { use proto::rpc::account_storage_details::account_storage_map_details::{ AllMapEntries, - Entries as ProtoEntries, - MapEntriesWithProofs, + PartialStorageMap, + Result as ProtoResult, }; let AccountStorageMapDetails { slot_name, entries } = value; - let (too_many_entries, proto_entries) = match entries { - StorageMapEntries::LimitExceeded => (true, None), + let result = match entries { + StorageMapEntries::LimitExceeded => ProtoResult::TooManyEntries(true), StorageMapEntries::AllEntries(entries) => { let all = AllMapEntries { entries: Vec::from_iter(entries.into_iter().map(|(key, value)| { @@ -605,40 +663,19 @@ impl From } })), }; - (false, Some(ProtoEntries::AllEntries(all))) + ProtoResult::AllEntries(all) }, - StorageMapEntries::EntriesWithProofs(proofs) => { - use miden_protocol::crypto::merkle::smt::SmtLeaf; - - let with_proofs = MapEntriesWithProofs { - entries: Vec::from_iter(proofs.into_iter().map(|proof| { - // Get key/value from the leaf before consuming the proof - let (key, value) = match proof.leaf() { - SmtLeaf::Empty(_) => { - (miden_protocol::EMPTY_WORD, miden_protocol::EMPTY_WORD) - }, - SmtLeaf::Single((k, v)) => (*k, *v), - SmtLeaf::Multiple(entries) => entries.iter().next().map_or( - (miden_protocol::EMPTY_WORD, miden_protocol::EMPTY_WORD), - |(k, v)| (*k, *v), - ), - }; - let smt_opening = proto::primitives::SmtOpening::from(proof); - proto::rpc::account_storage_details::account_storage_map_details::map_entries_with_proofs::StorageMapEntryWithProof { - key: Some(key.into()), - value: Some(value.into()), - proof: Some(smt_opening), - } - })), - }; - (false, Some(ProtoEntries::EntriesWithProofs(with_proofs))) + StorageMapEntries::PartialMap { map_keys, partial_smt } => { + ProtoResult::PartialMap(PartialStorageMap { + map_keys: map_keys.into_iter().map(Into::into).collect(), + partial_smt: Some(partial_smt.into()), + }) }, }; Self { slot_name: slot_name.to_string(), - too_many_entries, - entries: proto_entries, + result: Some(result), } } } @@ -671,11 +708,36 @@ impl TryFrom for AccountStorageDetails { let decoder = value.decoder(); let proto::rpc::AccountStorageDetails { header, map_details } = value; - let header = decode!(decoder, header)?; + let header: AccountStorageHeader = decode!(decoder, header)?; - let map_details = + let map_details: Vec = try_convert(map_details).collect::, _>>().context("map_details")?; + for map_detail in &map_details { + let StorageMapEntries::PartialMap { partial_smt, .. } = &map_detail.entries else { + continue; + }; + + let slot = header.find_slot_header_by_name(&map_detail.slot_name).ok_or_else(|| { + ConversionError::message(format!( + "partial storage map references unknown slot {}", + map_detail.slot_name + )) + })?; + if slot.slot_type() != StorageSlotType::Map { + return Err(ConversionError::message(format!( + "partial storage map references non-map slot {}", + map_detail.slot_name + ))); + } + if partial_smt.root() != slot.value() { + return Err(ConversionError::message(format!( + "partial storage map root for slot {} does not match storage header", + map_detail.slot_name + ))); + } + } + Ok(Self { header, map_details }) } } diff --git a/crates/proto/src/domain/account/tests.rs b/crates/proto/src/domain/account/tests.rs index ff431543ea..94702e3efa 100644 --- a/crates/proto/src/domain/account/tests.rs +++ b/crates/proto/src/domain/account/tests.rs @@ -45,6 +45,104 @@ fn account_storage_map_details_from_forest_entries_limit_exceeded() { assert_eq!(details.entries, StorageMapEntries::LimitExceeded); } +#[test] +fn account_storage_map_details_partial_map_round_trip() { + let slot_name = test_slot_name(); + let key0 = StorageMapKey::from_index(1); + let key1 = StorageMapKey::from_index(2); + let value0 = word_from_u32([1, 2, 3, 4]); + let storage_map = StorageMap::with_entries([(key0, value0)]).unwrap(); + let proofs = vec![storage_map.open(&key0).into(), storage_map.open(&key1).into()]; + + let details = AccountStorageMapDetails::from_proofs( + slot_name, + storage_map.root(), + vec![key0, key1], + proofs, + ) + .unwrap(); + let encoded: crate::generated::rpc::account_storage_details::AccountStorageMapDetails = + details.clone().into(); + let decoded = AccountStorageMapDetails::try_from(encoded).unwrap(); + + assert_eq!(decoded, details); + assert_matches::assert_matches!( + decoded.entries, + StorageMapEntries::PartialMap { map_keys, partial_smt } => { + assert_eq!(map_keys, vec![key0, key1]); + assert_eq!(partial_smt.get_value(&key0.hash().as_word()).unwrap(), value0); + assert_eq!(partial_smt.get_value(&key1.hash().as_word()).unwrap(), Word::empty()); + } + ); +} + +#[test] +fn account_storage_map_details_rejects_duplicate_partial_map_keys() { + let slot_name = test_slot_name(); + let key = StorageMapKey::from_index(1); + let storage_map = StorageMap::with_entries([(key, word_from_u32([1, 2, 3, 4]))]).unwrap(); + let proof: SmtProof = storage_map.open(&key).into(); + + let err = AccountStorageMapDetails::from_proofs( + slot_name, + storage_map.root(), + vec![key, key], + vec![proof.clone(), proof], + ) + .unwrap_err(); + + assert!(err.to_string().contains("duplicate keys")); +} + +#[test] +fn account_storage_map_details_rejects_missing_result() { + let encoded = crate::generated::rpc::account_storage_details::AccountStorageMapDetails { + slot_name: test_slot_name().to_string(), + result: None, + }; + + let err = AccountStorageMapDetails::try_from(encoded).unwrap_err(); + assert!(err.to_string().contains("result")); +} + +#[test] +fn account_storage_map_details_rejects_false_limit_marker() { + use crate::generated::rpc::account_storage_details::account_storage_map_details::Result; + + let encoded = crate::generated::rpc::account_storage_details::AccountStorageMapDetails { + slot_name: test_slot_name().to_string(), + result: Some(Result::TooManyEntries(false)), + }; + + let err = AccountStorageMapDetails::try_from(encoded).unwrap_err(); + assert!(err.to_string().contains("must be true")); +} + +#[test] +fn account_storage_details_rejects_partial_map_root_mismatch() { + let slot_name = test_slot_name(); + let key = StorageMapKey::from_index(1); + let storage_map = StorageMap::with_entries([(key, word_from_u32([1, 2, 3, 4]))]).unwrap(); + let map_details = AccountStorageMapDetails::from_proofs( + slot_name.clone(), + storage_map.root(), + vec![key], + vec![storage_map.open(&key).into()], + ) + .unwrap(); + let header = AccountStorageHeader::new(vec![StorageSlotHeader::new( + slot_name, + StorageSlotType::Map, + Word::empty(), + )]) + .unwrap(); + let encoded: crate::generated::rpc::AccountStorageDetails = + AccountStorageDetails { header, map_details: vec![map_details] }.into(); + + let err = AccountStorageDetails::try_from(encoded).unwrap_err(); + assert!(err.to_string().contains("does not match storage header")); +} + #[test] fn account_detail_request_converts_all_storage_maps() { use crate::generated::rpc::account_request::account_detail_request::StorageRequest; @@ -103,6 +201,35 @@ fn account_detail_request_converts_explicit_storage_maps() { )); } +#[test] +fn account_detail_request_rejects_duplicate_storage_map_keys() { + use crate::generated::rpc::account_request::account_detail_request::{ + StorageMapDetailRequest, + StorageMapDetailRequests, + StorageRequest, + storage_map_detail_request, + }; + use crate::generated::rpc::account_request::account_detail_request::storage_map_detail_request::MapKeys; + + let map_key: crate::generated::primitives::Digest = Word::from([1, 2, 3, 4u32]).into(); + let request = crate::generated::rpc::account_request::AccountDetailRequest { + code_commitment: None, + asset_vault_commitment: None, + storage_request: Some(StorageRequest::StorageMaps(StorageMapDetailRequests { + storage_maps: vec![StorageMapDetailRequest { + slot_name: "miden::test::storage::slot".to_string(), + slot_data: Some(storage_map_detail_request::SlotData::MapKeys(MapKeys { + map_keys: vec![map_key, map_key], + })), + }], + })), + }; + + let err = AccountDetailRequest::try_from(request).unwrap_err(); + + assert!(err.to_string().contains("duplicate keys")); +} + #[test] fn account_detail_request_allows_no_storage_slot_data() { let request = crate::generated::rpc::account_request::AccountDetailRequest { diff --git a/crates/proto/src/domain/merkle.rs b/crates/proto/src/domain/merkle.rs index 32036bddd1..0c3916ca81 100644 --- a/crates/proto/src/domain/merkle.rs +++ b/crates/proto/src/domain/merkle.rs @@ -1,7 +1,11 @@ +use std::collections::BTreeSet; + use miden_protocol::Word; use miden_protocol::crypto::merkle::mmr::{Forest, MmrDelta}; -use miden_protocol::crypto::merkle::smt::{LeafIndex, SmtLeaf, SmtProof}; -use miden_protocol::crypto::merkle::{MerklePath, SparseMerklePath}; +use miden_protocol::crypto::merkle::smt::{ + LeafIndex, NodeValue, PartialSmt, SMT_DEPTH, SmtLeaf, SmtProof, UniqueNodes, +}; +use miden_protocol::crypto::merkle::{MerklePath, NodeIndex, SparseMerklePath}; use crate::decode::{ConversionResultExt, GrpcDecodeExt}; use crate::domain::{convert, try_convert}; @@ -13,7 +17,11 @@ use crate::{decode, generated as proto}; impl From<&MerklePath> for proto::primitives::MerklePath { fn from(value: &MerklePath) -> Self { - let siblings = value.nodes().iter().map(proto::primitives::Digest::from).collect(); + let siblings = value + .nodes() + .iter() + .map(proto::primitives::Digest::from) + .collect(); proto::primitives::MerklePath { siblings } } } @@ -48,7 +56,10 @@ impl From for proto::primitives::SparseMerklePath { let (empty_nodes_mask, siblings) = value.into_parts(); proto::primitives::SparseMerklePath { empty_nodes_mask, - siblings: siblings.into_iter().map(proto::primitives::Digest::from).collect(), + siblings: siblings + .into_iter() + .map(proto::primitives::Digest::from) + .collect(), } } } @@ -74,7 +85,11 @@ impl TryFrom for SparseMerklePath { impl From for proto::primitives::MmrDelta { fn from(value: MmrDelta) -> Self { - let data = value.data.into_iter().map(proto::primitives::Digest::from).collect(); + let data = value + .data + .into_iter() + .map(proto::primitives::Digest::from) + .collect(); proto::primitives::MmrDelta { forest: value.forest.num_leaves() as u64, data, @@ -93,8 +108,10 @@ impl TryFrom for MmrDelta { .collect::>() .context("data")?; - let forest_size: usize = - value.forest.try_into().context("forest size does not fit in usize")?; + let forest_size: usize = value + .forest + .try_into() + .context("forest size does not fit in usize")?; let forest = Forest::new(forest_size).context("forest size out of range")?; Ok(MmrDelta { forest, data }) @@ -117,18 +134,19 @@ impl TryFrom for SmtLeaf { match leaf { proto::primitives::smt_leaf::Leaf::EmptyLeafIndex(leaf_index) => { Ok(Self::new_empty(LeafIndex::new_max_depth(leaf_index))) - }, + } proto::primitives::smt_leaf::Leaf::Single(entry) => { let (key, value): (Word, Word) = entry.try_into().context("entry")?; Ok(SmtLeaf::new_single(key, value)) - }, + } proto::primitives::smt_leaf::Leaf::Multiple(entries) => { - let domain_entries: Vec<(Word, Word)> = - try_convert(entries.entries).collect::>().context("entries")?; + let domain_entries: Vec<(Word, Word)> = try_convert(entries.entries) + .collect::>() + .context("entries")?; Ok(SmtLeaf::new_multiple(domain_entries)?) - }, + } } } } @@ -197,3 +215,330 @@ impl From for proto::primitives::SmtOpening { } } } + +// PARTIAL SMT +// ------------------------------------------------------------------------------------------------ + +impl From for proto::primitives::PartialSmt { + fn from(unique_nodes: UniqueNodes) -> Self { + use proto::primitives::partial_smt_node::Value; + + let UniqueNodes { + root, + nodes, + leaves, + value_only_leaves, + } = unique_nodes; + + let mut node_levels = nodes.into_iter().collect::>(); + node_levels.sort_by_key(|(depth, _)| *depth); + let node_levels = node_levels + .into_iter() + .map(|(depth, nodes)| { + let nodes = nodes + .into_iter() + .map(|(index, value)| { + let value = match value { + NodeValue::EmptySubtreeRoot => Value::EmptySubtreeRoot(true), + NodeValue::Present(value) => Value::Digest(value.into()), + }; + proto::primitives::PartialSmtNode { + index, + value: Some(value), + } + }) + .collect(); + + proto::primitives::PartialSmtNodeLevel { + depth: u32::from(depth), + nodes, + } + }) + .collect(); + + let leaves = leaves + .into_iter() + .map(|(index, leaf)| proto::primitives::IndexedSmtLeaf { + index, + leaf: Some(leaf.into()), + }) + .collect(); + + let value_only_leaves = value_only_leaves + .into_iter() + .map(|(index, value)| proto::primitives::IndexedDigest { + index, + value: Some(value.into()), + }) + .collect(); + + Self { + root: Some(root.into()), + node_levels, + leaves, + value_only_leaves, + } + } +} + +impl TryFrom for UniqueNodes { + type Error = ConversionError; + + fn try_from(value: proto::primitives::PartialSmt) -> Result { + use proto::primitives::partial_smt_node::Value; + + let decoder = value.decoder(); + let proto::primitives::PartialSmt { + root, + node_levels, + leaves, + value_only_leaves, + } = value; + + let root = decode!(decoder, root)?; + + let mut seen_depths = BTreeSet::new(); + let mut decoded_levels = Vec::with_capacity(node_levels.len()); + for level in node_levels { + let depth = u8::try_from(level.depth).context("node_levels.depth")?; + if depth == 0 || depth >= SMT_DEPTH { + return Err(ConversionError::message(format!( + "partial SMT node depth {depth} must be in the range 1..{SMT_DEPTH}" + ))); + } + if !seen_depths.insert(depth) { + return Err(ConversionError::message(format!( + "partial SMT contains duplicate node depth {depth}" + ))); + } + + let mut seen_indices = BTreeSet::new(); + let mut decoded_nodes = Vec::with_capacity(level.nodes.len()); + for node in level.nodes { + NodeIndex::new(depth, node.index).context("node_levels.nodes.index")?; + if !seen_indices.insert(node.index) { + return Err(ConversionError::message(format!( + "partial SMT contains duplicate node index {} at depth {depth}", + node.index + ))); + } + + let node_value = match node.value.ok_or_else(|| { + ConversionError::missing_field::("value") + })? { + Value::Digest(value) => NodeValue::Present(value.try_into().context("digest")?), + Value::EmptySubtreeRoot(true) => NodeValue::EmptySubtreeRoot, + Value::EmptySubtreeRoot(false) => { + return Err(ConversionError::message( + "partial SMT empty_subtree_root marker must be true", + )); + } + }; + decoded_nodes.push((node.index, node_value)); + } + decoded_levels.push((depth, decoded_nodes)); + } + + let mut seen_leaf_indices = BTreeSet::new(); + let mut decoded_leaves = Vec::with_capacity(leaves.len()); + for indexed_leaf in leaves { + if !seen_leaf_indices.insert(indexed_leaf.index) { + return Err(ConversionError::message(format!( + "partial SMT contains duplicate leaf index {}", + indexed_leaf.index + ))); + } + let decoder = indexed_leaf.decoder(); + let leaf = decode!(decoder, indexed_leaf.leaf)?; + decoded_leaves.push((indexed_leaf.index, leaf)); + } + + let mut seen_value_only_indices = BTreeSet::new(); + let mut decoded_value_only_leaves = Vec::with_capacity(value_only_leaves.len()); + for indexed_digest in value_only_leaves { + if !seen_value_only_indices.insert(indexed_digest.index) { + return Err(ConversionError::message(format!( + "partial SMT contains duplicate value-only leaf index {}", + indexed_digest.index + ))); + } + if seen_leaf_indices.contains(&indexed_digest.index) { + return Err(ConversionError::message(format!( + "partial SMT leaf index {} has both a leaf and a value-only leaf", + indexed_digest.index + ))); + } + let decoder = indexed_digest.decoder(); + let digest = decode!(decoder, indexed_digest.value)?; + decoded_value_only_leaves.push((indexed_digest.index, digest)); + } + + Ok(UniqueNodes { + root, + nodes: decoded_levels.into_iter().collect(), + leaves: decoded_leaves, + value_only_leaves: decoded_value_only_leaves, + }) + } +} + +impl From for proto::primitives::PartialSmt { + fn from(partial_smt: PartialSmt) -> Self { + let mut unique_nodes = partial_smt.to_unique_nodes(); + + // Workaround for a PartialSmt issue: https://github.com/0xMiden/miden-vm/issues/3470. + if PartialSmt::from_unique_nodes(unique_nodes.clone()).is_err() { + for (index, node) in partial_smt.inner_node_indices() { + if index.depth() == 0 { + continue; + } + + let level = unique_nodes.nodes.entry(index.depth()).or_default(); + if !level + .iter() + .any(|(position, _)| *position == index.position()) + { + level.push((index.position(), NodeValue::Present(node.hash()))); + } + } + + debug_assert!(PartialSmt::from_unique_nodes(unique_nodes.clone()).is_ok()); + } + + unique_nodes.into() + } +} + +impl TryFrom for PartialSmt { + type Error = ConversionError; + + fn try_from(value: proto::primitives::PartialSmt) -> Result { + let unique_nodes = UniqueNodes::try_from(value)?; + PartialSmt::from_unique_nodes(unique_nodes) + .map_err(|err| ConversionError::deserialization("PartialSmt", err)) + } +} + +#[cfg(test)] +mod tests { + use miden_protocol::crypto::merkle::smt::Smt; + + use super::*; + + // Test if the PartialSmt::from_unique_nodes() issue is still there: + // https://github.com/0xMiden/miden-vm/issues/3470 + // + // We can remove our workaround from the conversion above once that has been + // fixed. + #[test] + fn unique_nodes_cannot_reconstruct_mixed_inclusion_and_exclusion_branches() { + let included_key = Word::from([1, 2, 3, 4u32]); + let other_key = Word::from([5, 6, 7, 8u32]); + let missing_key = Word::from([9, 10, 11, 12u32]); + let included_value = Word::from([13, 14, 15, 16u32]); + let other_value = Word::from([17, 18, 19, 20u32]); + let smt = + Smt::with_entries([(included_key, included_value), (other_key, other_value)]).unwrap(); + let partial_smt = + PartialSmt::from_proofs([smt.open(&included_key), smt.open(&missing_key)]).unwrap(); + + // The partial tree itself is valid and tracks both the inclusion and exclusion. + assert_eq!( + partial_smt.get_value(&included_key).unwrap(), + included_value + ); + assert_eq!(partial_smt.get_value(&missing_key).unwrap(), Word::empty()); + + let err = PartialSmt::from_unique_nodes(partial_smt.to_unique_nodes()).expect_err( + "mixed inclusion and exclusion branches should expose the reconstruction bug", + ); + + assert!(err.to_string().contains("not found but is required")); + } + + #[test] + fn partial_smt_round_trip() { + let key0 = Word::from([1, 2, 3, 4u32]); + let key1 = Word::from([5, 6, 7, 8u32]); + let missing_key = Word::from([9, 10, 11, 12u32]); + let value0 = Word::from([13, 14, 15, 16u32]); + let value1 = Word::from([17, 18, 19, 20u32]); + let smt = Smt::with_entries([(key0, value0), (key1, value1)]).unwrap(); + let partial_smt = + PartialSmt::from_proofs([smt.open(&key0), smt.open(&missing_key)]).unwrap(); + + let encoded: proto::primitives::PartialSmt = partial_smt.clone().into(); + assert!(encoded.node_levels.is_sorted_by_key(|level| level.depth)); + + let decoded_unique_nodes = UniqueNodes::try_from(encoded).unwrap(); + let decoded = PartialSmt::from_unique_nodes(decoded_unique_nodes).unwrap(); + + assert_eq!(decoded, partial_smt); + assert_eq!(decoded.get_value(&key0).unwrap(), value0); + assert_eq!(decoded.get_value(&missing_key).unwrap(), Word::empty()); + } + + #[test] + fn partial_smt_rejects_false_empty_subtree_marker() { + use proto::primitives::partial_smt_node::Value; + + let encoded = proto::primitives::PartialSmt { + root: Some(PartialSmt::EMPTY_ROOT.into()), + node_levels: vec![proto::primitives::PartialSmtNodeLevel { + depth: 1, + nodes: vec![proto::primitives::PartialSmtNode { + index: 0, + value: Some(Value::EmptySubtreeRoot(false)), + }], + }], + leaves: vec![], + value_only_leaves: vec![], + }; + + let err = UniqueNodes::try_from(encoded).unwrap_err(); + assert!(err.to_string().contains("must be true")); + } + + #[test] + fn partial_smt_rejects_duplicate_depths() { + let encoded = proto::primitives::PartialSmt { + root: Some(PartialSmt::EMPTY_ROOT.into()), + node_levels: vec![ + proto::primitives::PartialSmtNodeLevel { + depth: 1, + nodes: vec![], + }, + proto::primitives::PartialSmtNodeLevel { + depth: 1, + nodes: vec![], + }, + ], + leaves: vec![], + value_only_leaves: vec![], + }; + + let err = UniqueNodes::try_from(encoded).unwrap_err(); + assert!(err.to_string().contains("duplicate node depth")); + } + + #[test] + fn partial_smt_rejects_invalid_node_index() { + use proto::primitives::partial_smt_node::Value; + + let encoded = proto::primitives::PartialSmt { + root: Some(PartialSmt::EMPTY_ROOT.into()), + node_levels: vec![proto::primitives::PartialSmtNodeLevel { + depth: 1, + nodes: vec![proto::primitives::PartialSmtNode { + index: 2, + value: Some(Value::EmptySubtreeRoot(true)), + }], + }], + leaves: vec![], + value_only_leaves: vec![], + }; + + let err = UniqueNodes::try_from(encoded).unwrap_err(); + assert!(err.to_string().contains("not valid for depth")); + } +} diff --git a/crates/store/src/account_state_forest/mod.rs b/crates/store/src/account_state_forest/mod.rs index 0f7c2775ed..512ba2d295 100644 --- a/crates/store/src/account_state_forest/mod.rs +++ b/crates/store/src/account_state_forest/mod.rs @@ -615,7 +615,7 @@ impl AccountStateForest { Ok(Some(AccountVaultDetails::from_assets(assets))) } - /// Opens a storage map and returns storage map details with SMT proofs for the given keys. + /// Opens a storage map and returns one partial SMT covering the requested keys. /// /// Returns `None` if no storage root is tracked for this account/slot/block combination. /// Returns a `MerkleError` if the forest doesn't contain sufficient data for the proofs. @@ -631,13 +631,19 @@ impl AccountStateForest { ) -> Option> { let lineage = Self::storage_lineage_id(account_id, &slot_name); let tree = self.get_tree_id(lineage, block_num)?; + let map_root = match self.forest.root_info(tree) { + RootInfo::LatestVersion(root) | RootInfo::HistoricalVersion(root) => root, + RootInfo::Missing => return None, + }; let proofs = Result::from_iter(raw_keys.iter().map(|raw_key| { let key_hashed = raw_key.hash().into(); self.forest.open(tree, key_hashed).map_err(Self::map_forest_error) })); - Some(proofs.map(|proofs| AccountStorageMapDetails::from_proofs(slot_name, proofs))) + Some(proofs.and_then(|proofs| { + AccountStorageMapDetails::from_proofs(slot_name, map_root, raw_keys.to_vec(), proofs) + })) } /// Enumerates a storage map as it is stored in the SMT. diff --git a/crates/store/src/account_state_forest/tests.rs b/crates/store/src/account_state_forest/tests.rs index b066b8ab80..90d95178a0 100644 --- a/crates/store/src/account_state_forest/tests.rs +++ b/crates/store/src/account_state_forest/tests.rs @@ -876,7 +876,7 @@ fn storage_map_empty_entries_query() { } #[test] -fn storage_map_open_returns_proofs() { +fn storage_map_open_returns_partial_map() { use std::collections::BTreeMap; use assert_matches::assert_matches; @@ -904,8 +904,15 @@ fn storage_map_open_returns_proofs() { forest.get_storage_map_details_for_keys(account_id, slot_name.clone(), block_num, &keys); let details = result.expect("Should return Some").expect("Should not error"); - assert_matches!(details.entries, StorageMapEntries::EntriesWithProofs(entries) => { - assert_eq!(entries.len(), keys.len()); + assert_matches!(details.entries, StorageMapEntries::PartialMap { map_keys, partial_smt } => { + assert_eq!(map_keys, keys); + for key in &map_keys { + assert!(partial_smt.get_value(&key.hash().as_word()).is_ok()); + } + assert_eq!( + partial_smt.root(), + forest.get_storage_map_root(account_id, &slot_name, block_num).unwrap() + ); }); } diff --git a/crates/store/src/state/account.rs b/crates/store/src/state/account.rs index d8dd4bf40d..7b30c7a9ff 100644 --- a/crates/store/src/state/account.rs +++ b/crates/store/src/state/account.rs @@ -433,7 +433,7 @@ fn estimate_storage_map_details_field_len(details: &AccountStorageMapDetails) -> }, // `apply_all_storage_maps_response_budget()` is only used for `all_storage_maps` requests, // which never request proofs. Be conservative and force the fallback path if this changes. - StorageMapEntries::EntriesWithProofs(_) => usize::MAX, + StorageMapEntries::PartialMap { .. } => usize::MAX, } } diff --git a/proto/proto/rpc.proto b/proto/proto/rpc.proto index 66f9468f45..f78075d041 100644 --- a/proto/proto/rpc.proto +++ b/proto/proto/rpc.proto @@ -360,19 +360,6 @@ message AccountVaultDetails { // Account storage details for AccountResponse message AccountStorageDetails { message AccountStorageMapDetails { - // Wrapper for repeated storage map entries including their proofs. - // Used when specific keys are requested to enable client-side verification. - message MapEntriesWithProofs { - // Definition of individual storage entries including a proof. - message StorageMapEntryWithProof { - primitives.Digest key = 1; - primitives.Digest value = 2; - primitives.SmtOpening proof = 3; - } - - repeated StorageMapEntryWithProof entries = 1; - } - // Wrapper for repeated storage map entries (without proofs). // Used when all entries are requested for small maps. message AllMapEntries { @@ -385,21 +372,32 @@ message AccountStorageDetails { repeated StorageMapEntry entries = 1; } + // A partial storage map containing one compact SMT proof shared by all requested keys. + message PartialStorageMap { + // Original, unhashed storage map keys covered by `partial_smt`. + repeated primitives.Digest map_keys = 1; + + // Compact partial SMT tracking all requested keys. + primitives.PartialSmt partial_smt = 2; + } + // Storage slot name. string slot_name = 1; - // True when the number of entries exceeds the response limit. - // When set, clients should use the `SyncAccountStorageMaps` endpoint. - bool too_many_entries = 2; + // Exactly one storage map result. + oneof result { + // The number of entries exceeds the response limit. This variant is always true. + bool too_many_entries = 2; - // The map entries (with or without proofs). Empty when too_many_entries is true. - oneof entries { // All storage entries without proofs (for small maps or full requests). AllMapEntries all_entries = 3; - // Specific entries with their SMT proofs (for partial requests). - MapEntriesWithProofs entries_with_proofs = 4; + // Specific keys covered by one compact partial SMT. + PartialStorageMap partial_map = 5; } + + reserved 4; + reserved "entries_with_proofs"; } // Account storage header (storage slot info for up to 256 slots) diff --git a/proto/proto/types/primitives.proto b/proto/proto/types/primitives.proto index a0c30b812a..0f0aabfef9 100644 --- a/proto/proto/types/primitives.proto +++ b/proto/proto/types/primitives.proto @@ -52,6 +52,57 @@ message SmtOpening { SmtLeaf leaf = 2; } +// A compact representation of a partial SMT containing the nodes needed to reconstruct the +// tracked openings. +message PartialSmt { + // The root expected after reconstruction. + Digest root = 1; + + // Boundary nodes needed to reconstruct intermediate nodes, grouped by depth. + repeated PartialSmtNodeLevel node_levels = 2; + + // Fully materialized, non-empty leaves. + repeated IndexedSmtLeaf leaves = 3; + + // Leaves for which only the hash is known. + repeated IndexedDigest value_only_leaves = 4; +} + +// Partial SMT boundary nodes at a single depth. +message PartialSmtNodeLevel { + // Depth of all nodes in this group. + uint32 depth = 1; + + // Boundary nodes at this depth. + repeated PartialSmtNode nodes = 2; +} + +// A boundary node needed to reconstruct a partial SMT. +message PartialSmtNode { + // Position of the node within its depth. + fixed64 index = 1; + + oneof value { + // Hash of a non-empty boundary node. + Digest digest = 2; + + // Indicates the canonical empty subtree root for the containing depth. + bool empty_subtree_root = 3; + } +} + +// A fully materialized SMT leaf together with its index. +message IndexedSmtLeaf { + fixed64 index = 1; + SmtLeaf leaf = 2; +} + +// A digest together with its SMT leaf index. +message IndexedDigest { + fixed64 index = 1; + Digest value = 2; +} + // A different representation of a Merkle path designed for memory efficiency. message SparseMerklePath { // A bitmask representing empty nodes. From 946e7e2ea3bfe7abc7e15129f4852b1d60f4b4e0 Mon Sep 17 00:00:00 2001 From: KOVACS Krisztian Date: Wed, 5 Aug 2026 17:33:21 +0200 Subject: [PATCH 2/2] fix: formatting --- crates/proto/src/domain/merkle.rs | 81 ++++++++++--------------------- 1 file changed, 25 insertions(+), 56 deletions(-) diff --git a/crates/proto/src/domain/merkle.rs b/crates/proto/src/domain/merkle.rs index 0c3916ca81..6cfddf09c5 100644 --- a/crates/proto/src/domain/merkle.rs +++ b/crates/proto/src/domain/merkle.rs @@ -3,7 +3,13 @@ use std::collections::BTreeSet; use miden_protocol::Word; use miden_protocol::crypto::merkle::mmr::{Forest, MmrDelta}; use miden_protocol::crypto::merkle::smt::{ - LeafIndex, NodeValue, PartialSmt, SMT_DEPTH, SmtLeaf, SmtProof, UniqueNodes, + LeafIndex, + NodeValue, + PartialSmt, + SMT_DEPTH, + SmtLeaf, + SmtProof, + UniqueNodes, }; use miden_protocol::crypto::merkle::{MerklePath, NodeIndex, SparseMerklePath}; @@ -17,11 +23,7 @@ use crate::{decode, generated as proto}; impl From<&MerklePath> for proto::primitives::MerklePath { fn from(value: &MerklePath) -> Self { - let siblings = value - .nodes() - .iter() - .map(proto::primitives::Digest::from) - .collect(); + let siblings = value.nodes().iter().map(proto::primitives::Digest::from).collect(); proto::primitives::MerklePath { siblings } } } @@ -56,10 +58,7 @@ impl From for proto::primitives::SparseMerklePath { let (empty_nodes_mask, siblings) = value.into_parts(); proto::primitives::SparseMerklePath { empty_nodes_mask, - siblings: siblings - .into_iter() - .map(proto::primitives::Digest::from) - .collect(), + siblings: siblings.into_iter().map(proto::primitives::Digest::from).collect(), } } } @@ -85,11 +84,7 @@ impl TryFrom for SparseMerklePath { impl From for proto::primitives::MmrDelta { fn from(value: MmrDelta) -> Self { - let data = value - .data - .into_iter() - .map(proto::primitives::Digest::from) - .collect(); + let data = value.data.into_iter().map(proto::primitives::Digest::from).collect(); proto::primitives::MmrDelta { forest: value.forest.num_leaves() as u64, data, @@ -108,10 +103,8 @@ impl TryFrom for MmrDelta { .collect::>() .context("data")?; - let forest_size: usize = value - .forest - .try_into() - .context("forest size does not fit in usize")?; + let forest_size: usize = + value.forest.try_into().context("forest size does not fit in usize")?; let forest = Forest::new(forest_size).context("forest size out of range")?; Ok(MmrDelta { forest, data }) @@ -134,19 +127,18 @@ impl TryFrom for SmtLeaf { match leaf { proto::primitives::smt_leaf::Leaf::EmptyLeafIndex(leaf_index) => { Ok(Self::new_empty(LeafIndex::new_max_depth(leaf_index))) - } + }, proto::primitives::smt_leaf::Leaf::Single(entry) => { let (key, value): (Word, Word) = entry.try_into().context("entry")?; Ok(SmtLeaf::new_single(key, value)) - } + }, proto::primitives::smt_leaf::Leaf::Multiple(entries) => { - let domain_entries: Vec<(Word, Word)> = try_convert(entries.entries) - .collect::>() - .context("entries")?; + let domain_entries: Vec<(Word, Word)> = + try_convert(entries.entries).collect::>().context("entries")?; Ok(SmtLeaf::new_multiple(domain_entries)?) - } + }, } } } @@ -223,12 +215,7 @@ impl From for proto::primitives::PartialSmt { fn from(unique_nodes: UniqueNodes) -> Self { use proto::primitives::partial_smt_node::Value; - let UniqueNodes { - root, - nodes, - leaves, - value_only_leaves, - } = unique_nodes; + let UniqueNodes { root, nodes, leaves, value_only_leaves } = unique_nodes; let mut node_levels = nodes.into_iter().collect::>(); node_levels.sort_by_key(|(depth, _)| *depth); @@ -242,17 +229,11 @@ impl From for proto::primitives::PartialSmt { NodeValue::EmptySubtreeRoot => Value::EmptySubtreeRoot(true), NodeValue::Present(value) => Value::Digest(value.into()), }; - proto::primitives::PartialSmtNode { - index, - value: Some(value), - } + proto::primitives::PartialSmtNode { index, value: Some(value) } }) .collect(); - proto::primitives::PartialSmtNodeLevel { - depth: u32::from(depth), - nodes, - } + proto::primitives::PartialSmtNodeLevel { depth: u32::from(depth), nodes } }) .collect(); @@ -332,7 +313,7 @@ impl TryFrom for UniqueNodes { return Err(ConversionError::message( "partial SMT empty_subtree_root marker must be true", )); - } + }, }; decoded_nodes.push((node.index, node_value)); } @@ -394,10 +375,7 @@ impl From for proto::primitives::PartialSmt { } let level = unique_nodes.nodes.entry(index.depth()).or_default(); - if !level - .iter() - .any(|(position, _)| *position == index.position()) - { + if !level.iter().any(|(position, _)| *position == index.position()) { level.push((index.position(), NodeValue::Present(node.hash()))); } } @@ -443,10 +421,7 @@ mod tests { PartialSmt::from_proofs([smt.open(&included_key), smt.open(&missing_key)]).unwrap(); // The partial tree itself is valid and tracks both the inclusion and exclusion. - assert_eq!( - partial_smt.get_value(&included_key).unwrap(), - included_value - ); + assert_eq!(partial_smt.get_value(&included_key).unwrap(), included_value); assert_eq!(partial_smt.get_value(&missing_key).unwrap(), Word::empty()); let err = PartialSmt::from_unique_nodes(partial_smt.to_unique_nodes()).expect_err( @@ -504,14 +479,8 @@ mod tests { let encoded = proto::primitives::PartialSmt { root: Some(PartialSmt::EMPTY_ROOT.into()), node_levels: vec![ - proto::primitives::PartialSmtNodeLevel { - depth: 1, - nodes: vec![], - }, - proto::primitives::PartialSmtNodeLevel { - depth: 1, - nodes: vec![], - }, + proto::primitives::PartialSmtNodeLevel { depth: 1, nodes: vec![] }, + proto::primitives::PartialSmtNodeLevel { depth: 1, nodes: vec![] }, ], leaves: vec![], value_only_leaves: vec![],