chainHead: Add support for storage closest merkle descendant #14818 (#1153)

This PR adds support for fetching the closest merkle value of some key.


Builds on top of
- https://github.com/paritytech/trie/pull/199

Migrates https://github.com/paritytech/substrate/pull/14818 to the
monorepo.
Closes: https://github.com/paritytech/substrate/issues/14550
Closes: https://github.com/paritytech/polkadot-sdk/issues/1506

// @paritytech/subxt-team

---------

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>
Co-authored-by: Sebastian Kunert <skunert49@gmail.com>

substrate/primitives/state-machine/src/backend.rs

@@ -30,7 +30,7 @@ use sp_core::storage::{ChildInfo, StateVersion, TrackedStorageKey};
 #[cfg(feature = "std")]
 use sp_core::traits::RuntimeCode;
 use sp_std::vec::Vec;
-use sp_trie::PrefixedMemoryDB;
+use sp_trie::{MerkleValue, PrefixedMemoryDB};
 
 /// A struct containing arguments for iterating over the storage.
 #[derive(Default)]
@@ -195,7 +195,17 @@ pub trait Backend<H: Hasher>: sp_std::fmt::Debug {
 	/// Get keyed storage value hash or None if there is nothing associated.
 	fn storage_hash(&self, key: &[u8]) -> Result<Option<H::Out>, Self::Error>;
 
-	/// Get keyed child storage or None if there is nothing associated.
+	/// Get the merkle value or None if there is nothing associated.
+	fn closest_merkle_value(&self, key: &[u8]) -> Result<Option<MerkleValue<H::Out>>, Self::Error>;
+
+	/// Get the child merkle value or None if there is nothing associated.
+	fn child_closest_merkle_value(
+		&self,
+		child_info: &ChildInfo,
+		key: &[u8],
+	) -> Result<Option<MerkleValue<H::Out>>, Self::Error>;
+
+	/// Get child keyed child storage or None if there is nothing associated.
 	fn child_storage(
 		&self,
 		child_info: &ChildInfo,

substrate/primitives/state-machine/src/trie_backend.rs

@@ -30,7 +30,6 @@ use codec::Codec;
 use hash_db::HashDB;
 use hash_db::Hasher;
 use sp_core::storage::{ChildInfo, StateVersion};
-use sp_trie::PrefixedMemoryDB;
 #[cfg(feature = "std")]
 use sp_trie::{
 	cache::{LocalTrieCache, TrieCache},
@@ -39,6 +38,7 @@ use sp_trie::{
 };
 #[cfg(not(feature = "std"))]
 use sp_trie::{Error, NodeCodec};
+use sp_trie::{MerkleValue, PrefixedMemoryDB};
 use trie_db::TrieCache as TrieCacheT;
 #[cfg(not(feature = "std"))]
 use trie_db::{node::NodeOwned, CachedValue};
@@ -405,6 +405,18 @@ where
 		self.essence.child_storage(child_info, key)
 	}
 
+	fn closest_merkle_value(&self, key: &[u8]) -> Result<Option<MerkleValue<H::Out>>, Self::Error> {
+		self.essence.closest_merkle_value(key)
+	}
+
+	fn child_closest_merkle_value(
+		&self,
+		child_info: &ChildInfo,
+		key: &[u8],
+	) -> Result<Option<MerkleValue<H::Out>>, Self::Error> {
+		self.essence.child_closest_merkle_value(child_info, key)
+	}
+
 	fn next_storage_key(&self, key: &[u8]) -> Result<Option<StorageKey>, Self::Error> {
 		let (is_cached, mut cache) = access_cache(&self.next_storage_key_cache, Option::take)
 			.map(|cache| (cache.last_key == key, cache))

substrate/primitives/state-machine/src/trie_backend_essence.rs

@@ -32,11 +32,12 @@ use sp_std::{boxed::Box, marker::PhantomData, vec::Vec};
 #[cfg(feature = "std")]
 use sp_trie::recorder::Recorder;
 use sp_trie::{
-	child_delta_trie_root, delta_trie_root, empty_child_trie_root, read_child_trie_hash,
-	read_child_trie_value, read_trie_value,
+	child_delta_trie_root, delta_trie_root, empty_child_trie_root,
+	read_child_trie_first_descedant_value, read_child_trie_hash, read_child_trie_value,
+	read_trie_first_descedant_value, read_trie_value,
 	trie_types::{TrieDBBuilder, TrieError},
-	DBValue, KeySpacedDB, NodeCodec, PrefixedMemoryDB, Trie, TrieCache, TrieDBRawIterator,
-	TrieRecorder,
+	DBValue, KeySpacedDB, MerkleValue, NodeCodec, PrefixedMemoryDB, Trie, TrieCache,
+	TrieDBRawIterator, TrieRecorder,
 };
 #[cfg(feature = "std")]
 use std::{collections::HashMap, sync::Arc};
@@ -574,6 +575,39 @@ where
 		})
 	}
 
+	/// Get the closest merkle value at given key.
+	pub fn closest_merkle_value(&self, key: &[u8]) -> Result<Option<MerkleValue<H::Out>>> {
+		let map_e = |e| format!("Trie lookup error: {}", e);
+
+		self.with_recorder_and_cache(None, |recorder, cache| {
+			read_trie_first_descedant_value::<Layout<H>, _>(self, &self.root, key, recorder, cache)
+				.map_err(map_e)
+		})
+	}
+
+	/// Get the child closest merkle value at given key.
+	pub fn child_closest_merkle_value(
+		&self,
+		child_info: &ChildInfo,
+		key: &[u8],
+	) -> Result<Option<MerkleValue<H::Out>>> {
+		let Some(child_root) = self.child_root(child_info)? else { return Ok(None) };
+
+		let map_e = |e| format!("Trie lookup error: {}", e);
+
+		self.with_recorder_and_cache(Some(child_root), |recorder, cache| {
+			read_child_trie_first_descedant_value::<Layout<H>, _>(
+				child_info.keyspace(),
+				self,
+				&child_root,
+				key,
+				recorder,
+				cache,
+			)
+			.map_err(map_e)
+		})
+	}
+
 	/// Create a raw iterator over the storage.
 	pub fn raw_iter(&self, args: IterArgs) -> Result<RawIter<S, H, C>> {
 		let root = if let Some(child_info) = args.child_info.as_ref() {