Set StateBackend::Transaction to PrefixedMemoryDB (#14612)

* Yep

* Try to get it working everywhere

* Make `from_raw_storage` start with an empty db

* More fixes!

* Make everything compile

* Fix `child_storage_root`

* Fix after merge

* Cleanups

* Update primitives/state-machine/src/overlayed_changes/mod.rs

Co-authored-by: Davide Galassi <davxy@datawok.net>

* Review comments

* Fix issues

* Silence warning

* FMT

* Clippy

---------

Co-authored-by: Davide Galassi <davxy@datawok.net>

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

@@ -30,6 +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;
 
 /// A struct containing arguments for iterating over the storage.
 #[derive(Default)]
@@ -168,6 +169,12 @@ where
 	}
 }
 
+/// The transaction type used by [`Backend`].
+///
+/// This transaction contains all the changes that need to be applied to the backend to create the
+/// state for a new block.
+pub type BackendTransaction<H> = PrefixedMemoryDB<H>;
+
 /// A state backend is used to read state data and can have changes committed
 /// to it.
 ///
@@ -176,11 +183,8 @@ pub trait Backend<H: Hasher>: sp_std::fmt::Debug {
 	/// An error type when fetching data is not possible.
 	type Error: super::Error;
 
-	/// Storage changes to be applied if committing
-	type Transaction: Consolidate + Default + Send;
-
 	/// Type of trie backend storage.
-	type TrieBackendStorage: TrieBackendStorage<H, Overlay = Self::Transaction>;
+	type TrieBackendStorage: TrieBackendStorage<H>;
 
 	/// Type of the raw storage iterator.
 	type RawIter: StorageIterator<H, Backend = Self, Error = Self::Error>;
@@ -236,7 +240,7 @@ pub trait Backend<H: Hasher>: sp_std::fmt::Debug {
 		&self,
 		delta: impl Iterator<Item = (&'a [u8], Option<&'a [u8]>)>,
 		state_version: StateVersion,
-	) -> (H::Out, Self::Transaction)
+	) -> (H::Out, BackendTransaction<H>)
 	where
 		H::Out: Ord;
 
@@ -248,7 +252,7 @@ pub trait Backend<H: Hasher>: sp_std::fmt::Debug {
 		child_info: &ChildInfo,
 		delta: impl Iterator<Item = (&'a [u8], Option<&'a [u8]>)>,
 		state_version: StateVersion,
-	) -> (H::Out, bool, Self::Transaction)
+	) -> (H::Out, bool, BackendTransaction<H>)
 	where
 		H::Out: Ord;
 
@@ -283,11 +287,11 @@ pub trait Backend<H: Hasher>: sp_std::fmt::Debug {
 			Item = (&'a ChildInfo, impl Iterator<Item = (&'a [u8], Option<&'a [u8]>)>),
 		>,
 		state_version: StateVersion,
-	) -> (H::Out, Self::Transaction)
+	) -> (H::Out, BackendTransaction<H>)
 	where
 		H::Out: Ord + Encode,
 	{
-		let mut txs: Self::Transaction = Default::default();
+		let mut txs = BackendTransaction::default();
 		let mut child_roots: Vec<_> = Default::default();
 		// child first
 		for (child_info, child_delta) in child_deltas {
@@ -308,6 +312,7 @@ pub trait Backend<H: Hasher>: sp_std::fmt::Debug {
 			state_version,
 		);
 		txs.consolidate(parent_txs);
+
 		(root, txs)
 	}
 
@@ -331,7 +336,7 @@ pub trait Backend<H: Hasher>: sp_std::fmt::Debug {
 	fn commit(
 		&self,
 		_: H::Out,
-		_: Self::Transaction,
+		_: BackendTransaction<H>,
 		_: StorageCollection,
 		_: ChildStorageCollection,
 	) -> Result<(), Self::Error> {
@@ -377,34 +382,6 @@ pub trait AsTrieBackend<H: Hasher, C = sp_trie::cache::LocalTrieCache<H>> {
 	fn as_trie_backend(&self) -> &TrieBackend<Self::TrieBackendStorage, H, C>;
 }
 
-/// Trait that allows consolidate two transactions together.
-pub trait Consolidate {
-	/// Consolidate two transactions into one.
-	fn consolidate(&mut self, other: Self);
-}
-
-impl Consolidate for () {
-	fn consolidate(&mut self, _: Self) {
-		()
-	}
-}
-
-impl Consolidate for Vec<(Option<ChildInfo>, StorageCollection)> {
-	fn consolidate(&mut self, mut other: Self) {
-		self.append(&mut other);
-	}
-}
-
-impl<H, KF> Consolidate for sp_trie::GenericMemoryDB<H, KF>
-where
-	H: Hasher,
-	KF: sp_trie::KeyFunction<H>,
-{
-	fn consolidate(&mut self, other: Self) {
-		sp_trie::GenericMemoryDB::consolidate(self, other)
-	}
-}
-
 /// Wrapper to create a [`RuntimeCode`] from a type that implements [`Backend`].
 #[cfg(feature = "std")]
 pub struct BackendRuntimeCode<'a, B, H> {

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

@@ -29,7 +29,7 @@ use sp_core::{
 	Blake2Hasher,
 };
 use sp_externalities::{Extension, Extensions, MultiRemovalResults};
-use sp_trie::{empty_child_trie_root, HashKey, LayoutV0, LayoutV1, TrieConfiguration};
+use sp_trie::{empty_child_trie_root, LayoutV0, LayoutV1, TrieConfiguration};
 use std::{
 	any::{Any, TypeId},
 	collections::BTreeMap,
@@ -39,7 +39,7 @@ use std::{
 /// Simple Map-based Externalities impl.
 #[derive(Debug)]
 pub struct BasicExternalities {
-	overlay: OverlayedChanges,
+	overlay: OverlayedChanges<Blake2Hasher>,
 	extensions: Extensions,
 }
 
@@ -282,7 +282,7 @@ impl Externalities for BasicExternalities {
 		if let Some((data, child_info)) = self.overlay.child_changes(child_info.storage_key()) {
 			let delta =
 				data.into_iter().map(|(k, v)| (k.as_ref(), v.value().map(|v| v.as_slice())));
-			crate::in_memory_backend::new_in_mem::<Blake2Hasher, HashKey<_>>()
+			crate::in_memory_backend::new_in_mem::<Blake2Hasher>()
 				.child_storage_root(&child_info, delta, state_version)
 				.0
 		} else {

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

@@ -22,7 +22,7 @@ use crate::overlayed_changes::OverlayedExtensions;
 use crate::{
 	backend::Backend, IndexOperation, IterArgs, OverlayedChanges, StorageKey, StorageValue,
 };
-use codec::{Decode, Encode, EncodeAppend};
+use codec::{Encode, EncodeAppend};
 use hash_db::Hasher;
 #[cfg(feature = "std")]
 use sp_core::hexdisplay::HexDisplay;
@@ -30,9 +30,8 @@ use sp_core::storage::{
 	well_known_keys::is_child_storage_key, ChildInfo, StateVersion, TrackedStorageKey,
 };
 use sp_externalities::{Extension, ExtensionStore, Externalities, MultiRemovalResults};
-use sp_trie::{empty_child_trie_root, LayoutV1};
 
-use crate::{log_error, trace, warn, StorageTransactionCache};
+use crate::{log_error, trace, warn};
 use sp_std::{
 	any::{Any, TypeId},
 	boxed::Box,
@@ -98,11 +97,9 @@ where
 	B: 'a + Backend<H>,
 {
 	/// The overlayed changes to write to.
-	overlay: &'a mut OverlayedChanges,
+	overlay: &'a mut OverlayedChanges<H>,
 	/// The storage backend to read from.
 	backend: &'a B,
-	/// The cache for the storage transactions.
-	storage_transaction_cache: &'a mut StorageTransactionCache<B::Transaction, H>,
 	/// Pseudo-unique id used for tracing.
 	pub id: u16,
 	/// Extensions registered with this instance.
@@ -117,37 +114,24 @@ where
 {
 	/// Create a new `Ext`.
 	#[cfg(not(feature = "std"))]
-	pub fn new(
-		overlay: &'a mut OverlayedChanges,
-		storage_transaction_cache: &'a mut StorageTransactionCache<B::Transaction, H>,
-		backend: &'a B,
-	) -> Self {
-		Ext { overlay, backend, id: 0, storage_transaction_cache }
+	pub fn new(overlay: &'a mut OverlayedChanges<H>, backend: &'a B) -> Self {
+		Ext { overlay, backend, id: 0 }
 	}
 
 	/// Create a new `Ext` from overlayed changes and read-only backend
 	#[cfg(feature = "std")]
 	pub fn new(
-		overlay: &'a mut OverlayedChanges,
-		storage_transaction_cache: &'a mut StorageTransactionCache<B::Transaction, H>,
+		overlay: &'a mut OverlayedChanges<H>,
 		backend: &'a B,
 		extensions: Option<&'a mut sp_externalities::Extensions>,
 	) -> Self {
 		Self {
 			overlay,
 			backend,
-			storage_transaction_cache,
 			id: rand::random(),
 			extensions: extensions.map(OverlayedExtensions::new),
 		}
 	}
-
-	/// Invalidates the currently cached storage root and the db transaction.
-	///
-	/// Called when there are changes that likely will invalidate the storage root.
-	fn mark_dirty(&mut self) {
-		self.storage_transaction_cache.reset();
-	}
 }
 
 #[cfg(test)]
@@ -412,7 +396,6 @@ where
 			),
 		);
 
-		self.mark_dirty();
 		self.overlay.set_storage(key, value);
 	}
 
@@ -432,7 +415,6 @@ where
 		);
 		let _guard = guard();
 
-		self.mark_dirty();
 		self.overlay.set_child_storage(child_info, key, value);
 	}
 
@@ -449,7 +431,6 @@ where
 			child_info = %HexDisplay::from(&child_info.storage_key()),
 		);
 		let _guard = guard();
-		self.mark_dirty();
 		let overlay = self.overlay.clear_child_storage(child_info);
 		let (maybe_cursor, backend, loops) =
 			self.limit_remove_from_backend(Some(child_info), None, maybe_limit, maybe_cursor);
@@ -478,7 +459,6 @@ where
 			return MultiRemovalResults { maybe_cursor: None, backend: 0, unique: 0, loops: 0 }
 		}
 
-		self.mark_dirty();
 		let overlay = self.overlay.clear_prefix(prefix);
 		let (maybe_cursor, backend, loops) =
 			self.limit_remove_from_backend(None, Some(prefix), maybe_limit, maybe_cursor);
@@ -501,7 +481,6 @@ where
 		);
 		let _guard = guard();
 
-		self.mark_dirty();
 		let overlay = self.overlay.clear_child_prefix(child_info, prefix);
 		let (maybe_cursor, backend, loops) = self.limit_remove_from_backend(
 			Some(child_info),
@@ -522,7 +501,6 @@ where
 		);
 
 		let _guard = guard();
-		self.mark_dirty();
 
 		let backend = &mut self.backend;
 		let current_value = self.overlay.value_mut_or_insert_with(&key, || {
@@ -533,27 +511,17 @@ where
 
 	fn storage_root(&mut self, state_version: StateVersion) -> Vec<u8> {
 		let _guard = guard();
-		if let Some(ref root) = self.storage_transaction_cache.transaction_storage_root {
-			trace!(
-				target: "state",
-				method = "StorageRoot",
-				ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
-				storage_root = %HexDisplay::from(&root.as_ref()),
-				cached = true,
-			);
-			return root.encode()
-		}
 
-		let root =
-			self.overlay
-				.storage_root(self.backend, self.storage_transaction_cache, state_version);
+		let (root, _cached) = self.overlay.storage_root(self.backend, state_version);
+
 		trace!(
 			target: "state",
 			method = "StorageRoot",
 			ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
 			storage_root = %HexDisplay::from(&root.as_ref()),
-			cached = false,
+			cached = %_cached,
 		);
+
 		root.encode()
 	}
 
@@ -563,74 +531,22 @@ where
 		state_version: StateVersion,
 	) -> Vec<u8> {
 		let _guard = guard();
-		let storage_key = child_info.storage_key();
-		let prefixed_storage_key = child_info.prefixed_storage_key();
-		if self.storage_transaction_cache.transaction_storage_root.is_some() {
-			let root = self
-				.storage(prefixed_storage_key.as_slice())
-				.and_then(|k| Decode::decode(&mut &k[..]).ok())
-				// V1 is equivalent to V0 on empty root.
-				.unwrap_or_else(empty_child_trie_root::<LayoutV1<H>>);
-			trace!(
-				target: "state",
-				method = "ChildStorageRoot",
-				ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
-				child_info = %HexDisplay::from(&storage_key),
-				storage_root = %HexDisplay::from(&root.as_ref()),
-				cached = true,
-			);
-			root.encode()
-		} else {
-			let root = if let Some((changes, info)) = self.overlay.child_changes(storage_key) {
-				let delta = changes.map(|(k, v)| (k.as_ref(), v.value().map(AsRef::as_ref)));
-				Some(self.backend.child_storage_root(info, delta, state_version))
-			} else {
-				None
-			};
 
-			if let Some((root, is_empty, _)) = root {
-				let root = root.encode();
-				// We store update in the overlay in order to be able to use
-				// 'self.storage_transaction' cache. This is brittle as it rely on Ext only querying
-				// the trie backend for storage root.
-				// A better design would be to manage 'child_storage_transaction' in a
-				// similar way as 'storage_transaction' but for each child trie.
-				if is_empty {
-					self.overlay.set_storage(prefixed_storage_key.into_inner(), None);
-				} else {
-					self.overlay.set_storage(prefixed_storage_key.into_inner(), Some(root.clone()));
-				}
+		let (root, _cached) = self
+			.overlay
+			.child_storage_root(child_info, self.backend, state_version)
+			.expect(EXT_NOT_ALLOWED_TO_FAIL);
 
-				trace!(
-					target: "state",
-					method = "ChildStorageRoot",
-					ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
-					child_info = %HexDisplay::from(&storage_key),
-					storage_root = %HexDisplay::from(&root.as_ref()),
-					cached = false,
-				);
+		trace!(
+			target: "state",
+			method = "ChildStorageRoot",
+			ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
+			child_info = %HexDisplay::from(&child_info.storage_key()),
+			storage_root = %HexDisplay::from(&root.as_ref()),
+			cached = %_cached,
+		);
 
-				root
-			} else {
-				// empty overlay
-				let root = self
-					.storage(prefixed_storage_key.as_slice())
-					.and_then(|k| Decode::decode(&mut &k[..]).ok())
-					// V1 is equivalent to V0 on empty root.
-					.unwrap_or_else(empty_child_trie_root::<LayoutV1<H>>);
-
-				trace!(
-					target: "state",
-					method = "ChildStorageRoot",
-					ext_id = %HexDisplay::from(&self.id.to_le_bytes()),
-					child_info = %HexDisplay::from(&storage_key),
-					storage_root = %HexDisplay::from(&root.as_ref()),
-					cached = false,
-				);
-
-				root.encode()
-			}
-		}
+		root.encode()
 	}
 
 	fn storage_index_transaction(&mut self, index: u32, hash: &[u8], size: u32) {
@@ -669,7 +585,6 @@ where
 	}
 
 	fn storage_rollback_transaction(&mut self) -> Result<(), ()> {
-		self.mark_dirty();
 		self.overlay.rollback_transaction().map_err(|_| ())
 	}
 
@@ -682,14 +597,9 @@ where
 			self.overlay.rollback_transaction().expect(BENCHMARKING_FN);
 		}
 		self.overlay
-			.drain_storage_changes(
-				self.backend,
-				self.storage_transaction_cache,
-				Default::default(), // using any state
-			)
+			.drain_storage_changes(self.backend, Default::default())
 			.expect(EXT_NOT_ALLOWED_TO_FAIL);
 		self.backend.wipe().expect(EXT_NOT_ALLOWED_TO_FAIL);
-		self.mark_dirty();
 		self.overlay
 			.enter_runtime()
 			.expect("We have reset the overlay above, so we can not be in the runtime; qed");
@@ -703,7 +613,7 @@ where
 		}
 		let changes = self
 			.overlay
-			.drain_storage_changes(self.backend, self.storage_transaction_cache, state_version)
+			.drain_storage_changes(self.backend, state_version)
 			.expect(EXT_NOT_ALLOWED_TO_FAIL);
 		self.backend
 			.commit(
@@ -713,7 +623,6 @@ where
 				changes.child_storage_changes,
 			)
 			.expect(EXT_NOT_ALLOWED_TO_FAIL);
-		self.mark_dirty();
 		self.overlay
 			.enter_runtime()
 			.expect("We have reset the overlay above, so we can not be in the runtime; qed");
@@ -914,7 +823,7 @@ where
 mod tests {
 	use super::*;
 	use crate::InMemoryBackend;
-	use codec::Encode;
+	use codec::{Decode, Encode};
 	use sp_core::{
 		map,
 		storage::{Storage, StorageChild},
@@ -926,7 +835,6 @@ mod tests {
 
 	#[test]
 	fn next_storage_key_works() {
-		let mut cache = StorageTransactionCache::default();
 		let mut overlay = OverlayedChanges::default();
 		overlay.set_storage(vec![20], None);
 		overlay.set_storage(vec![30], Some(vec![31]));
@@ -943,7 +851,7 @@ mod tests {
 		)
 			.into();
 
-		let ext = TestExt::new(&mut overlay, &mut cache, &backend, None);
+		let ext = TestExt::new(&mut overlay, &backend, None);
 
 		// next_backend < next_overlay
 		assert_eq!(ext.next_storage_key(&[5]), Some(vec![10]));
@@ -959,7 +867,7 @@ mod tests {
 
 		drop(ext);
 		overlay.set_storage(vec![50], Some(vec![50]));
-		let ext = TestExt::new(&mut overlay, &mut cache, &backend, None);
+		let ext = TestExt::new(&mut overlay, &backend, None);
 
 		// next_overlay exist but next_backend doesn't exist
 		assert_eq!(ext.next_storage_key(&[40]), Some(vec![50]));
@@ -967,7 +875,6 @@ mod tests {
 
 	#[test]
 	fn next_storage_key_works_with_a_lot_empty_values_in_overlay() {
-		let mut cache = StorageTransactionCache::default();
 		let mut overlay = OverlayedChanges::default();
 		overlay.set_storage(vec![20], None);
 		overlay.set_storage(vec![21], None);
@@ -990,7 +897,7 @@ mod tests {
 		)
 			.into();
 
-		let ext = TestExt::new(&mut overlay, &mut cache, &backend, None);
+		let ext = TestExt::new(&mut overlay, &backend, None);
 
 		assert_eq!(ext.next_storage_key(&[5]), Some(vec![30]));
 
@@ -1002,7 +909,6 @@ mod tests {
 		let child_info = ChildInfo::new_default(b"Child1");
 		let child_info = &child_info;
 
-		let mut cache = StorageTransactionCache::default();
 		let mut overlay = OverlayedChanges::default();
 		overlay.set_child_storage(child_info, vec![20], None);
 		overlay.set_child_storage(child_info, vec![30], Some(vec![31]));
@@ -1024,7 +930,7 @@ mod tests {
 		)
 			.into();
 
-		let ext = TestExt::new(&mut overlay, &mut cache, &backend, None);
+		let ext = TestExt::new(&mut overlay, &backend, None);
 
 		// next_backend < next_overlay
 		assert_eq!(ext.next_child_storage_key(child_info, &[5]), Some(vec![10]));
@@ -1040,7 +946,7 @@ mod tests {
 
 		drop(ext);
 		overlay.set_child_storage(child_info, vec![50], Some(vec![50]));
-		let ext = TestExt::new(&mut overlay, &mut cache, &backend, None);
+		let ext = TestExt::new(&mut overlay, &backend, None);
 
 		// next_overlay exist but next_backend doesn't exist
 		assert_eq!(ext.next_child_storage_key(child_info, &[40]), Some(vec![50]));
@@ -1050,7 +956,6 @@ mod tests {
 	fn child_storage_works() {
 		let child_info = ChildInfo::new_default(b"Child1");
 		let child_info = &child_info;
-		let mut cache = StorageTransactionCache::default();
 		let mut overlay = OverlayedChanges::default();
 		overlay.set_child_storage(child_info, vec![20], None);
 		overlay.set_child_storage(child_info, vec![30], Some(vec![31]));
@@ -1072,7 +977,7 @@ mod tests {
 		)
 			.into();
 
-		let ext = TestExt::new(&mut overlay, &mut cache, &backend, None);
+		let ext = TestExt::new(&mut overlay, &backend, None);
 
 		assert_eq!(ext.child_storage(child_info, &[10]), Some(vec![10]));
 		assert_eq!(
@@ -1094,7 +999,6 @@ mod tests {
 	fn clear_prefix_cannot_delete_a_child_root() {
 		let child_info = ChildInfo::new_default(b"Child1");
 		let child_info = &child_info;
-		let mut cache = StorageTransactionCache::default();
 		let mut overlay = OverlayedChanges::default();
 		let backend = (
 			Storage {
@@ -1112,7 +1016,7 @@ mod tests {
 		)
 			.into();
 
-		let ext = TestExt::new(&mut overlay, &mut cache, &backend, None);
+		let ext = TestExt::new(&mut overlay, &backend, None);
 
 		use sp_core::storage::well_known_keys;
 		let mut ext = ext;

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

@@ -24,36 +24,22 @@ use crate::{
 use codec::Codec;
 use hash_db::Hasher;
 use sp_core::storage::{ChildInfo, StateVersion, Storage};
-use sp_trie::{empty_trie_root, GenericMemoryDB, HashKey, KeyFunction, LayoutV1, MemoryDB};
+use sp_trie::{empty_trie_root, LayoutV1, PrefixedMemoryDB};
 use std::collections::{BTreeMap, HashMap};
 
 /// Create a new empty instance of in-memory backend.
-///
-/// It will use [`HashKey`] to store the keys internally.
-pub fn new_in_mem_hash_key<H>() -> TrieBackend<MemoryDB<H>, H>
+pub fn new_in_mem<H>() -> TrieBackend<PrefixedMemoryDB<H>, H>
 where
 	H: Hasher,
 	H::Out: Codec + Ord,
 {
-	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.
-	TrieBackendBuilder::new(db, empty_trie_root::<LayoutV1<H>>()).build()
+	TrieBackendBuilder::new(Default::default(), empty_trie_root::<LayoutV1<H>>()).build()
 }
 
-impl<H: Hasher, KF> TrieBackend<GenericMemoryDB<H, KF>, H>
+impl<H: Hasher> TrieBackend<PrefixedMemoryDB<H>, 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)>>(
@@ -85,15 +71,16 @@ where
 	}
 
 	/// Merge trie nodes into this backend.
-	pub fn update_backend(&self, root: H::Out, changes: GenericMemoryDB<H, KF>) -> Self {
+	pub fn update_backend(&self, root: H::Out, changes: PrefixedMemoryDB<H>) -> Self {
 		let mut clone = self.backend_storage().clone();
 		clone.consolidate(changes);
 		TrieBackendBuilder::new(clone, root).build()
 	}
 
 	/// 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: GenericMemoryDB<H, KF>) {
+	pub fn apply_transaction(&mut self, root: H::Out, transaction: PrefixedMemoryDB<H>) {
 		let mut storage = sp_std::mem::take(self).into_storage();
+
 		storage.consolidate(transaction);
 		*self = TrieBackendBuilder::new(storage, root).build();
 	}
@@ -104,33 +91,29 @@ where
 	}
 }
 
-impl<H: Hasher, KF> Clone for TrieBackend<GenericMemoryDB<H, KF>, H>
+impl<H: Hasher> Clone for TrieBackend<PrefixedMemoryDB<H>, H>
 where
 	H::Out: Codec + Ord,
-	KF: KeyFunction<H> + Send + Sync,
 {
 	fn clone(&self) -> Self {
 		TrieBackendBuilder::new(self.backend_storage().clone(), *self.root()).build()
 	}
 }
 
-impl<H, KF> Default for TrieBackend<GenericMemoryDB<H, KF>, H>
+impl<H> Default for TrieBackend<PrefixedMemoryDB<H>, H>
 where
 	H: Hasher,
 	H::Out: Codec + Ord,
-	KF: KeyFunction<H> + Send + Sync,
 {
 	fn default() -> Self {
 		new_in_mem()
 	}
 }
 
-impl<H: Hasher, KF>
-	From<(HashMap<Option<ChildInfo>, BTreeMap<StorageKey, StorageValue>>, StateVersion)>
-	for TrieBackend<GenericMemoryDB<H, KF>, H>
+impl<H: Hasher> From<(HashMap<Option<ChildInfo>, BTreeMap<StorageKey, StorageValue>>, StateVersion)>
+	for TrieBackend<PrefixedMemoryDB<H>, H>
 where
 	H::Out: Codec + Ord,
-	KF: KeyFunction<H> + Send + Sync,
 {
 	fn from(
 		(inner, state_version): (
@@ -149,10 +132,9 @@ where
 	}
 }
 
-impl<H: Hasher, KF> From<(Storage, StateVersion)> for TrieBackend<GenericMemoryDB<H, KF>, H>
+impl<H: Hasher> From<(Storage, StateVersion)> for TrieBackend<PrefixedMemoryDB<H>, 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
@@ -165,11 +147,10 @@ where
 	}
 }
 
-impl<H: Hasher, KF> From<(BTreeMap<StorageKey, StorageValue>, StateVersion)>
-	for TrieBackend<GenericMemoryDB<H, KF>, H>
+impl<H: Hasher> From<(BTreeMap<StorageKey, StorageValue>, StateVersion)>
+	for TrieBackend<PrefixedMemoryDB<H>, 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();
@@ -178,11 +159,10 @@ where
 	}
 }
 
-impl<H: Hasher, KF> From<(Vec<(Option<ChildInfo>, StorageCollection)>, StateVersion)>
-	for TrieBackend<GenericMemoryDB<H, KF>, H>
+impl<H: Hasher> From<(Vec<(Option<ChildInfo>, StorageCollection)>, StateVersion)>
+	for TrieBackend<PrefixedMemoryDB<H>, H>
 where
 	H::Out: Codec + Ord,
-	KF: KeyFunction<H> + Send + Sync,
 {
 	fn from(
 		(inner, state_version): (Vec<(Option<ChildInfo>, StorageCollection)>, StateVersion),
@@ -212,7 +192,7 @@ mod tests {
 	#[test]
 	fn in_memory_with_child_trie_only() {
 		let state_version = StateVersion::default();
-		let storage = new_in_mem_hash_key::<BlakeTwo256>();
+		let storage = new_in_mem::<BlakeTwo256>();
 		let child_info = ChildInfo::new_default(b"1");
 		let child_info = &child_info;
 		let storage = storage.update(
@@ -228,7 +208,7 @@ mod tests {
 	#[test]
 	fn insert_multiple_times_child_data_works() {
 		let state_version = StateVersion::default();
-		let mut storage = new_in_mem_hash_key::<BlakeTwo256>();
+		let mut storage = new_in_mem::<BlakeTwo256>();
 		let child_info = ChildInfo::new_default(b"1");
 
 		storage.insert(

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

@@ -125,13 +125,13 @@ impl sp_std::fmt::Display for DefaultError {
 }
 
 pub use crate::{
-	backend::{Backend, IterArgs, KeysIter, PairsIter, StorageIterator},
+	backend::{Backend, BackendTransaction, IterArgs, KeysIter, PairsIter, StorageIterator},
 	error::{Error, ExecutionError},
 	ext::Ext,
 	overlayed_changes::{
 		ChildStorageCollection, IndexOperation, OffchainChangesCollection,
 		OffchainOverlayedChanges, OverlayedChanges, StorageChanges, StorageCollection, StorageKey,
-		StorageTransactionCache, StorageValue,
+		StorageValue,
 	},
 	stats::{StateMachineStats, UsageInfo, UsageUnit},
 	trie_backend::{TrieBackend, TrieBackendBuilder},
@@ -143,7 +143,7 @@ mod std_reexport {
 	pub use crate::{
 		basic::BasicExternalities,
 		error::{Error, ExecutionError},
-		in_memory_backend::{new_in_mem, new_in_mem_hash_key},
+		in_memory_backend::new_in_mem,
 		read_only::{InspectState, ReadOnlyExternalities},
 		testing::TestExternalities,
 		trie_backend::create_proof_check_backend,
@@ -168,6 +168,7 @@ mod execution {
 		traits::{CallContext, CodeExecutor, RuntimeCode},
 	};
 	use sp_externalities::Extensions;
+	use sp_trie::PrefixedMemoryDB;
 	use std::collections::{HashMap, HashSet};
 
 	pub(crate) type CallResult<E> = Result<Vec<u8>, E>;
@@ -176,7 +177,7 @@ mod execution {
 	pub type DefaultHandler<E> = fn(CallResult<E>, CallResult<E>) -> CallResult<E>;
 
 	/// Trie backend with in-memory storage.
-	pub type InMemoryBackend<H> = TrieBackend<MemoryDB<H>, H>;
+	pub type InMemoryBackend<H> = TrieBackend<PrefixedMemoryDB<H>, H>;
 
 	/// Storage backend trust level.
 	#[derive(Debug, Clone)]
@@ -199,9 +200,8 @@ mod execution {
 		exec: &'a Exec,
 		method: &'a str,
 		call_data: &'a [u8],
-		overlay: &'a mut OverlayedChanges,
+		overlay: &'a mut OverlayedChanges<H>,
 		extensions: &'a mut Extensions,
-		storage_transaction_cache: Option<&'a mut StorageTransactionCache<B::Transaction, H>>,
 		runtime_code: &'a RuntimeCode<'a>,
 		stats: StateMachineStats,
 		/// The hash of the block the state machine will be executed on.
@@ -231,7 +231,7 @@ mod execution {
 		/// Creates new substrate state machine.
 		pub fn new(
 			backend: &'a B,
-			overlay: &'a mut OverlayedChanges,
+			overlay: &'a mut OverlayedChanges<H>,
 			exec: &'a Exec,
 			method: &'a str,
 			call_data: &'a [u8],
@@ -246,7 +246,6 @@ mod execution {
 				call_data,
 				extensions,
 				overlay,
-				storage_transaction_cache: None,
 				runtime_code,
 				stats: StateMachineStats::default(),
 				parent_hash: None,
@@ -254,19 +253,6 @@ mod execution {
 			}
 		}
 
-		/// Use given `cache` as storage transaction cache.
-		///
-		/// The cache will be used to cache storage transactions that can be build while executing a
-		/// function in the runtime. For example, when calculating the storage root a transaction is
-		/// build that will be cached.
-		pub fn with_storage_transaction_cache(
-			mut self,
-			cache: Option<&'a mut StorageTransactionCache<B::Transaction, H>>,
-		) -> Self {
-			self.storage_transaction_cache = cache;
-			self
-		}
-
 		/// Set the given `parent_hash` as the hash of the parent block.
 		///
 		/// This will be used for improved logging.
@@ -284,18 +270,11 @@ mod execution {
 		///
 		/// Returns the SCALE encoded result of the executed function.
 		pub fn execute(&mut self) -> Result<Vec<u8>, Box<dyn Error>> {
-			let mut cache = StorageTransactionCache::default();
-
-			let cache = match self.storage_transaction_cache.as_mut() {
-				Some(cache) => cache,
-				None => &mut cache,
-			};
-
 			self.overlay
 				.enter_runtime()
 				.expect("StateMachine is never called from the runtime; qed");
 
-			let mut ext = Ext::new(self.overlay, cache, self.backend, Some(self.extensions));
+			let mut ext = Ext::new(self.overlay, self.backend, Some(self.extensions));
 
 			let ext_id = ext.id;
 
@@ -331,7 +310,7 @@ mod execution {
 	/// Prove execution using the given state backend, overlayed changes, and call executor.
 	pub fn prove_execution<B, H, Exec>(
 		backend: &mut B,
-		overlay: &mut OverlayedChanges,
+		overlay: &mut OverlayedChanges<H>,
 		exec: &Exec,
 		method: &str,
 		call_data: &[u8],
@@ -366,7 +345,7 @@ mod execution {
 	/// blocks (e.g. a transaction at a time), ensure a different method is used.
 	pub fn prove_execution_on_trie_backend<S, H, Exec>(
 		trie_backend: &TrieBackend<S, H>,
-		overlay: &mut OverlayedChanges,
+		overlay: &mut OverlayedChanges<H>,
 		exec: &Exec,
 		method: &str,
 		call_data: &[u8],
@@ -405,7 +384,7 @@ mod execution {
 	pub fn execution_proof_check<H, Exec>(
 		root: H::Out,
 		proof: StorageProof,
-		overlay: &mut OverlayedChanges,
+		overlay: &mut OverlayedChanges<H>,
 		exec: &Exec,
 		method: &str,
 		call_data: &[u8],
@@ -430,7 +409,7 @@ mod execution {
 	/// Check execution proof on proving backend, generated by `prove_execution` call.
 	pub fn execution_proof_check_on_trie_backend<H, Exec>(
 		trie_backend: &TrieBackend<MemoryDB<H>, H>,
-		overlay: &mut OverlayedChanges,
+		overlay: &mut OverlayedChanges<H>,
 		exec: &Exec,
 		method: &str,
 		call_data: &[u8],
@@ -1109,7 +1088,7 @@ mod execution {
 #[cfg(test)]
 mod tests {
 	use super::{backend::AsTrieBackend, ext::Ext, *};
-	use crate::{execution::CallResult, in_memory_backend::new_in_mem_hash_key};
+	use crate::{execution::CallResult, in_memory_backend::new_in_mem};
 	use assert_matches::assert_matches;
 	use codec::Encode;
 	use sp_core::{
@@ -1287,8 +1266,7 @@ mod tests {
 
 		let overlay_limit = overlay.clone();
 		{
-			let mut cache = StorageTransactionCache::default();
-			let mut ext = Ext::new(&mut overlay, &mut cache, backend, None);
+			let mut ext = Ext::new(&mut overlay, backend, None);
 			let _ = ext.clear_prefix(b"ab", None, None);
 		}
 		overlay.commit_transaction().unwrap();
@@ -1311,8 +1289,7 @@ mod tests {
 
 		let mut overlay = overlay_limit;
 		{
-			let mut cache = StorageTransactionCache::default();
-			let mut ext = Ext::new(&mut overlay, &mut cache, backend, None);
+			let mut ext = Ext::new(&mut overlay, backend, None);
 			assert_matches!(
 				ext.clear_prefix(b"ab", Some(1), None).deconstruct(),
 				(Some(_), 1, 3, 1)
@@ -1356,8 +1333,7 @@ mod tests {
 		overlay.set_child_storage(&child_info, b"4".to_vec(), Some(b"1312".to_vec()));
 
 		{
-			let mut cache = StorageTransactionCache::default();
-			let mut ext = Ext::new(&mut overlay, &mut cache, &backend, None);
+			let mut ext = Ext::new(&mut overlay, &backend, None);
 			let r = ext.kill_child_storage(&child_info, Some(2), None);
 			assert_matches!(r.deconstruct(), (Some(_), 2, 6, 2));
 		}
@@ -1392,8 +1368,7 @@ mod tests {
 		];
 		let backend = InMemoryBackend::<BlakeTwo256>::from((initial, StateVersion::default()));
 		let mut overlay = OverlayedChanges::default();
-		let mut cache = StorageTransactionCache::default();
-		let mut ext = Ext::new(&mut overlay, &mut cache, &backend, None);
+		let mut ext = Ext::new(&mut overlay, &backend, None);
 		let r = ext.kill_child_storage(&child_info, Some(0), None).deconstruct();
 		assert_matches!(r, (Some(_), 0, 0, 0));
 		let r = ext
@@ -1422,8 +1397,7 @@ mod tests {
 		];
 		let backend = InMemoryBackend::<BlakeTwo256>::from((initial, StateVersion::default()));
 		let mut overlay = OverlayedChanges::default();
-		let mut cache = StorageTransactionCache::default();
-		let mut ext = Ext::new(&mut overlay, &mut cache, &backend, None);
+		let mut ext = Ext::new(&mut overlay, &backend, None);
 		assert_eq!(ext.kill_child_storage(&child_info, None, None).deconstruct(), (None, 4, 4, 4));
 	}
 
@@ -1431,11 +1405,10 @@ 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_hash_key::<BlakeTwo256>();
+		let state = new_in_mem::<BlakeTwo256>();
 		let backend = state.as_trie_backend();
 		let mut overlay = OverlayedChanges::default();
-		let mut cache = StorageTransactionCache::default();
-		let mut ext = Ext::new(&mut overlay, &mut cache, backend, None);
+		let mut ext = Ext::new(&mut overlay, backend, None);
 
 		ext.set_child_storage(child_info, b"abc".to_vec(), b"def".to_vec());
 		assert_eq!(ext.child_storage(child_info, b"abc"), Some(b"def".to_vec()));
@@ -1447,19 +1420,18 @@ 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_hash_key::<BlakeTwo256>();
+		let state = new_in_mem::<BlakeTwo256>();
 		let backend = state.as_trie_backend();
 		let mut overlay = OverlayedChanges::default();
-		let mut cache = StorageTransactionCache::default();
 		{
-			let mut ext = Ext::new(&mut overlay, &mut cache, backend, None);
+			let mut ext = Ext::new(&mut overlay, backend, None);
 
 			ext.storage_append(key.clone(), reference_data[0].encode());
 			assert_eq!(ext.storage(key.as_slice()), Some(vec![reference_data[0].clone()].encode()));
 		}
 		overlay.start_transaction();
 		{
-			let mut ext = Ext::new(&mut overlay, &mut cache, backend, None);
+			let mut ext = Ext::new(&mut overlay, backend, None);
 
 			for i in reference_data.iter().skip(1) {
 				ext.storage_append(key.clone(), i.encode());
@@ -1468,7 +1440,7 @@ mod tests {
 		}
 		overlay.rollback_transaction().unwrap();
 		{
-			let ext = Ext::new(&mut overlay, &mut cache, backend, None);
+			let ext = Ext::new(&mut overlay, backend, None);
 			assert_eq!(ext.storage(key.as_slice()), Some(vec![reference_data[0].clone()].encode()));
 		}
 	}
@@ -1483,14 +1455,13 @@ mod tests {
 		}
 
 		let key = b"events".to_vec();
-		let mut cache = StorageTransactionCache::default();
-		let state = new_in_mem_hash_key::<BlakeTwo256>();
+		let state = new_in_mem::<BlakeTwo256>();
 		let backend = state.as_trie_backend();
 		let mut overlay = OverlayedChanges::default();
 
 		// For example, block initialization with event.
 		{
-			let mut ext = Ext::new(&mut overlay, &mut cache, backend, None);
+			let mut ext = Ext::new(&mut overlay, backend, None);
 			ext.clear_storage(key.as_slice());
 			ext.storage_append(key.clone(), Item::InitializationItem.encode());
 		}
@@ -1498,7 +1469,7 @@ mod tests {
 
 		// For example, first transaction resulted in panic during block building
 		{
-			let mut ext = Ext::new(&mut overlay, &mut cache, backend, None);
+			let mut ext = Ext::new(&mut overlay, backend, None);
 
 			assert_eq!(ext.storage(key.as_slice()), Some(vec![Item::InitializationItem].encode()));
 
@@ -1513,7 +1484,7 @@ mod tests {
 
 		// Then we apply next transaction which is valid this time.
 		{
-			let mut ext = Ext::new(&mut overlay, &mut cache, backend, None);
+			let mut ext = Ext::new(&mut overlay, backend, None);
 
 			assert_eq!(ext.storage(key.as_slice()), Some(vec![Item::InitializationItem].encode()));
 
@@ -1528,7 +1499,7 @@ mod tests {
 
 		// Then only initlaization item and second (committed) item should persist.
 		{
-			let ext = Ext::new(&mut overlay, &mut cache, backend, None);
+			let ext = Ext::new(&mut overlay, backend, None);
 			assert_eq!(
 				ext.storage(key.as_slice()),
 				Some(vec![Item::InitializationItem, Item::CommitedItem].encode()),
@@ -1945,12 +1916,11 @@ mod tests {
 
 		let mut transaction = {
 			let backend = test_trie(state_version, None, None);
-			let mut cache = StorageTransactionCache::default();
-			let mut ext = Ext::new(&mut overlay, &mut cache, &backend, None);
+			let mut ext = Ext::new(&mut overlay, &backend, None);
 			ext.set_child_storage(&child_info_1, b"abc".to_vec(), b"def".to_vec());
 			ext.set_child_storage(&child_info_2, b"abc".to_vec(), b"def".to_vec());
 			ext.storage_root(state_version);
-			cache.transaction.unwrap()
+			overlay.drain_storage_changes(&backend, state_version).unwrap().transaction
 		};
 		let mut duplicate = false;
 		for (k, (value, rc)) in transaction.drain().iter() {
@@ -1982,8 +1952,7 @@ mod tests {
 		assert_eq!(overlay.storage(b"bbb"), None);
 
 		{
-			let mut cache = StorageTransactionCache::default();
-			let mut ext = Ext::new(&mut overlay, &mut cache, backend, None);
+			let mut ext = Ext::new(&mut overlay, backend, None);
 			assert_eq!(ext.storage(b"bbb"), Some(vec![]));
 			assert_eq!(ext.storage(b"ccc"), Some(vec![]));
 			ext.clear_storage(b"ccc");

substrate/primitives/state-machine/src/overlayed_changes/mod.rs

@@ -21,7 +21,7 @@ mod changeset;
 mod offchain;
 
 use self::changeset::OverlayedChangeSet;
-use crate::{backend::Backend, stats::StateMachineStats, DefaultError};
+use crate::{backend::Backend, stats::StateMachineStats, BackendTransaction, DefaultError};
 use codec::{Decode, Encode};
 use hash_db::Hasher;
 pub use offchain::OffchainOverlayedChanges;
@@ -34,6 +34,7 @@ use sp_externalities::{Extension, Extensions};
 #[cfg(not(feature = "std"))]
 use sp_std::collections::btree_map::BTreeMap as Map;
 use sp_std::{collections::btree_set::BTreeSet, vec::Vec};
+use sp_trie::{empty_child_trie_root, LayoutV1};
 #[cfg(feature = "std")]
 use std::collections::{hash_map::Entry as MapEntry, HashMap as Map};
 #[cfg(feature = "std")]
@@ -88,8 +89,7 @@ impl Extrinsics {
 /// The set of changes that are overlaid onto the backend.
 ///
 /// It allows changes to be modified using nestable transactions.
-#[derive(Debug, Default, Clone)]
-pub struct OverlayedChanges {
+pub struct OverlayedChanges<H: Hasher> {
 	/// Top level storage changes.
 	top: OverlayedChangeSet,
 	/// Child storage changes. The map key is the child storage key without the common prefix.
@@ -102,6 +102,52 @@ pub struct OverlayedChanges {
 	collect_extrinsics: bool,
 	/// Collect statistic on this execution.
 	stats: StateMachineStats,
+	/// Caches the "storage transaction" that is created while calling `storage_root`.
+	///
+	/// This transaction can be applied to the backend to persist the state changes.
+	storage_transaction_cache: Option<StorageTransactionCache<H>>,
+}
+
+impl<H: Hasher> Default for OverlayedChanges<H> {
+	fn default() -> Self {
+		Self {
+			top: Default::default(),
+			children: Default::default(),
+			offchain: Default::default(),
+			transaction_index_ops: Default::default(),
+			collect_extrinsics: Default::default(),
+			stats: Default::default(),
+			storage_transaction_cache: None,
+		}
+	}
+}
+
+impl<H: Hasher> Clone for OverlayedChanges<H> {
+	fn clone(&self) -> Self {
+		Self {
+			top: self.top.clone(),
+			children: self.children.clone(),
+			offchain: self.offchain.clone(),
+			transaction_index_ops: self.transaction_index_ops.clone(),
+			collect_extrinsics: self.collect_extrinsics,
+			stats: self.stats.clone(),
+			storage_transaction_cache: self.storage_transaction_cache.clone(),
+		}
+	}
+}
+
+impl<H: Hasher> sp_std::fmt::Debug for OverlayedChanges<H> {
+	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+		f.debug_struct("OverlayedChanges")
+			.field("top", &self.top)
+			.field("children", &self.children)
+			.field("offchain", &self.offchain)
+			.field("transaction_index_ops", &self.transaction_index_ops)
+			.field("collect_extrinsics", &self.collect_extrinsics)
+			.field("stats", &self.stats)
+			.field("storage_transaction_cache", &self.storage_transaction_cache)
+			.finish()
+	}
 }
 
 /// Transaction index operation.
@@ -129,7 +175,7 @@ pub enum IndexOperation {
 ///
 /// This contains all the changes to the storage and transactions to apply theses changes to the
 /// backend.
-pub struct StorageChanges<Transaction, H: Hasher> {
+pub struct StorageChanges<H: Hasher> {
 	/// All changes to the main storage.
 	///
 	/// A value of `None` means that it was deleted.
@@ -142,7 +188,7 @@ pub struct StorageChanges<Transaction, H: Hasher> {
 	/// [`main_storage_changes`](StorageChanges::main_storage_changes) and from
 	/// [`child_storage_changes`](StorageChanges::child_storage_changes).
 	/// [`offchain_storage_changes`](StorageChanges::offchain_storage_changes).
-	pub transaction: Transaction,
+	pub transaction: BackendTransaction<H>,
 	/// The storage root after applying the transaction.
 	pub transaction_storage_root: H::Out,
 	/// Changes to the transaction index,
@@ -151,7 +197,7 @@ pub struct StorageChanges<Transaction, H: Hasher> {
 }
 
 #[cfg(feature = "std")]
-impl<Transaction, H: Hasher> StorageChanges<Transaction, H> {
+impl<H: Hasher> StorageChanges<H> {
 	/// Deconstruct into the inner values
 	pub fn into_inner(
 		self,
@@ -159,7 +205,7 @@ impl<Transaction, H: Hasher> StorageChanges<Transaction, H> {
 		StorageCollection,
 		ChildStorageCollection,
 		OffchainChangesCollection,
-		Transaction,
+		BackendTransaction<H>,
 		H::Out,
 		Vec<IndexOperation>,
 	) {
@@ -174,30 +220,7 @@ impl<Transaction, H: Hasher> StorageChanges<Transaction, H> {
 	}
 }
 
-/// Storage transactions are calculated as part of the `storage_root`.
-/// These transactions can be reused for importing the block into the
-/// storage. So, we cache them to not require a recomputation of those transactions.
-pub struct StorageTransactionCache<Transaction, H: Hasher> {
-	/// Contains the changes for the main and the child storages as one transaction.
-	pub(crate) transaction: Option<Transaction>,
-	/// The storage root after applying the transaction.
-	pub(crate) transaction_storage_root: Option<H::Out>,
-}
-
-impl<Transaction, H: Hasher> StorageTransactionCache<Transaction, H> {
-	/// Reset the cached transactions.
-	pub fn reset(&mut self) {
-		*self = Self::default();
-	}
-}
-
-impl<Transaction, H: Hasher> Default for StorageTransactionCache<Transaction, H> {
-	fn default() -> Self {
-		Self { transaction: None, transaction_storage_root: None }
-	}
-}
-
-impl<Transaction: Default, H: Hasher> Default for StorageChanges<Transaction, H> {
+impl<H: Hasher> Default for StorageChanges<H> {
 	fn default() -> Self {
 		Self {
 			main_storage_changes: Default::default(),
@@ -211,7 +234,46 @@ impl<Transaction: Default, H: Hasher> Default for StorageChanges<Transaction, H>
 	}
 }
 
-impl OverlayedChanges {
+/// Storage transactions are calculated as part of the `storage_root`.
+/// These transactions can be reused for importing the block into the
+/// storage. So, we cache them to not require a recomputation of those transactions.
+struct StorageTransactionCache<H: Hasher> {
+	/// Contains the changes for the main and the child storages as one transaction.
+	transaction: BackendTransaction<H>,
+	/// The storage root after applying the transaction.
+	transaction_storage_root: H::Out,
+}
+
+impl<H: Hasher> StorageTransactionCache<H> {
+	fn into_inner(self) -> (BackendTransaction<H>, H::Out) {
+		(self.transaction, self.transaction_storage_root)
+	}
+}
+
+impl<H: Hasher> Clone for StorageTransactionCache<H> {
+	fn clone(&self) -> Self {
+		Self {
+			transaction: self.transaction.clone(),
+			transaction_storage_root: self.transaction_storage_root,
+		}
+	}
+}
+
+impl<H: Hasher> sp_std::fmt::Debug for StorageTransactionCache<H> {
+	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+		let mut debug = f.debug_struct("StorageTransactionCache");
+
+		#[cfg(feature = "std")]
+		debug.field("transaction_storage_root", &self.transaction_storage_root);
+
+		#[cfg(not(feature = "std"))]
+		debug.field("transaction_storage_root", &self.transaction_storage_root.as_ref());
+
+		debug.finish()
+	}
+}
+
+impl<H: Hasher> OverlayedChanges<H> {
 	/// Whether no changes are contained in the top nor in any of the child changes.
 	pub fn is_empty(&self) -> bool {
 		self.top.is_empty() && self.children.is_empty()
@@ -234,6 +296,12 @@ impl OverlayedChanges {
 		})
 	}
 
+	/// Should be called when there are changes that require to reset the
+	/// `storage_transaction_cache`.
+	fn mark_dirty(&mut self) {
+		self.storage_transaction_cache = None;
+	}
+
 	/// Returns mutable reference to current value.
 	/// If there is no value in the overlay, the given callback is used to initiate the value.
 	/// Warning this function registers a change, so the mutable reference MUST be modified.
@@ -245,6 +313,8 @@ impl OverlayedChanges {
 		key: &[u8],
 		init: impl Fn() -> StorageValue,
 	) -> &mut StorageValue {
+		self.mark_dirty();
+
 		let value = self.top.modify(key.to_vec(), init, self.extrinsic_index());
 
 		// if the value was deleted initialise it back with an empty vec
@@ -266,6 +336,8 @@ impl OverlayedChanges {
 	///
 	/// Can be rolled back or committed when called inside a transaction.
 	pub fn set_storage(&mut self, key: StorageKey, val: Option<StorageValue>) {
+		self.mark_dirty();
+
 		let size_write = val.as_ref().map(|x| x.len() as u64).unwrap_or(0);
 		self.stats.tally_write_overlay(size_write);
 		self.top.set(key, val, self.extrinsic_index());
@@ -282,6 +354,8 @@ impl OverlayedChanges {
 		key: StorageKey,
 		val: Option<StorageValue>,
 	) {
+		self.mark_dirty();
+
 		let extrinsic_index = self.extrinsic_index();
 		let size_write = val.as_ref().map(|x| x.len() as u64).unwrap_or(0);
 		self.stats.tally_write_overlay(size_write);
@@ -300,6 +374,8 @@ impl OverlayedChanges {
 	///
 	/// Can be rolled back or committed when called inside a transaction.
 	pub(crate) fn clear_child_storage(&mut self, child_info: &ChildInfo) -> u32 {
+		self.mark_dirty();
+
 		let extrinsic_index = self.extrinsic_index();
 		let storage_key = child_info.storage_key().to_vec();
 		let top = &self.top;
@@ -316,6 +392,8 @@ impl OverlayedChanges {
 	///
 	/// Can be rolled back or committed when called inside a transaction.
 	pub(crate) fn clear_prefix(&mut self, prefix: &[u8]) -> u32 {
+		self.mark_dirty();
+
 		self.top.clear_where(|key, _| key.starts_with(prefix), self.extrinsic_index())
 	}
 
@@ -323,6 +401,8 @@ impl OverlayedChanges {
 	///
 	/// Can be rolled back or committed when called inside a transaction
 	pub(crate) fn clear_child_prefix(&mut self, child_info: &ChildInfo, prefix: &[u8]) -> u32 {
+		self.mark_dirty();
+
 		let extrinsic_index = self.extrinsic_index();
 		let storage_key = child_info.storage_key().to_vec();
 		let top = &self.top;
@@ -364,6 +444,8 @@ impl OverlayedChanges {
 	/// Any changes made during that transaction are discarded. Returns an error if
 	/// there is no open transaction that can be rolled back.
 	pub fn rollback_transaction(&mut self) -> Result<(), NoOpenTransaction> {
+		self.mark_dirty();
+
 		self.top.rollback_transaction()?;
 		retain_map(&mut self.children, |_, (changeset, _)| {
 			changeset
@@ -432,32 +514,6 @@ impl OverlayedChanges {
 		Ok(())
 	}
 
-	/// Consume all changes (top + children) and return them.
-	///
-	/// After calling this function no more changes are contained in this changeset.
-	///
-	/// Panics:
-	/// Panics if `transaction_depth() > 0`
-	fn drain_committed(
-		&mut self,
-	) -> (
-		impl Iterator<Item = (StorageKey, Option<StorageValue>)>,
-		impl Iterator<
-			Item = (
-				StorageKey,
-				(impl Iterator<Item = (StorageKey, Option<StorageValue>)>, ChildInfo),
-			),
-		>,
-	) {
-		use sp_std::mem::take;
-		(
-			take(&mut self.top).drain_commited(),
-			take(&mut self.children)
-				.into_iter()
-				.map(|(key, (val, info))| (key, (val.drain_commited(), info))),
-		)
-	}
-
 	/// Consume all changes (top + children) and return them.
 	///
 	/// After calling this function no more changes are contained in this changeset.
@@ -495,42 +551,33 @@ impl OverlayedChanges {
 		&self.transaction_index_ops
 	}
 
-	/// Convert this instance with all changes into a [`StorageChanges`] instance.
-	#[cfg(feature = "std")]
-	pub fn into_storage_changes<B: Backend<H>, H: Hasher>(
-		mut self,
-		backend: &B,
-		mut cache: StorageTransactionCache<B::Transaction, H>,
-		state_version: StateVersion,
-	) -> Result<StorageChanges<B::Transaction, H>, DefaultError>
-	where
-		H::Out: Ord + Encode + 'static,
-	{
-		self.drain_storage_changes(backend, &mut cache, state_version)
-	}
-
 	/// Drain all changes into a [`StorageChanges`] instance. Leave empty overlay in place.
-	pub fn drain_storage_changes<B: Backend<H>, H: Hasher>(
+	pub fn drain_storage_changes<B: Backend<H>>(
 		&mut self,
 		backend: &B,
-		cache: &mut StorageTransactionCache<B::Transaction, H>,
 		state_version: StateVersion,
-	) -> Result<StorageChanges<B::Transaction, H>, DefaultError>
+	) -> Result<StorageChanges<H>, DefaultError>
 	where
 		H::Out: Ord + Encode + 'static,
 	{
-		// If the transaction does not exist, we generate it.
-		if cache.transaction.is_none() {
-			self.storage_root(backend, cache, state_version);
-		}
+		let (transaction, transaction_storage_root) = match self.storage_transaction_cache.take() {
+			Some(cache) => cache.into_inner(),
+			// If the transaction does not exist, we generate it.
+			None => {
+				self.storage_root(backend, state_version);
+				self.storage_transaction_cache
+					.take()
+					.expect("`storage_transaction_cache` was just initialized; qed")
+					.into_inner()
+			},
+		};
 
-		let (transaction, transaction_storage_root) = cache
-			.transaction
-			.take()
-			.and_then(|t| cache.transaction_storage_root.take().map(|tr| (t, tr)))
-			.expect("Transaction was be generated as part of `storage_root`; qed");
+		use sp_std::mem::take;
+		let main_storage_changes = take(&mut self.top).drain_commited();
+		let child_storage_changes = take(&mut self.children)
+			.into_iter()
+			.map(|(key, (val, info))| (key, (val.drain_commited(), info)));
 
-		let (main_storage_changes, child_storage_changes) = self.drain_committed();
 		let offchain_storage_changes = self.offchain_drain_committed().collect();
 
 		#[cfg(feature = "std")]
@@ -562,29 +609,29 @@ impl OverlayedChanges {
 	/// Changes that are made outside of extrinsics, are marked with
 	/// `NO_EXTRINSIC_INDEX` index.
 	fn extrinsic_index(&self) -> Option<u32> {
-		match self.collect_extrinsics {
-			true => Some(
-				self.storage(EXTRINSIC_INDEX)
-					.and_then(|idx| idx.and_then(|idx| Decode::decode(&mut &*idx).ok()))
-					.unwrap_or(NO_EXTRINSIC_INDEX),
-			),
-			false => None,
-		}
+		self.collect_extrinsics.then(|| {
+			self.storage(EXTRINSIC_INDEX)
+				.and_then(|idx| idx.and_then(|idx| Decode::decode(&mut &*idx).ok()))
+				.unwrap_or(NO_EXTRINSIC_INDEX)
+		})
 	}
 
 	/// Generate the storage root using `backend` and all changes
 	/// as seen by the current transaction.
 	///
-	/// Returns the storage root and caches storage transaction in the given `cache`.
-	pub fn storage_root<H: Hasher, B: Backend<H>>(
-		&self,
+	/// Returns the storage root and whether it was already cached.
+	pub fn storage_root<B: Backend<H>>(
+		&mut self,
 		backend: &B,
-		cache: &mut StorageTransactionCache<B::Transaction, H>,
 		state_version: StateVersion,
-	) -> H::Out
+	) -> (H::Out, bool)
 	where
 		H::Out: Ord + Encode,
 	{
+		if let Some(cache) = &self.storage_transaction_cache {
+			return (cache.transaction_storage_root, true)
+		}
+
 		let delta = self.changes().map(|(k, v)| (&k[..], v.value().map(|v| &v[..])));
 		let child_delta = self.children().map(|(changes, info)| {
 			(info, changes.map(|(k, v)| (&k[..], v.value().map(|v| &v[..]))))
@@ -592,10 +639,72 @@ impl OverlayedChanges {
 
 		let (root, transaction) = backend.full_storage_root(delta, child_delta, state_version);
 
-		cache.transaction = Some(transaction);
-		cache.transaction_storage_root = Some(root);
+		self.storage_transaction_cache =
+			Some(StorageTransactionCache { transaction, transaction_storage_root: root });
 
-		root
+		(root, false)
+	}
+
+	/// Generate the child storage root using `backend` and all child changes
+	/// as seen by the current transaction.
+	///
+	/// Returns the child storage root and whether it was already cached.
+	pub fn child_storage_root<B: Backend<H>>(
+		&mut self,
+		child_info: &ChildInfo,
+		backend: &B,
+		state_version: StateVersion,
+	) -> Result<(H::Out, bool), B::Error>
+	where
+		H::Out: Ord + Encode + Decode,
+	{
+		let storage_key = child_info.storage_key();
+		let prefixed_storage_key = child_info.prefixed_storage_key();
+
+		if self.storage_transaction_cache.is_some() {
+			let root = self
+				.storage(prefixed_storage_key.as_slice())
+				.map(|v| Ok(v.map(|v| v.to_vec())))
+				.or_else(|| backend.storage(prefixed_storage_key.as_slice()).map(Some).transpose())
+				.transpose()?
+				.flatten()
+				.and_then(|k| Decode::decode(&mut &k[..]).ok())
+				// V1 is equivalent to V0 on empty root.
+				.unwrap_or_else(empty_child_trie_root::<LayoutV1<H>>);
+
+			return Ok((root, true))
+		}
+
+		let root = if let Some((changes, info)) = self.child_changes(storage_key) {
+			let delta = changes.map(|(k, v)| (k.as_ref(), v.value().map(AsRef::as_ref)));
+			Some(backend.child_storage_root(info, delta, state_version))
+		} else {
+			None
+		};
+
+		let root = if let Some((root, is_empty, _)) = root {
+			// We store update in the overlay in order to be able to use
+			// 'self.storage_transaction' cache. This is brittle as it rely on Ext only querying
+			// the trie backend for storage root.
+			// A better design would be to manage 'child_storage_transaction' in a
+			// similar way as 'storage_transaction' but for each child trie.
+			self.set_storage(prefixed_storage_key.into_inner(), (!is_empty).then(|| root.encode()));
+
+			self.mark_dirty();
+
+			root
+		} else {
+			// empty overlay
+			let root = backend
+				.storage(prefixed_storage_key.as_slice())?
+				.and_then(|k| Decode::decode(&mut &k[..]).ok())
+				// V1 is equivalent to V0 on empty root.
+				.unwrap_or_else(empty_child_trie_root::<LayoutV1<H>>);
+
+			root
+		};
+
+		Ok((root, false))
 	}
 
 	/// Returns an iterator over the keys (in lexicographic order) following `key` (excluding `key`)
@@ -639,7 +748,7 @@ impl OverlayedChanges {
 }
 
 #[cfg(feature = "std")]
-impl From<sp_core::storage::Storage> for OverlayedChanges {
+impl<H: Hasher> From<sp_core::storage::Storage> for OverlayedChanges<H> {
 	fn from(storage: sp_core::storage::Storage) -> Self {
 		Self {
 			top: storage.top.into(),
@@ -742,7 +851,8 @@ impl<'a> OverlayedExtensions<'a> {
 #[cfg(test)]
 mod tests {
 	use super::*;
-	use crate::{ext::Ext, InMemoryBackend};
+	use crate::{ext::Ext, new_in_mem, InMemoryBackend};
+	use array_bytes::bytes2hex;
 	use sp_core::{traits::Externalities, Blake2Hasher};
 	use std::collections::BTreeMap;
 
@@ -755,7 +865,7 @@ mod tests {
 
 	#[test]
 	fn overlayed_storage_works() {
-		let mut overlayed = OverlayedChanges::default();
+		let mut overlayed = OverlayedChanges::<Blake2Hasher>::default();
 
 		let key = vec![42, 69, 169, 142];
 
@@ -790,7 +900,7 @@ mod tests {
 	fn offchain_overlayed_storage_transactions_works() {
 		use sp_core::offchain::STORAGE_PREFIX;
 		fn check_offchain_content(
-			state: &OverlayedChanges,
+			state: &OverlayedChanges<Blake2Hasher>,
 			nb_commit: usize,
 			expected: Vec<(Vec<u8>, Option<Vec<u8>>)>,
 		) {
@@ -867,18 +977,61 @@ mod tests {
 		overlay.set_storage(b"dogglesworth".to_vec(), Some(b"cat".to_vec()));
 		overlay.set_storage(b"doug".to_vec(), None);
 
-		let mut cache = StorageTransactionCache::default();
-		let mut ext = Ext::new(&mut overlay, &mut cache, &backend, None);
-		let root = array_bytes::hex2bytes_unchecked(
-			"39245109cef3758c2eed2ccba8d9b370a917850af3824bc8348d505df2c298fa",
-		);
+		{
+			let mut ext = Ext::new(&mut overlay, &backend, None);
+			let root = "39245109cef3758c2eed2ccba8d9b370a917850af3824bc8348d505df2c298fa";
 
-		assert_eq!(&ext.storage_root(state_version)[..], &root);
+			assert_eq!(bytes2hex("", &ext.storage_root(state_version)), root);
+			// Calling a second time should use it from the cache
+			assert_eq!(bytes2hex("", &ext.storage_root(state_version)), root);
+		}
+
+		// Check that the storage root is recalculated
+		overlay.set_storage(b"doug2".to_vec(), Some(b"yes".to_vec()));
+
+		let mut ext = Ext::new(&mut overlay, &backend, None);
+		let root = "5c0a4e35cb967de785e1cb8743e6f24b6ff6d45155317f2078f6eb3fc4ff3e3d";
+		assert_eq!(bytes2hex("", &ext.storage_root(state_version)), root);
+	}
+
+	#[test]
+	fn overlayed_child_storage_root_works() {
+		let state_version = StateVersion::default();
+		let child_info = ChildInfo::new_default(b"Child1");
+		let child_info = &child_info;
+		let backend = new_in_mem::<Blake2Hasher>();
+		let mut overlay = OverlayedChanges::<Blake2Hasher>::default();
+		overlay.start_transaction();
+		overlay.set_child_storage(child_info, vec![20], Some(vec![20]));
+		overlay.set_child_storage(child_info, vec![30], Some(vec![30]));
+		overlay.set_child_storage(child_info, vec![40], Some(vec![40]));
+		overlay.commit_transaction().unwrap();
+		overlay.set_child_storage(child_info, vec![10], Some(vec![10]));
+		overlay.set_child_storage(child_info, vec![30], None);
+
+		{
+			let mut ext = Ext::new(&mut overlay, &backend, None);
+			let child_root = "c02965e1df4dc5baf6977390ce67dab1d7a9b27a87c1afe27b50d29cc990e0f5";
+			let root = "eafb765909c3ed5afd92a0c564acf4620d0234b31702e8e8e9b48da72a748838";
+
+			assert_eq!(
+				bytes2hex("", &ext.child_storage_root(child_info, state_version)),
+				child_root,
+			);
+
+			assert_eq!(bytes2hex("", &ext.storage_root(state_version)), root);
+
+			// Calling a second time should use it from the cache
+			assert_eq!(
+				bytes2hex("", &ext.child_storage_root(child_info, state_version)),
+				child_root,
+			);
+		}
 	}
 
 	#[test]
 	fn extrinsic_changes_are_collected() {
-		let mut overlay = OverlayedChanges::default();
+		let mut overlay = OverlayedChanges::<Blake2Hasher>::default();
 		overlay.set_collect_extrinsics(true);
 
 		overlay.start_transaction();
@@ -919,7 +1072,7 @@ mod tests {
 
 	#[test]
 	fn next_storage_key_change_works() {
-		let mut overlay = OverlayedChanges::default();
+		let mut overlay = OverlayedChanges::<Blake2Hasher>::default();
 		overlay.start_transaction();
 		overlay.set_storage(vec![20], Some(vec![20]));
 		overlay.set_storage(vec![30], Some(vec![30]));
@@ -960,7 +1113,7 @@ mod tests {
 		let child_info = ChildInfo::new_default(b"Child1");
 		let child_info = &child_info;
 		let child = child_info.storage_key();
-		let mut overlay = OverlayedChanges::default();
+		let mut overlay = OverlayedChanges::<Blake2Hasher>::default();
 		overlay.start_transaction();
 		overlay.set_child_storage(child_info, vec![20], Some(vec![20]));
 		overlay.set_child_storage(child_info, vec![30], Some(vec![30]));

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

@@ -23,8 +23,8 @@ use std::{
 };
 
 use crate::{
-	backend::Backend, ext::Ext, InMemoryBackend, OverlayedChanges, StorageKey,
-	StorageTransactionCache, StorageValue, TrieBackendBuilder,
+	backend::Backend, ext::Ext, InMemoryBackend, OverlayedChanges, StorageKey, StorageValue,
+	TrieBackendBuilder,
 };
 
 use hash_db::{HashDB, Hasher};
@@ -36,7 +36,7 @@ use sp_core::{
 	},
 };
 use sp_externalities::{Extension, ExtensionStore, Extensions};
-use sp_trie::StorageProof;
+use sp_trie::{PrefixedMemoryDB, StorageProof};
 
 /// Simple HashMap-based Externalities impl.
 pub struct TestExternalities<H>
@@ -45,10 +45,8 @@ where
 	H::Out: codec::Codec + Ord,
 {
 	/// The overlay changed storage.
-	overlay: OverlayedChanges,
+	overlay: OverlayedChanges<H>,
 	offchain_db: TestPersistentOffchainDB,
-	storage_transaction_cache:
-		StorageTransactionCache<<InMemoryBackend<H> as Backend<H>>::Transaction, H>,
 	/// Storage backend.
 	pub backend: InMemoryBackend<H>,
 	/// Extensions.
@@ -64,12 +62,7 @@ where
 {
 	/// Get externalities implementation.
 	pub fn ext(&mut self) -> Ext<H, InMemoryBackend<H>> {
-		Ext::new(
-			&mut self.overlay,
-			&mut self.storage_transaction_cache,
-			&self.backend,
-			Some(&mut self.extensions),
-		)
+		Ext::new(&mut self.overlay, &self.backend, Some(&mut self.extensions))
 	}
 
 	/// Create a new instance of `TestExternalities` with storage.
@@ -112,13 +105,12 @@ where
 			offchain_db,
 			extensions: Default::default(),
 			backend,
-			storage_transaction_cache: Default::default(),
 			state_version,
 		}
 	}
 
 	/// Returns the overlayed changes.
-	pub fn overlayed_changes(&self) -> &OverlayedChanges {
+	pub fn overlayed_changes(&self) -> &OverlayedChanges<H> {
 		&self.overlay
 	}
 
@@ -165,32 +157,50 @@ where
 	/// This can be used as a fast way to restore the storage state from a backup because the trie
 	/// does not need to be computed.
 	pub fn from_raw_snapshot(
-		&mut self,
-		raw_storage: Vec<(H::Out, (Vec<u8>, i32))>,
+		raw_storage: Vec<(Vec<u8>, (Vec<u8>, i32))>,
 		storage_root: H::Out,
-	) {
-		for (k, (v, ref_count)) in raw_storage {
+		state_version: StateVersion,
+	) -> Self {
+		let mut backend = PrefixedMemoryDB::default();
+
+		for (key, (v, ref_count)) in raw_storage {
+			let mut hash = H::Out::default();
+			let hash_len = hash.as_ref().len();
+
+			if key.len() < hash_len {
+				log::warn!("Invalid key in `from_raw_snapshot`: {key:?}");
+				continue
+			}
+
+			hash.as_mut().copy_from_slice(&key[(key.len() - hash_len)..]);
+
 			// Each time .emplace is called the internal MemoryDb ref count increments.
 			// Repeatedly call emplace to initialise the ref count to the correct value.
 			for _ in 0..ref_count {
-				self.backend.backend_storage_mut().emplace(k, hash_db::EMPTY_PREFIX, v.clone());
+				backend.emplace(hash, (&key[..(key.len() - hash_len)], None), v.clone());
 			}
 		}
-		self.backend.set_root(storage_root);
+
+		Self {
+			backend: TrieBackendBuilder::new(backend, storage_root).build(),
+			overlay: Default::default(),
+			offchain_db: Default::default(),
+			extensions: Default::default(),
+			state_version,
+		}
 	}
 
 	/// Drains the underlying raw storage key/values and returns the root hash.
 	///
 	/// Useful for backing up the storage in a format that can be quickly re-loaded.
-	///
-	/// Note: This DB will be inoperable after this call.
-	pub fn into_raw_snapshot(mut self) -> (Vec<(H::Out, (Vec<u8>, i32))>, H::Out) {
+	pub fn into_raw_snapshot(mut self) -> (Vec<(Vec<u8>, (Vec<u8>, i32))>, H::Out) {
 		let raw_key_values = self
 			.backend
 			.backend_storage_mut()
 			.drain()
 			.into_iter()
-			.collect::<Vec<(H::Out, (Vec<u8>, i32))>>();
+			.filter(|(_, (_, r))| *r > 0)
+			.collect::<Vec<(Vec<u8>, (Vec<u8>, i32))>>();
 
 		(raw_key_values, *self.backend.root())
 	}
@@ -220,11 +230,7 @@ where
 	///
 	/// This will panic if there are still open transactions.
 	pub fn commit_all(&mut self) -> Result<(), String> {
-		let changes = self.overlay.drain_storage_changes::<_, _>(
-			&self.backend,
-			&mut Default::default(),
-			self.state_version,
-		)?;
+		let changes = self.overlay.drain_storage_changes(&self.backend, self.state_version)?;
 
 		self.backend
 			.apply_transaction(changes.transaction_storage_root, changes.transaction);
@@ -248,12 +254,8 @@ where
 		let proving_backend = TrieBackendBuilder::wrap(&self.backend)
 			.with_recorder(Default::default())
 			.build();
-		let mut proving_ext = Ext::new(
-			&mut self.overlay,
-			&mut self.storage_transaction_cache,
-			&proving_backend,
-			Some(&mut self.extensions),
-		);
+		let mut proving_ext =
+			Ext::new(&mut self.overlay, &proving_backend, Some(&mut self.extensions));
 
 		let outcome = sp_externalities::set_and_run_with_externalities(&mut proving_ext, execute);
 		let proof = proving_backend.extract_proof().expect("Failed to extract storage proof");
@@ -409,36 +411,25 @@ mod tests {
 		original_ext.insert_child(child_info.clone(), b"cattytown".to_vec(), b"is_dark".to_vec());
 		original_ext.insert_child(child_info.clone(), b"doggytown".to_vec(), b"is_sunny".to_vec());
 
-		// Call emplace on one of the keys to increment the MemoryDb refcount, so we can check
-		// that it is intact in the recovered_ext.
-		let keys = original_ext.backend.backend_storage_mut().keys();
-		let expected_ref_count = 5;
-		let ref_count_key = keys.into_iter().next().unwrap().0;
-		for _ in 0..expected_ref_count - 1 {
-			original_ext.backend.backend_storage_mut().emplace(
-				ref_count_key,
-				hash_db::EMPTY_PREFIX,
-				// We can use anything for the 'value' because it does not affect behavior when
-				// emplacing an existing key.
-				(&[0u8; 32]).to_vec(),
-			);
-		}
-		let refcount = original_ext
-			.backend
-			.backend_storage()
-			.raw(&ref_count_key, hash_db::EMPTY_PREFIX)
-			.unwrap()
-			.1;
-		assert_eq!(refcount, expected_ref_count);
+		// Apply the backend to itself again to increase the ref count of all nodes.
+		original_ext.backend.apply_transaction(
+			*original_ext.backend.root(),
+			original_ext.backend.clone().into_storage(),
+		);
+
+		// Ensure all have the correct ref counrt
+		assert!(original_ext.backend.backend_storage().keys().values().all(|r| *r == 2));
 
 		// Drain the raw storage and root.
 		let root = *original_ext.backend.root();
 		let (raw_storage, storage_root) = original_ext.into_raw_snapshot();
 
 		// Load the raw storage and root into a new TestExternalities.
-		let mut recovered_ext =
-			TestExternalities::<BlakeTwo256>::from((Default::default(), Default::default()));
-		recovered_ext.from_raw_snapshot(raw_storage, storage_root);
+		let recovered_ext = TestExternalities::<BlakeTwo256>::from_raw_snapshot(
+			raw_storage,
+			storage_root,
+			Default::default(),
+		);
 
 		// Check the storage root is the same as the original
 		assert_eq!(root, *recovered_ext.backend.root());
@@ -458,14 +449,8 @@ mod tests {
 			Some(b"is_sunny".to_vec())
 		);
 
-		// Check the refcount of the key with > 1 refcount is correct.
-		let refcount = recovered_ext
-			.backend
-			.backend_storage()
-			.raw(&ref_count_key, hash_db::EMPTY_PREFIX)
-			.unwrap()
-			.1;
-		assert_eq!(refcount, expected_ref_count);
+		// Ensure all have the correct ref count after importing
+		assert!(recovered_ext.backend.backend_storage().keys().values().all(|r| *r == 2));
 	}
 
 	#[test]

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

@@ -30,6 +30,7 @@ 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},
@@ -377,7 +378,6 @@ where
 	H::Out: Ord + Codec,
 {
 	type Error = crate::DefaultError;
-	type Transaction = S::Overlay;
 	type TrieBackendStorage = S;
 	type RawIter = crate::trie_backend_essence::RawIter<S, H, C>;
 
@@ -458,7 +458,7 @@ where
 		&self,
 		delta: impl Iterator<Item = (&'a [u8], Option<&'a [u8]>)>,
 		state_version: StateVersion,
-	) -> (H::Out, Self::Transaction)
+	) -> (H::Out, PrefixedMemoryDB<H>)
 	where
 		H::Out: Ord,
 	{
@@ -470,7 +470,7 @@ where
 		child_info: &ChildInfo,
 		delta: impl Iterator<Item = (&'a [u8], Option<&'a [u8]>)>,
 		state_version: StateVersion,
-	) -> (H::Out, bool, Self::Transaction)
+	) -> (H::Out, bool, PrefixedMemoryDB<H>)
 	where
 		H::Out: Ord,
 	{
@@ -529,7 +529,7 @@ pub mod tests {
 	use sp_trie::{
 		cache::{CacheSize, SharedTrieCache},
 		trie_types::{TrieDBBuilder, TrieDBMutBuilderV0, TrieDBMutBuilderV1},
-		KeySpacedDBMut, PrefixedKey, PrefixedMemoryDB, Trie, TrieCache, TrieMut,
+		KeySpacedDBMut, PrefixedMemoryDB, Trie, TrieCache, TrieMut,
 	};
 	use std::iter;
 	use trie_db::NodeCodec;
@@ -1187,7 +1187,7 @@ pub mod tests {
 			(Some(child_info_1.clone()), (28..65).map(|i| (vec![i], Some(vec![i]))).collect()),
 			(Some(child_info_2.clone()), (10..15).map(|i| (vec![i], Some(vec![i]))).collect()),
 		];
-		let in_memory = new_in_mem::<BlakeTwo256, PrefixedKey<BlakeTwo256>>();
+		let in_memory = new_in_mem::<BlakeTwo256>();
 		let in_memory = in_memory.update(contents, state_version);
 		let child_storage_keys = vec![child_info_1.to_owned(), child_info_2.to_owned()];
 		let in_memory_root = in_memory
@@ -1292,7 +1292,7 @@ pub mod tests {
 					.collect(),
 			),
 		];
-		let in_memory = new_in_mem::<BlakeTwo256, PrefixedKey<BlakeTwo256>>();
+		let in_memory = new_in_mem::<BlakeTwo256>();
 		let in_memory = in_memory.update(contents, state_version);
 		let child_storage_keys = vec![child_info_1.to_owned()];
 		let in_memory_root = in_memory
@@ -1480,7 +1480,7 @@ pub mod tests {
 			(Some(child_info_1.clone()), vec![(key.clone(), Some(child_trie_1_val.clone()))]),
 			(Some(child_info_2.clone()), vec![(key.clone(), Some(child_trie_2_val.clone()))]),
 		];
-		let in_memory = new_in_mem::<BlakeTwo256, PrefixedKey<BlakeTwo256>>();
+		let in_memory = new_in_mem::<BlakeTwo256>();
 		let in_memory = in_memory.update(contents, state_version);
 		let child_storage_keys = vec![child_info_1.to_owned(), child_info_2.to_owned()];
 		let in_memory_root = in_memory

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

@@ -19,7 +19,7 @@
 //! from storage.
 
 use crate::{
-	backend::{Consolidate, IterArgs, StorageIterator},
+	backend::{IterArgs, StorageIterator},
 	trie_backend::TrieCacheProvider,
 	warn, StorageKey, StorageValue,
 };
@@ -35,7 +35,8 @@ 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,
 	trie_types::{TrieDBBuilder, TrieError},
-	DBValue, KeySpacedDB, NodeCodec, Trie, TrieCache, TrieDBRawIterator, TrieRecorder,
+	DBValue, KeySpacedDB, NodeCodec, PrefixedMemoryDB, Trie, TrieCache, TrieDBRawIterator,
+	TrieRecorder,
 };
 #[cfg(feature = "std")]
 use std::{collections::HashMap, sync::Arc};
@@ -621,8 +622,8 @@ where
 		&self,
 		delta: impl Iterator<Item = (&'a [u8], Option<&'a [u8]>)>,
 		state_version: StateVersion,
-	) -> (H::Out, S::Overlay) {
-		let mut write_overlay = S::Overlay::default();
+	) -> (H::Out, PrefixedMemoryDB<H>) {
+		let mut write_overlay = PrefixedMemoryDB::default();
 
 		let root = self.with_recorder_and_cache_for_storage_root(None, |recorder, cache| {
 			let mut eph = Ephemeral::new(self.backend_storage(), &mut write_overlay);
@@ -654,11 +655,11 @@ where
 		child_info: &ChildInfo,
 		delta: impl Iterator<Item = (&'a [u8], Option<&'a [u8]>)>,
 		state_version: StateVersion,
-	) -> (H::Out, bool, S::Overlay) {
+	) -> (H::Out, bool, PrefixedMemoryDB<H>) {
 		let default_root = match child_info.child_type() {
 			ChildType::ParentKeyId => empty_child_trie_root::<sp_trie::LayoutV1<H>>(),
 		};
-		let mut write_overlay = S::Overlay::default();
+		let mut write_overlay = PrefixedMemoryDB::default();
 		let child_root = match self.child_root(child_info) {
 			Ok(Some(hash)) => hash,
 			Ok(None) => default_root,
@@ -707,7 +708,7 @@ where
 
 pub(crate) struct Ephemeral<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> {
 	storage: &'a S,
-	overlay: &'a mut S::Overlay,
+	overlay: &'a mut PrefixedMemoryDB<H>,
 }
 
 impl<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> AsHashDB<H, DBValue>
@@ -722,7 +723,7 @@ impl<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> AsHashDB<H, DBValue>
 }
 
 impl<'a, S: TrieBackendStorage<H>, H: Hasher> Ephemeral<'a, S, H> {
-	pub fn new(storage: &'a S, overlay: &'a mut S::Overlay) -> Self {
+	pub fn new(storage: &'a S, overlay: &'a mut PrefixedMemoryDB<H>) -> Self {
 		Ephemeral { storage, overlay }
 	}
 }
@@ -768,16 +769,11 @@ impl<'a, S: 'a + TrieBackendStorage<H>, H: Hasher> HashDBRef<H, DBValue> for Eph
 
 /// Key-value pairs storage that is used by trie backend essence.
 pub trait TrieBackendStorage<H: Hasher>: Send + Sync {
-	/// Type of in-memory overlay.
-	type Overlay: HashDB<H, DBValue> + Default + Consolidate;
-
 	/// Get the value stored at key.
 	fn get(&self, key: &H::Out, prefix: Prefix) -> Result<Option<DBValue>>;
 }
 
 impl<T: TrieBackendStorage<H>, H: Hasher> TrieBackendStorage<H> for &T {
-	type Overlay = T::Overlay;
-
 	fn get(&self, key: &H::Out, prefix: Prefix) -> Result<Option<DBValue>> {
 		(*self).get(key, prefix)
 	}
@@ -786,8 +782,6 @@ impl<T: TrieBackendStorage<H>, H: Hasher> TrieBackendStorage<H> for &T {
 // This implementation is used by normal storage trie clients.
 #[cfg(feature = "std")]
 impl<H: Hasher> TrieBackendStorage<H> for Arc<dyn Storage<H>> {
-	type Overlay = sp_trie::PrefixedMemoryDB<H>;
-
 	fn get(&self, key: &H::Out, prefix: Prefix) -> Result<Option<DBValue>> {
 		Storage::<H>::get(std::ops::Deref::deref(self), key, prefix)
 	}
@@ -798,8 +792,6 @@ where
 	H: Hasher,
 	KF: sp_trie::KeyFunction<H> + Send + Sync,
 {
-	type Overlay = Self;
-
 	fn get(&self, key: &H::Out, prefix: Prefix) -> Result<Option<DBValue>> {
 		Ok(hash_db::HashDB::get(self, key, prefix))
 	}