chainHead: Add support for storage closest merkle descendant #14818

substrate/client/api/Cargo.toml

@@ -35,6 +35,7 @@ sp-runtime = { path = "../../primitives/runtime", default-features = false}
 sp-state-machine = { path = "../../primitives/state-machine" }
 sp-statement-store = { path = "../../primitives/statement-store" }
 sp-storage = { path = "../../primitives/storage" }
+sp-trie = { path = "../../primitives/trie" }
 
 [dev-dependencies]
 thiserror = "1.0.48"

substrate/client/api/src/backend.rs

@@ -33,6 +33,7 @@ use sp_state_machine::{
 	OffchainChangesCollection, StorageCollection, StorageIterator,
 };
 use sp_storage::{ChildInfo, StorageData, StorageKey};
+pub use sp_trie::MerkleValue;
 
 use crate::{blockchain::Backend as BlockchainBackend, UsageInfo};
 
@@ -470,6 +471,21 @@ pub trait StorageProvider<Block: BlockT, B: Backend<Block>> {
 		child_info: &ChildInfo,
 		key: &StorageKey,
 	) -> sp_blockchain::Result<Option<Block::Hash>>;
+
+	/// Given a block's `Hash` and a key, return the closest merkle value.
+	fn closest_merkle_value(
+		&self,
+		hash: Block::Hash,
+		key: &StorageKey,
+	) -> sp_blockchain::Result<Option<MerkleValue<Block::Hash>>>;
+
+	/// Given a block's `Hash`, a key and a child storage key, return the closest merkle value.
+	fn child_closest_merkle_value(
+		&self,
+		hash: Block::Hash,
+		child_info: &ChildInfo,
+		key: &StorageKey,
+	) -> sp_blockchain::Result<Option<MerkleValue<Block::Hash>>>;
 }
 
 /// Client backend.

substrate/client/db/src/bench.rs

@@ -37,7 +37,7 @@ use sp_state_machine::{
 };
 use sp_trie::{
 	cache::{CacheSize, SharedTrieCache},
-	prefixed_key, MemoryDB,
+	prefixed_key, MemoryDB, MerkleValue,
 };
 use std::{
 	cell::{Cell, RefCell},
@@ -382,6 +382,27 @@ impl<B: BlockT> StateBackend<HashingFor<B>> for BenchmarkingState<B> {
 			.child_storage_hash(child_info, key)
 	}
 
+	fn closest_merkle_value(
+		&self,
+		key: &[u8],
+	) -> Result<Option<MerkleValue<B::Hash>>, Self::Error> {
+		self.add_read_key(None, key);
+		self.state.borrow().as_ref().ok_or_else(state_err)?.closest_merkle_value(key)
+	}
+
+	fn child_closest_merkle_value(
+		&self,
+		child_info: &ChildInfo,
+		key: &[u8],
+	) -> Result<Option<MerkleValue<B::Hash>>, Self::Error> {
+		self.add_read_key(None, key);
+		self.state
+			.borrow()
+			.as_ref()
+			.ok_or_else(state_err)?
+			.child_closest_merkle_value(child_info, key)
+	}
+
 	fn exists_storage(&self, key: &[u8]) -> Result<bool, Self::Error> {
 		self.add_read_key(None, key);
 		self.state.borrow().as_ref().ok_or_else(state_err)?.exists_storage(key)

substrate/client/db/src/lib.rs

@@ -90,7 +90,7 @@ use sp_state_machine::{
 	OffchainChangesCollection, StateMachineStats, StorageCollection, StorageIterator, StorageKey,
 	StorageValue, UsageInfo as StateUsageInfo,
 };
-use sp_trie::{cache::SharedTrieCache, prefixed_key, MemoryDB, PrefixedMemoryDB};
+use sp_trie::{cache::SharedTrieCache, prefixed_key, MemoryDB, MerkleValue, PrefixedMemoryDB};
 
 // Re-export the Database trait so that one can pass an implementation of it.
 pub use sc_state_db::PruningMode;
@@ -214,6 +214,21 @@ impl<B: BlockT> StateBackend<HashingFor<B>> for RefTrackingState<B> {
 		self.state.child_storage_hash(child_info, key)
 	}
 
+	fn closest_merkle_value(
+		&self,
+		key: &[u8],
+	) -> Result<Option<MerkleValue<B::Hash>>, Self::Error> {
+		self.state.closest_merkle_value(key)
+	}
+
+	fn child_closest_merkle_value(
+		&self,
+		child_info: &ChildInfo,
+		key: &[u8],
+	) -> Result<Option<MerkleValue<B::Hash>>, Self::Error> {
+		self.state.child_closest_merkle_value(child_info, key)
+	}
+
 	fn exists_storage(&self, key: &[u8]) -> Result<bool, Self::Error> {
 		self.state.exists_storage(key)
 	}

substrate/client/db/src/record_stats_state.rs

@@ -28,6 +28,7 @@ use sp_state_machine::{
 	backend::{AsTrieBackend, Backend as StateBackend},
 	BackendTransaction, IterArgs, StorageIterator, StorageKey, StorageValue, TrieBackend,
 };
+use sp_trie::MerkleValue;
 use std::sync::Arc;
 
 /// State abstraction for recording stats about state access.
@@ -144,6 +145,21 @@ impl<S: StateBackend<HashingFor<B>>, B: BlockT> StateBackend<HashingFor<B>>
 		self.state.child_storage_hash(child_info, key)
 	}
 
+	fn closest_merkle_value(
+		&self,
+		key: &[u8],
+	) -> Result<Option<MerkleValue<B::Hash>>, Self::Error> {
+		self.state.closest_merkle_value(key)
+	}
+
+	fn child_closest_merkle_value(
+		&self,
+		child_info: &ChildInfo,
+		key: &[u8],
+	) -> Result<Option<MerkleValue<B::Hash>>, Self::Error> {
+		self.state.child_closest_merkle_value(child_info, key)
+	}
+
 	fn exists_storage(&self, key: &[u8]) -> Result<bool, Self::Error> {
 		self.state.exists_storage(key)
 	}

substrate/client/rpc-spec-v2/src/chain_head/chain_head.rs

@@ -27,7 +27,7 @@ use crate::{
 		api::ChainHeadApiServer,
 		chain_head_follow::ChainHeadFollower,
 		error::Error as ChainHeadRpcError,
-		event::{FollowEvent, MethodResponse, OperationError, StorageQuery, StorageQueryType},
+		event::{FollowEvent, MethodResponse, OperationError, StorageQuery},
 		hex_string,
 		subscription::{SubscriptionManagement, SubscriptionManagementError},
 	},
@@ -329,19 +329,10 @@ where
 		let items = items
 			.into_iter()
 			.map(|query| {
-				if query.query_type == StorageQueryType::ClosestDescendantMerkleValue {
-					// Note: remove this once all types are implemented.
-					return Err(ChainHeadRpcError::InvalidParam(
-						"Storage query type not supported".into(),
-					))
-				}
-
-				Ok(StorageQuery {
-					key: StorageKey(parse_hex_param(query.key)?),
-					query_type: query.query_type,
-				})
+				let key = StorageKey(parse_hex_param(query.key)?);
+				Ok(StorageQuery { key, query_type: query.query_type })
 			})
-			.collect::<Result<Vec<_>, _>>()?;
+			.collect::<Result<Vec<_>, ChainHeadRpcError>>()?;
 
 		let child_trie = child_trie
 			.map(|child_trie| parse_hex_param(child_trie))

substrate/client/rpc-spec-v2/src/chain_head/chain_head_storage.rs

@@ -145,6 +145,36 @@ where
 			.unwrap_or_else(|error| QueryResult::Err(error.to_string()))
 	}
 
+	/// Fetch the closest merkle value.
+	fn query_storage_merkle_value(
+		&self,
+		hash: Block::Hash,
+		key: &StorageKey,
+		child_key: Option<&ChildInfo>,
+	) -> QueryResult {
+		let result = if let Some(child_key) = child_key {
+			self.client.child_closest_merkle_value(hash, child_key, key)
+		} else {
+			self.client.closest_merkle_value(hash, key)
+		};
+
+		result
+			.map(|opt| {
+				QueryResult::Ok(opt.map(|storage_data| {
+					let result = match &storage_data {
+						sc_client_api::MerkleValue::Node(data) => hex_string(&data.as_slice()),
+						sc_client_api::MerkleValue::Hash(hash) => hex_string(&hash.as_ref()),
+					};
+
+					StorageResult {
+						key: hex_string(&key.0),
+						result: StorageResultType::ClosestDescendantMerkleValue(result),
+					}
+				}))
+			})
+			.unwrap_or_else(|error| QueryResult::Err(error.to_string()))
+	}
+
 	/// Iterate over at most `operation_max_storage_items` keys.
 	///
 	/// Returns the storage result with a potential next key to resume iteration.
@@ -286,13 +316,21 @@ where
 							return
 						},
 					},
+				StorageQueryType::ClosestDescendantMerkleValue =>
+					match self.query_storage_merkle_value(hash, &item.key, child_key.as_ref()) {
+						Ok(Some(value)) => storage_results.push(value),
+						Ok(None) => continue,
+						Err(error) => {
+							send_error::<Block>(&sender, operation.operation_id(), error);
+							return
+						},
+					},
 				StorageQueryType::DescendantsValues => self
 					.iter_operations
 					.push_back(QueryIter { next_key: item.key, ty: IterQueryType::Value }),
 				StorageQueryType::DescendantsHashes => self
 					.iter_operations
 					.push_back(QueryIter { next_key: item.key, ty: IterQueryType::Hash }),
-				_ => continue,
 			};
 		}
 

substrate/client/rpc-spec-v2/src/chain_head/test_utils.rs

@@ -20,8 +20,8 @@ use parking_lot::Mutex;
 use sc_client_api::{
 	execution_extensions::ExecutionExtensions, BlockBackend, BlockImportNotification,
 	BlockchainEvents, CallExecutor, ChildInfo, ExecutorProvider, FinalityNotification,
-	FinalityNotifications, FinalizeSummary, ImportNotifications, KeysIter, PairsIter, StorageData,
-	StorageEventStream, StorageKey, StorageProvider,
+	FinalityNotifications, FinalizeSummary, ImportNotifications, KeysIter, MerkleValue, PairsIter,
+	StorageData, StorageEventStream, StorageKey, StorageProvider,
 };
 use sc_utils::mpsc::{tracing_unbounded, TracingUnboundedSender};
 use sp_api::{CallApiAt, CallApiAtParams, NumberFor, RuntimeVersion};
@@ -198,6 +198,23 @@ impl<
 	) -> sp_blockchain::Result<Option<Block::Hash>> {
 		self.client.child_storage_hash(hash, child_info, key)
 	}
+
+	fn closest_merkle_value(
+		&self,
+		hash: Block::Hash,
+		key: &StorageKey,
+	) -> sp_blockchain::Result<Option<MerkleValue<Block::Hash>>> {
+		self.client.closest_merkle_value(hash, key)
+	}
+
+	fn child_closest_merkle_value(
+		&self,
+		hash: Block::Hash,
+		child_info: &ChildInfo,
+		key: &StorageKey,
+	) -> sp_blockchain::Result<Option<MerkleValue<Block::Hash>>> {
+		self.client.child_closest_merkle_value(hash, child_info, key)
+	}
 }
 
 impl<Block: BlockT, Client: CallApiAt<Block>> CallApiAt<Block> for ChainHeadMockClient<Client> {