InMemoryBackend: Make it generic over the key hasher (paritytech#11488)

* InMemoryBackend: Make it generic over the key hasher

* Update primitives/state-machine/src/in_memory_backend.rs

Co-authored-by: cheme <emericchevalier.pro@gmail.com>

* Update primitives/state-machine/src/in_memory_backend.rs

Co-authored-by: cheme <emericchevalier.pro@gmail.com>

* FMT

Co-authored-by: cheme <emericchevalier.pro@gmail.com>

primitives/state-machine/src/basic.rs

@@ -30,7 +30,7 @@ use sp_core::{
 	Blake2Hasher,
 };
 use sp_externalities::{Extension, Extensions};
-use sp_trie::{empty_child_trie_root, LayoutV0, LayoutV1, TrieConfiguration};
+use sp_trie::{empty_child_trie_root, HashKey, LayoutV0, LayoutV1, TrieConfiguration};
 use std::{
 	any::{Any, TypeId},
 	collections::BTreeMap,
@@ -310,7 +310,7 @@ impl Externalities for BasicExternalities {
 	) -> Vec<u8> {
 		if let Some(child) = self.inner.children_default.get(child_info.storage_key()) {
 			let delta = child.data.iter().map(|(k, v)| (k.as_ref(), Some(v.as_ref())));
-			crate::in_memory_backend::new_in_mem::<Blake2Hasher>()
+			crate::in_memory_backend::new_in_mem::<Blake2Hasher, HashKey<_>>()
 				.child_storage_root(&child.child_info, delta, state_version)
 				.0
 		} else {

primitives/state-machine/src/in_memory_backend.rs

@@ -23,22 +23,36 @@ use crate::{
 use codec::Codec;
 use hash_db::Hasher;
 use sp_core::storage::{ChildInfo, StateVersion, Storage};
-use sp_trie::{empty_trie_root, LayoutV1, MemoryDB};
+use sp_trie::{empty_trie_root, GenericMemoryDB, HashKey, KeyFunction, LayoutV1, MemoryDB};
 use std::collections::{BTreeMap, HashMap};
 
 /// Create a new empty instance of in-memory backend.
-pub fn new_in_mem<H: Hasher>() -> TrieBackend<MemoryDB<H>, H>
+///
+/// It will use [`HashKey`] to store the keys internally.
+pub fn new_in_mem_hash_key<H>() -> TrieBackend<MemoryDB<H>, H>
 where
+	H: Hasher,
 	H::Out: Codec + Ord,
 {
-	let db = MemoryDB::default();
+	new_in_mem::<H, HashKey<H>>()
+}
+
+/// Create a new empty instance of in-memory backend.
+pub fn new_in_mem<H, KF>() -> TrieBackend<GenericMemoryDB<H, KF>, H>
+where
+	H: Hasher,
+	H::Out: Codec + Ord,
+	KF: KeyFunction<H> + Send + Sync,
+{
+	let db = GenericMemoryDB::default();
 	// V1 is same as V0 for an empty trie.
 	TrieBackend::new(db, empty_trie_root::<LayoutV1<H>>())
 }
 
-impl<H: Hasher> TrieBackend<MemoryDB<H>, H>
+impl<H: Hasher, KF> TrieBackend<GenericMemoryDB<H, KF>, H>
 where
 	H::Out: Codec + Ord,
+	KF: KeyFunction<H> + Send + Sync,
 {
 	/// Copy the state, with applied updates
 	pub fn update<T: IntoIterator<Item = (Option<ChildInfo>, StorageCollection)>>(
@@ -70,14 +84,14 @@ where
 	}
 
 	/// Merge trie nodes into this backend.
-	pub fn update_backend(&self, root: H::Out, changes: MemoryDB<H>) -> Self {
+	pub fn update_backend(&self, root: H::Out, changes: GenericMemoryDB<H, KF>) -> Self {
 		let mut clone = self.backend_storage().clone();
 		clone.consolidate(changes);
 		Self::new(clone, root)
 	}
 
 	/// Apply the given transaction to this backend and set the root to the given value.
-	pub fn apply_transaction(&mut self, root: H::Out, transaction: MemoryDB<H>) {
+	pub fn apply_transaction(&mut self, root: H::Out, transaction: GenericMemoryDB<H, KF>) {
 		let mut storage = sp_std::mem::take(self).into_storage();
 		storage.consolidate(transaction);
 		*self = TrieBackend::new(storage, root);
@@ -89,28 +103,33 @@ where
 	}
 }
 
-impl<H: Hasher> Clone for TrieBackend<MemoryDB<H>, H>
+impl<H: Hasher, KF> Clone for TrieBackend<GenericMemoryDB<H, KF>, H>
 where
 	H::Out: Codec + Ord,
+	KF: KeyFunction<H> + Send + Sync,
 {
 	fn clone(&self) -> Self {
 		TrieBackend::new(self.backend_storage().clone(), *self.root())
 	}
 }
 
-impl<H: Hasher> Default for TrieBackend<MemoryDB<H>, H>
+impl<H, KF> Default for TrieBackend<GenericMemoryDB<H, KF>, H>
 where
+	H: Hasher,
 	H::Out: Codec + Ord,
+	KF: KeyFunction<H> + Send + Sync,
 {
 	fn default() -> Self {
 		new_in_mem()
 	}
 }
 
-impl<H: Hasher> From<(HashMap<Option<ChildInfo>, BTreeMap<StorageKey, StorageValue>>, StateVersion)>
-	for TrieBackend<MemoryDB<H>, H>
+impl<H: Hasher, KF>
+	From<(HashMap<Option<ChildInfo>, BTreeMap<StorageKey, StorageValue>>, StateVersion)>
+	for TrieBackend<GenericMemoryDB<H, KF>, H>
 where
 	H::Out: Codec + Ord,
+	KF: KeyFunction<H> + Send + Sync,
 {
 	fn from(
 		(inner, state_version): (
@@ -129,9 +148,10 @@ where
 	}
 }
 
-impl<H: Hasher> From<(Storage, StateVersion)> for TrieBackend<MemoryDB<H>, H>
+impl<H: Hasher, KF> From<(Storage, StateVersion)> for TrieBackend<GenericMemoryDB<H, KF>, H>
 where
 	H::Out: Codec + Ord,
+	KF: KeyFunction<H> + Send + Sync,
 {
 	fn from((inners, state_version): (Storage, StateVersion)) -> Self {
 		let mut inner: HashMap<Option<ChildInfo>, BTreeMap<StorageKey, StorageValue>> = inners
@@ -144,10 +164,11 @@ where
 	}
 }
 
-impl<H: Hasher> From<(BTreeMap<StorageKey, StorageValue>, StateVersion)>
-	for TrieBackend<MemoryDB<H>, H>
+impl<H: Hasher, KF> From<(BTreeMap<StorageKey, StorageValue>, StateVersion)>
+	for TrieBackend<GenericMemoryDB<H, KF>, H>
 where
 	H::Out: Codec + Ord,
+	KF: KeyFunction<H> + Send + Sync,
 {
 	fn from((inner, state_version): (BTreeMap<StorageKey, StorageValue>, StateVersion)) -> Self {
 		let mut expanded = HashMap::new();
@@ -156,10 +177,11 @@ where
 	}
 }
 
-impl<H: Hasher> From<(Vec<(Option<ChildInfo>, StorageCollection)>, StateVersion)>
-	for TrieBackend<MemoryDB<H>, H>
+impl<H: Hasher, KF> From<(Vec<(Option<ChildInfo>, StorageCollection)>, StateVersion)>
+	for TrieBackend<GenericMemoryDB<H, KF>, H>
 where
 	H::Out: Codec + Ord,
+	KF: KeyFunction<H> + Send + Sync,
 {
 	fn from(
 		(inner, state_version): (Vec<(Option<ChildInfo>, StorageCollection)>, StateVersion),
@@ -189,7 +211,7 @@ mod tests {
 	#[test]
 	fn in_memory_with_child_trie_only() {
 		let state_version = StateVersion::default();
-		let storage = new_in_mem::<BlakeTwo256>();
+		let storage = new_in_mem_hash_key::<BlakeTwo256>();
 		let child_info = ChildInfo::new_default(b"1");
 		let child_info = &child_info;
 		let storage = storage.update(
@@ -205,7 +227,7 @@ mod tests {
 	#[test]
 	fn insert_multiple_times_child_data_works() {
 		let state_version = StateVersion::default();
-		let mut storage = new_in_mem::<BlakeTwo256>();
+		let mut storage = new_in_mem_hash_key::<BlakeTwo256>();
 		let child_info = ChildInfo::new_default(b"1");
 
 		storage.insert(

primitives/state-machine/src/lib.rs

@@ -143,7 +143,7 @@ mod std_reexport {
 	pub use crate::{
 		basic::BasicExternalities,
 		error::{Error, ExecutionError},
-		in_memory_backend::new_in_mem,
+		in_memory_backend::{new_in_mem, new_in_mem_hash_key},
 		proving_backend::{
 			create_proof_check_backend, ProofRecorder, ProvingBackend, ProvingBackendRecorder,
 		},
@@ -1347,7 +1347,7 @@ mod execution {
 #[cfg(test)]
 mod tests {
 	use super::{ext::Ext, *};
-	use crate::execution::CallResult;
+	use crate::{execution::CallResult, in_memory_backend::new_in_mem_hash_key};
 	use codec::{Decode, Encode};
 	use sp_core::{
 		map,
@@ -1687,7 +1687,7 @@ mod tests {
 	fn set_child_storage_works() {
 		let child_info = ChildInfo::new_default(b"sub1");
 		let child_info = &child_info;
-		let state = new_in_mem::<BlakeTwo256>();
+		let state = new_in_mem_hash_key::<BlakeTwo256>();
 		let backend = state.as_trie_backend().unwrap();
 		let mut overlay = OverlayedChanges::default();
 		let mut cache = StorageTransactionCache::default();
@@ -1703,7 +1703,7 @@ mod tests {
 	fn append_storage_works() {
 		let reference_data = vec![b"data1".to_vec(), b"2".to_vec(), b"D3".to_vec(), b"d4".to_vec()];
 		let key = b"key".to_vec();
-		let state = new_in_mem::<BlakeTwo256>();
+		let state = new_in_mem_hash_key::<BlakeTwo256>();
 		let backend = state.as_trie_backend().unwrap();
 		let mut overlay = OverlayedChanges::default();
 		let mut cache = StorageTransactionCache::default();
@@ -1740,7 +1740,7 @@ mod tests {
 
 		let key = b"events".to_vec();
 		let mut cache = StorageTransactionCache::default();
-		let state = new_in_mem::<BlakeTwo256>();
+		let state = new_in_mem_hash_key::<BlakeTwo256>();
 		let backend = state.as_trie_backend().unwrap();
 		let mut overlay = OverlayedChanges::default();
 

primitives/trie/src/lib.rs

@@ -31,9 +31,8 @@ pub use error::Error;
 /// Various re-exports from the `hash-db` crate.
 pub use hash_db::{HashDB as HashDBT, EMPTY_PREFIX};
 use hash_db::{Hasher, Prefix};
-pub use memory_db::prefixed_key;
 /// Various re-exports from the `memory-db` crate.
-pub use memory_db::KeyFunction;
+pub use memory_db::{prefixed_key, HashKey, KeyFunction, PrefixedKey};
 /// The Substrate format implementation of `NodeCodec`.
 pub use node_codec::NodeCodec;
 use sp_std::{borrow::Borrow, boxed::Box, marker::PhantomData, vec::Vec};