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/utils/frame/benchmarking-cli/src/block/bench.rs

@@ -77,7 +77,7 @@ where
 		+ UsageProvider<Block>
 		+ BlockBackend<Block>
 		+ HeaderBackend<Block>,
-	C::Api: ApiExt<Block, StateBackend = BA::State> + BlockBuilderApi<Block>,
+	C::Api: ApiExt<Block> + BlockBuilderApi<Block>,
 {
 	/// Returns a new [`Self`] from the arguments.
 	pub fn new(client: Arc<C>, params: BenchmarkParams) -> Self {

substrate/utils/frame/benchmarking-cli/src/block/cmd.rs

@@ -90,7 +90,7 @@ impl BlockCmd {
 			+ StorageProvider<Block, BA>
 			+ UsageProvider<Block>
 			+ HeaderBackend<Block>,
-		C::Api: ApiExt<Block, StateBackend = BA::State> + BlockBuilderApi<Block>,
+		C::Api: ApiExt<Block> + BlockBuilderApi<Block>,
 	{
 		// Put everything in the benchmark type to have the generic types handy.
 		Benchmark::new(client, self.params.clone()).run()

substrate/utils/frame/benchmarking-cli/src/extrinsic/bench.rs

@@ -76,7 +76,7 @@ where
 	C: BlockBuilderProvider<BA, Block, C>
 		+ ProvideRuntimeApi<Block>
 		+ sp_blockchain::HeaderBackend<Block>,
-	C::Api: ApiExt<Block, StateBackend = BA::State> + BlockBuilderApi<Block>,
+	C::Api: ApiExt<Block> + BlockBuilderApi<Block>,
 {
 	/// Create a new [`Self`] from the arguments.
 	pub fn new(

substrate/utils/frame/benchmarking-cli/src/extrinsic/cmd.rs

@@ -97,7 +97,7 @@ impl ExtrinsicCmd {
 		C: BlockBuilderProvider<BA, Block, C>
 			+ ProvideRuntimeApi<Block>
 			+ sp_blockchain::HeaderBackend<Block>,
-		C::Api: ApiExt<Block, StateBackend = BA::State> + BlockBuilderApi<Block>,
+		C::Api: ApiExt<Block> + BlockBuilderApi<Block>,
 	{
 		// Short circuit if --list was specified.
 		if self.params.list {

substrate/utils/frame/benchmarking-cli/src/overhead/cmd.rs

@@ -111,7 +111,7 @@ impl OverheadCmd {
 		C: BlockBuilderProvider<BA, Block, C>
 			+ ProvideRuntimeApi<Block>
 			+ sp_blockchain::HeaderBackend<Block>,
-		C::Api: ApiExt<Block, StateBackend = BA::State> + BlockBuilderApi<Block>,
+		C::Api: ApiExt<Block> + BlockBuilderApi<Block>,
 	{
 		if ext_builder.pallet() != "system" || ext_builder.extrinsic() != "remark" {
 			return Err(format!("The extrinsic builder is required to build `System::Remark` extrinsics but builds `{}` extrinsics instead", ext_builder.name()).into());

substrate/utils/frame/remote-externalities/Cargo.toml

@@ -16,8 +16,8 @@ jsonrpsee = { version = "0.16.2", features = ["http-client"] }
 codec = { package = "parity-scale-codec", version = "3.6.1" }
 log = "0.4.17"
 serde = "1.0.163"
-frame-support = { version = "4.0.0-dev", optional = true, path = "../../../frame/support" }
 sp-core = { version = "21.0.0", path = "../../../primitives/core" }
+sp-state-machine = { version = "0.28.0", path = "../../../primitives/state-machine" }
 sp-io = { version = "23.0.0", path = "../../../primitives/io" }
 sp-runtime = { version = "24.0.0", path = "../../../primitives/runtime" }
 tokio = { version = "1.22.0", features = ["macros", "rt-multi-thread"] }
@@ -29,9 +29,7 @@ spinners = "4.1.0"
 tokio-retry = "0.3.0"
 
 [dev-dependencies]
-frame-support = { version = "4.0.0-dev", path = "../../../frame/support" }
-pallet-elections-phragmen = { version = "5.0.0-dev", path = "../../../frame/elections-phragmen" }
-tracing-subscriber = { version = "0.3.16", features = ["env-filter"] }
+sp-tracing = { version = "10.0.0", path = "../../../primitives/tracing" }
 
 [features]
-remote-test = ["frame-support"]
+remote-test = []

substrate/utils/frame/remote-externalities/src/lib.rs

@@ -36,10 +36,12 @@ use sp_core::{
 		well_known_keys::{is_default_child_storage_key, DEFAULT_CHILD_STORAGE_KEY_PREFIX},
 		ChildInfo, ChildType, PrefixedStorageKey, StorageData, StorageKey,
 	},
-	H256,
 };
-pub use sp_io::TestExternalities;
-use sp_runtime::{traits::Block as BlockT, StateVersion};
+use sp_runtime::{
+	traits::{Block as BlockT, HashingFor},
+	StateVersion,
+};
+use sp_state_machine::TestExternalities;
 use spinners::{Spinner, Spinners};
 use std::{
 	cmp::max,
@@ -58,7 +60,7 @@ type SnapshotVersion = Compact<u16>;
 
 const LOG_TARGET: &str = "remote-ext";
 const DEFAULT_HTTP_ENDPOINT: &str = "https://rpc.polkadot.io:443";
-const SNAPSHOT_VERSION: SnapshotVersion = Compact(2);
+const SNAPSHOT_VERSION: SnapshotVersion = Compact(3);
 
 /// The snapshot that we store on disk.
 #[derive(Decode, Encode)]
@@ -67,16 +69,16 @@ struct Snapshot<B: BlockT> {
 	state_version: StateVersion,
 	block_hash: B::Hash,
 	// <Vec<Key, (Value, MemoryDbRefCount)>>
-	raw_storage: Vec<(H256, (Vec<u8>, i32))>,
-	storage_root: H256,
+	raw_storage: Vec<(Vec<u8>, (Vec<u8>, i32))>,
+	storage_root: B::Hash,
 }
 
 impl<B: BlockT> Snapshot<B> {
 	pub fn new(
 		state_version: StateVersion,
 		block_hash: B::Hash,
-		raw_storage: Vec<(H256, (Vec<u8>, i32))>,
-		storage_root: H256,
+		raw_storage: Vec<(Vec<u8>, (Vec<u8>, i32))>,
+		storage_root: B::Hash,
 	) -> Self {
 		Self {
 			snapshot_version: SNAPSHOT_VERSION,
@@ -91,21 +93,14 @@ impl<B: BlockT> Snapshot<B> {
 		let bytes = fs::read(path).map_err(|_| "fs::read failed.")?;
 		// The first item in the SCALE encoded struct bytes is the snapshot version. We decode and
 		// check that first, before proceeding to decode the rest of the snapshot.
-		let maybe_version: Result<SnapshotVersion, _> = Decode::decode(&mut &*bytes);
-		match maybe_version {
-			Ok(snapshot_version) => {
-				if snapshot_version != SNAPSHOT_VERSION {
-					return Err(
-						"Unsupported snapshot version detected. Please create a new snapshot.",
-					)
-				}
-				match Decode::decode(&mut &*bytes) {
-					Ok(snapshot) => return Ok(snapshot),
-					Err(_) => Err("Decode failed"),
-				}
-			},
-			Err(_) => Err("Decode failed"),
+		let snapshot_version = SnapshotVersion::decode(&mut &*bytes)
+			.map_err(|_| "Failed to decode snapshot version")?;
+
+		if snapshot_version != SNAPSHOT_VERSION {
+			return Err("Unsupported snapshot version detected. Please create a new snapshot.")
 		}
+
+		Decode::decode(&mut &*bytes).map_err(|_| "Decode failed")
 	}
 }
 
@@ -113,13 +108,13 @@ impl<B: BlockT> Snapshot<B> {
 /// bits and pieces to it, and can be loaded remotely.
 pub struct RemoteExternalities<B: BlockT> {
 	/// The inner externalities.
-	pub inner_ext: TestExternalities,
+	pub inner_ext: TestExternalities<HashingFor<B>>,
 	/// The block hash it which we created this externality env.
 	pub block_hash: B::Hash,
 }
 
 impl<B: BlockT> Deref for RemoteExternalities<B> {
-	type Target = TestExternalities;
+	type Target = TestExternalities<HashingFor<B>>;
 	fn deref(&self) -> &Self::Target {
 		&self.inner_ext
 	}
@@ -319,8 +314,6 @@ pub struct Builder<B: BlockT> {
 	overwrite_state_version: Option<StateVersion>,
 }
 
-// NOTE: ideally we would use `DefaultNoBound` here, but not worth bringing in frame-support for
-// that.
 impl<B: BlockT> Default for Builder<B> {
 	fn default() -> Self {
 		Self {
@@ -576,7 +569,7 @@ where
 		&self,
 		prefix: StorageKey,
 		at: B::Hash,
-		pending_ext: &mut TestExternalities,
+		pending_ext: &mut TestExternalities<HashingFor<B>>,
 	) -> Result<Vec<KeyValue>, &'static str> {
 		let start = Instant::now();
 		let mut sp = Spinner::with_timer(Spinners::Dots, "Scraping keys...".into());
@@ -768,7 +761,7 @@ where
 	async fn load_child_remote(
 		&self,
 		top_kv: &[KeyValue],
-		pending_ext: &mut TestExternalities,
+		pending_ext: &mut TestExternalities<HashingFor<B>>,
 	) -> Result<ChildKeyValues, &'static str> {
 		let child_roots = top_kv
 			.into_iter()
@@ -826,7 +819,7 @@ where
 	/// cache, we can also optimize further.
 	async fn load_top_remote(
 		&self,
-		pending_ext: &mut TestExternalities,
+		pending_ext: &mut TestExternalities<HashingFor<B>>,
 	) -> Result<TopKeyValues, &'static str> {
 		let config = self.as_online();
 		let at = self
@@ -926,7 +919,9 @@ where
 	/// `load_child_remote`.
 	///
 	/// Must be called after `init_remote_client`.
-	async fn load_remote_and_maybe_save(&mut self) -> Result<TestExternalities, &'static str> {
+	async fn load_remote_and_maybe_save(
+		&mut self,
+	) -> Result<TestExternalities<HashingFor<B>>, &'static str> {
 		let state_version =
 			StateApi::<B::Hash>::runtime_version(self.as_online().rpc_client(), None)
 				.await
@@ -966,13 +961,11 @@ where
 			std::fs::write(path, encoded).map_err(|_| "fs::write failed")?;
 
 			// pending_ext was consumed when creating the snapshot, need to reinitailize it
-			let mut pending_ext = TestExternalities::new_with_code_and_state(
-				Default::default(),
-				Default::default(),
+			return Ok(TestExternalities::from_raw_snapshot(
+				raw_storage,
+				storage_root,
 				self.overwrite_state_version.unwrap_or(state_version),
-			);
-			pending_ext.from_raw_snapshot(raw_storage, storage_root);
-			return Ok(pending_ext)
+			))
 		}
 
 		Ok(pending_ext)
@@ -995,12 +988,11 @@ where
 		let Snapshot { snapshot_version: _, block_hash, state_version, raw_storage, storage_root } =
 			Snapshot::<B>::load(&config.state_snapshot.path)?;
 
-		let mut inner_ext = TestExternalities::new_with_code_and_state(
-			Default::default(),
-			Default::default(),
+		let inner_ext = TestExternalities::from_raw_snapshot(
+			raw_storage,
+			storage_root,
 			self.overwrite_state_version.unwrap_or(state_version),
 		);
-		inner_ext.from_raw_snapshot(raw_storage, storage_root);
 		sp.stop_with_message(format!("✅ Loaded snapshot ({:.2}s)", start.elapsed().as_secs_f32()));
 
 		Ok(RemoteExternalities { inner_ext, block_hash })
@@ -1099,17 +1091,12 @@ where
 
 #[cfg(test)]
 mod test_prelude {
-	use tracing_subscriber::EnvFilter;
-
 	pub(crate) use super::*;
 	pub(crate) use sp_runtime::testing::{Block as RawBlock, ExtrinsicWrapper, H256 as Hash};
 	pub(crate) type Block = RawBlock<ExtrinsicWrapper<Hash>>;
 
 	pub(crate) fn init_logger() {
-		let _ = tracing_subscriber::fmt()
-			.with_env_filter(EnvFilter::from_default_env())
-			.with_level(true)
-			.try_init();
+		let _ = sp_tracing::try_init_simple();
 	}
 }
 
@@ -1369,9 +1356,6 @@ mod remote_tests {
 			.filter(|p| p.path().file_name().unwrap_or_default() == CACHE)
 			.collect::<Vec<_>>();
 
-		let snap: Snapshot<Block> = Builder::<Block>::new().load_snapshot(CACHE.into()).unwrap();
-		assert!(matches!(snap, Snapshot { raw_storage, .. } if raw_storage.len() > 0));
-
 		assert!(to_delete.len() == 1);
 		let to_delete = to_delete.first().unwrap();
 		assert!(std::fs::metadata(to_delete.path()).unwrap().size() > 1);
@@ -1401,9 +1385,6 @@ mod remote_tests {
 			.filter(|p| p.path().file_name().unwrap_or_default() == CACHE)
 			.collect::<Vec<_>>();
 
-		let snap: Snapshot<Block> = Builder::<Block>::new().load_snapshot(CACHE.into()).unwrap();
-		assert!(matches!(snap, Snapshot { raw_storage, .. } if raw_storage.len() > 0));
-
 		assert!(to_delete.len() == 1);
 		let to_delete = to_delete.first().unwrap();
 		assert!(std::fs::metadata(to_delete.path()).unwrap().size() > 1);

substrate/utils/frame/try-runtime/cli/src/commands/fast_forward.rs

@@ -25,11 +25,11 @@ use sc_executor::{sp_wasm_interface::HostFunctions, WasmExecutor};
 use serde::de::DeserializeOwned;
 use sp_core::H256;
 use sp_inherents::{InherentData, InherentDataProvider};
-use sp_io::TestExternalities;
 use sp_runtime::{
-	traits::{Header, NumberFor, One},
+	traits::{HashingFor, Header, NumberFor, One},
 	Digest,
 };
+use sp_state_machine::TestExternalities;
 use std::{fmt::Debug, str::FromStr};
 use substrate_rpc_client::{ws_client, ChainApi};
 
@@ -92,8 +92,8 @@ where
 }
 
 /// Call `method` with `data` and return the result. `externalities` will not change.
-async fn dry_run<T: Decode, Block: BlockT, HostFns: HostFunctions>(
-	externalities: &TestExternalities,
+fn dry_run<T: Decode, Block: BlockT, HostFns: HostFunctions>(
+	externalities: &TestExternalities<HashingFor<Block>>,
 	executor: &WasmExecutor<HostFns>,
 	method: &'static str,
 	data: &[u8],
@@ -111,7 +111,7 @@ async fn dry_run<T: Decode, Block: BlockT, HostFns: HostFunctions>(
 
 /// Call `method` with `data` and actually save storage changes to `externalities`.
 async fn run<Block: BlockT, HostFns: HostFunctions>(
-	externalities: &mut TestExternalities,
+	externalities: &mut TestExternalities<HashingFor<Block>>,
 	executor: &WasmExecutor<HostFns>,
 	method: &'static str,
 	data: &[u8],
@@ -124,11 +124,8 @@ async fn run<Block: BlockT, HostFns: HostFunctions>(
 		full_extensions(executor.clone()),
 	)?;
 
-	let storage_changes = changes.drain_storage_changes(
-		&externalities.backend,
-		&mut Default::default(),
-		externalities.state_version,
-	)?;
+	let storage_changes =
+		changes.drain_storage_changes(&externalities.backend, externalities.state_version)?;
 
 	externalities
 		.backend
@@ -143,7 +140,7 @@ async fn next_empty_block<
 	HostFns: HostFunctions,
 	BBIP: BlockBuildingInfoProvider<Block, Option<(InherentData, Digest)>>,
 >(
-	externalities: &mut TestExternalities,
+	externalities: &mut TestExternalities<HashingFor<Block>>,
 	executor: &WasmExecutor<HostFns>,
 	parent_height: NumberFor<Block>,
 	parent_hash: Block::Hash,
@@ -182,8 +179,7 @@ async fn next_empty_block<
 			executor,
 			"BlockBuilder_inherent_extrinsics",
 			&inherent_data.encode(),
-		)
-		.await?;
+		)?;
 	}
 
 	for xt in &extrinsics {
@@ -196,8 +192,7 @@ async fn next_empty_block<
 		executor,
 		"BlockBuilder_finalize_block",
 		&[0u8; 0],
-	)
-	.await?;
+	)?;
 
 	run::<Block, _>(externalities, executor, "BlockBuilder_finalize_block", &[0u8; 0]).await?;
 

substrate/utils/frame/try-runtime/cli/src/commands/follow_chain.rs

@@ -177,7 +177,6 @@ where
 		let storage_changes = changes
 			.drain_storage_changes(
 				&state_ext.backend,
-				&mut Default::default(),
 				// Note that in case a block contains a runtime upgrade, state version could
 				// potentially be incorrect here, this is very niche and would only result in
 				// unaligned roots, so this use case is ignored for now.

substrate/utils/frame/try-runtime/cli/src/lib.rs

@@ -28,7 +28,6 @@ use crate::block_building_info::BlockBuildingInfoProvider;
 use parity_scale_codec::Decode;
 use remote_externalities::{
 	Builder, Mode, OfflineConfig, OnlineConfig, RemoteExternalities, SnapshotConfig,
-	TestExternalities,
 };
 use sc_cli::{
 	execution_method_from_cli, CliConfiguration, RuntimeVersion, WasmExecutionMethod,
@@ -38,7 +37,6 @@ use sc_cli::{
 use sc_executor::{
 	sp_wasm_interface::HostFunctions, HeapAllocStrategy, WasmExecutor, DEFAULT_HEAP_ALLOC_STRATEGY,
 };
-use sp_api::HashT;
 use sp_core::{
 	hexdisplay::HexDisplay,
 	offchain::{
@@ -53,10 +51,12 @@ use sp_externalities::Extensions;
 use sp_inherents::InherentData;
 use sp_keystore::{testing::MemoryKeystore, KeystoreExt};
 use sp_runtime::{
-	traits::{BlakeTwo256, Block as BlockT, NumberFor},
+	traits::{BlakeTwo256, Block as BlockT, Hash as HashT, HashingFor, NumberFor},
 	DeserializeOwned, Digest,
 };
-use sp_state_machine::{CompactProof, OverlayedChanges, StateMachine, TrieBackendBuilder};
+use sp_state_machine::{
+	CompactProof, OverlayedChanges, StateMachine, TestExternalities, TrieBackendBuilder,
+};
 use sp_version::StateVersion;
 use std::{fmt::Debug, path::PathBuf, str::FromStr};
 
@@ -514,7 +514,7 @@ pub(crate) fn build_executor<H: HostFunctions>(shared: &SharedParams) -> WasmExe
 /// Ensure that the given `ext` is compiled with `try-runtime`
 fn ensure_try_runtime<Block: BlockT, HostFns: HostFunctions>(
 	executor: &WasmExecutor<HostFns>,
-	ext: &mut TestExternalities,
+	ext: &mut TestExternalities<HashingFor<Block>>,
 ) -> bool {
 	use sp_api::RuntimeApiInfo;
 	let final_code = ext
@@ -532,12 +532,12 @@ fn ensure_try_runtime<Block: BlockT, HostFns: HostFunctions>(
 /// Execute the given `method` and `data` on top of `ext`, returning the results (encoded) and the
 /// state `changes`.
 pub(crate) fn state_machine_call<Block: BlockT, HostFns: HostFunctions>(
-	ext: &TestExternalities,
+	ext: &TestExternalities<HashingFor<Block>>,
 	executor: &WasmExecutor<HostFns>,
 	method: &'static str,
 	data: &[u8],
 	mut extensions: Extensions,
-) -> sc_cli::Result<(OverlayedChanges, Vec<u8>)> {
+) -> sc_cli::Result<(OverlayedChanges<HashingFor<Block>>, Vec<u8>)> {
 	let mut changes = Default::default();
 	let encoded_results = StateMachine::new(
 		&ext.backend,
@@ -561,13 +561,13 @@ pub(crate) fn state_machine_call<Block: BlockT, HostFns: HostFunctions>(
 ///
 /// Make sure [`LOG_TARGET`] is enabled in logging.
 pub(crate) fn state_machine_call_with_proof<Block: BlockT, HostFns: HostFunctions>(
-	ext: &TestExternalities,
+	ext: &TestExternalities<HashingFor<Block>>,
 	executor: &WasmExecutor<HostFns>,
 	method: &'static str,
 	data: &[u8],
 	mut extensions: Extensions,
 	maybe_export_proof: Option<PathBuf>,
-) -> sc_cli::Result<(OverlayedChanges, Vec<u8>)> {
+) -> sc_cli::Result<(OverlayedChanges<HashingFor<Block>>, Vec<u8>)> {
 	use parity_scale_codec::Encode;
 
 	let mut changes = Default::default();
@@ -624,7 +624,7 @@ pub(crate) fn state_machine_call_with_proof<Block: BlockT, HostFns: HostFunction
 	let proof_size = proof.encoded_size();
 	let compact_proof = proof
 		.clone()
-		.into_compact_proof::<sp_runtime::traits::BlakeTwo256>(pre_root)
+		.into_compact_proof::<HashingFor<Block>>(pre_root)
 		.map_err(|e| {
 			log::error!(target: LOG_TARGET, "failed to generate compact proof {}: {:?}", method, e);
 			e