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() {

substrate/primitives/trie/src/lib.rs

@@ -44,7 +44,6 @@ pub use storage_proof::{CompactProof, StorageProof};
 /// Trie codec reexport, mainly child trie support
 /// for trie compact proof.
 pub use trie_codec::{decode_compact, encode_compact, Error as CompactProofError};
-pub use trie_db::proof::VerifyError;
 use trie_db::proof::{generate_proof, verify_proof};
 /// Various re-exports from the `trie-db` crate.
 pub use trie_db::{
@@ -53,6 +52,7 @@ pub use trie_db::{
 	CError, DBValue, Query, Recorder, Trie, TrieCache, TrieConfiguration, TrieDBIterator,
 	TrieDBKeyIterator, TrieDBRawIterator, TrieLayout, TrieMut, TrieRecorder,
 };
+pub use trie_db::{proof::VerifyError, MerkleValue};
 /// The Substrate format implementation of `TrieStream`.
 pub use trie_stream::TrieStream;
 
@@ -295,6 +295,25 @@ pub fn read_trie_value<L: TrieLayout, DB: hash_db::HashDBRef<L::Hash, trie_db::D
 		.get(key)
 }
 
+/// Read the [`trie_db::MerkleValue`] of the node that is the closest descendant for
+/// the provided key.
+pub fn read_trie_first_descedant_value<L: TrieLayout, DB>(
+	db: &DB,
+	root: &TrieHash<L>,
+	key: &[u8],
+	recorder: Option<&mut dyn TrieRecorder<TrieHash<L>>>,
+	cache: Option<&mut dyn TrieCache<L::Codec>>,
+) -> Result<Option<MerkleValue<TrieHash<L>>>, Box<TrieError<L>>>
+where
+	DB: hash_db::HashDBRef<L::Hash, trie_db::DBValue>,
+{
+	TrieDBBuilder::<L>::new(db, root)
+		.with_optional_cache(cache)
+		.with_optional_recorder(recorder)
+		.build()
+		.lookup_first_descendant(key)
+}
+
 /// Read a value from the trie with given Query.
 pub fn read_trie_value_with<
 	L: TrieLayout,
@@ -397,6 +416,27 @@ where
 		.get_hash(key)
 }
 
+/// Read the [`trie_db::MerkleValue`] of the node that is the closest descendant for
+/// the provided child key.
+pub fn read_child_trie_first_descedant_value<L: TrieConfiguration, DB>(
+	keyspace: &[u8],
+	db: &DB,
+	root: &TrieHash<L>,
+	key: &[u8],
+	recorder: Option<&mut dyn TrieRecorder<TrieHash<L>>>,
+	cache: Option<&mut dyn TrieCache<L::Codec>>,
+) -> Result<Option<MerkleValue<TrieHash<L>>>, Box<TrieError<L>>>
+where
+	DB: hash_db::HashDBRef<L::Hash, trie_db::DBValue>,
+{
+	let db = KeySpacedDB::new(db, keyspace);
+	TrieDBBuilder::<L>::new(&db, &root)
+		.with_optional_recorder(recorder)
+		.with_optional_cache(cache)
+		.build()
+		.lookup_first_descendant(key)
+}
+
 /// Read a value from the child trie with given query.
 pub fn read_child_trie_value_with<L, Q, DB>(
 	keyspace: &[u8],