allow try-runtime and TestExternalities to report PoV size (#10372)

* allow try-runtime and test-externalities to report proof size

* self review

* fix test

* Fix humanized dispaly of bytes

* Fix some test

* Fix some review grumbles

* last of the review comments

* fmt

* remove unused import

* move test

* fix import

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

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* last touches

* fix

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

substrate/Cargo.lock

@@ -9650,6 +9650,7 @@ dependencies = [
  "sp-std",
  "sp-tracing",
  "substrate-test-runtime-client",
+ "zstd",
 ]
 
 [[package]]
@@ -10925,6 +10926,7 @@ dependencies = [
  "sp-state-machine",
  "sp-version",
  "structopt",
+ "zstd",
 ]
 
 [[package]]

substrate/primitives/runtime/Cargo.toml

@@ -38,6 +38,7 @@ sp-state-machine = { version = "0.10.0-dev", path = "../state-machine" }
 sp-api = { version = "4.0.0-dev", path = "../api" }
 substrate-test-runtime-client = { version = "2.0.0", path = "../../test-utils/runtime/client" }
 sp-tracing = { version = "4.0.0-dev", path = "../../primitives/tracing" }
+zstd = "0.9"
 
 [features]
 bench = []

substrate/primitives/runtime/src/lib.rs

@@ -916,9 +916,13 @@ impl<R> TransactionOutcome<R> {
 
 #[cfg(test)]
 mod tests {
+	use crate::traits::BlakeTwo256;
+
 	use super::*;
 	use codec::{Decode, Encode};
 	use sp_core::crypto::Pair;
+	use sp_io::TestExternalities;
+	use sp_state_machine::create_proof_check_backend;
 
 	#[test]
 	fn opaque_extrinsic_serialization() {
@@ -1019,4 +1023,47 @@ mod tests {
 			panic!("Hey, I'm an error");
 		});
 	}
+
+	#[test]
+	fn execute_and_generate_proof_works() {
+		use codec::Encode;
+		use sp_state_machine::Backend;
+		let mut ext = TestExternalities::default();
+
+		ext.insert(b"a".to_vec(), vec![1u8; 33]);
+		ext.insert(b"b".to_vec(), vec![2u8; 33]);
+		ext.insert(b"c".to_vec(), vec![3u8; 33]);
+		ext.insert(b"d".to_vec(), vec![4u8; 33]);
+
+		let pre_root = ext.backend.root().clone();
+		let (_, proof) = ext.execute_and_prove(|| {
+			sp_io::storage::get(b"a");
+			sp_io::storage::get(b"b");
+			sp_io::storage::get(b"v");
+			sp_io::storage::get(b"d");
+		});
+
+		let compact_proof = proof.clone().into_compact_proof::<BlakeTwo256>(pre_root).unwrap();
+		let compressed_proof = zstd::stream::encode_all(&compact_proof.encode()[..], 0).unwrap();
+
+		// just an example of how you'd inspect the size of the proof.
+		println!("proof size: {:?}", proof.encoded_size());
+		println!("compact proof size: {:?}", compact_proof.encoded_size());
+		println!("zstd-compressed compact proof size: {:?}", &compressed_proof.len());
+
+		// create a new trie-backed from the proof and make sure it contains everything
+		let proof_check = create_proof_check_backend::<BlakeTwo256>(pre_root, proof).unwrap();
+		assert_eq!(proof_check.storage(b"a",).unwrap().unwrap(), vec![1u8; 33]);
+
+		let _ = ext.execute_and_prove(|| {
+			sp_io::storage::set(b"a", &vec![1u8; 44]);
+		});
+
+		// ensure that these changes are propagated to the backend.
+
+		ext.execute_with(|| {
+			assert_eq!(sp_io::storage::get(b"a").unwrap(), vec![1u8; 44]);
+			assert_eq!(sp_io::storage::get(b"b").unwrap(), vec![2u8; 33]);
+		});
+	}
 }

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

@@ -188,6 +188,9 @@ mod execution {
 	/// Trie backend with in-memory storage.
 	pub type InMemoryBackend<H> = TrieBackend<MemoryDB<H>, H>;
 
+	/// Proving Trie backend with in-memory storage.
+	pub type InMemoryProvingBackend<'a, H> = ProvingBackend<'a, MemoryDB<H>, H>;
+
 	/// Strategy for executing a call into the runtime.
 	#[derive(Copy, Clone, Eq, PartialEq, Debug)]
 	pub enum ExecutionStrategy {

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

@@ -221,6 +221,11 @@ where
 	pub fn estimate_encoded_size(&self) -> usize {
 		self.0.essence().backend_storage().proof_recorder.estimate_encoded_size()
 	}
+
+	/// Clear the proof recorded data.
+	pub fn clear_recorder(&self) {
+		self.0.essence().backend_storage().proof_recorder.reset()
+	}
 }
 
 impl<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> TrieBackendStorage<H>
@@ -358,7 +363,9 @@ where
 	}
 }
 
-/// Create proof check backend.
+/// Create a backend used for checking the proof., using `H` as hasher.
+///
+/// `proof` and `root` must match, i.e. `root` must be the correct root of `proof` nodes.
 pub fn create_proof_check_backend<H>(
 	root: H::Out,
 	proof: StorageProof,

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

@@ -23,8 +23,8 @@ use std::{
 };
 
 use crate::{
-	backend::Backend, ext::Ext, InMemoryBackend, OverlayedChanges, StorageKey,
-	StorageTransactionCache, StorageValue,
+	backend::Backend, ext::Ext, InMemoryBackend, InMemoryProvingBackend, OverlayedChanges,
+	StorageKey, StorageTransactionCache, StorageValue,
 };
 
 use hash_db::Hasher;
@@ -38,6 +38,7 @@ use sp_core::{
 	traits::TaskExecutorExt,
 };
 use sp_externalities::{Extension, ExtensionStore, Extensions};
+use sp_trie::StorageProof;
 
 /// Simple HashMap-based Externalities impl.
 pub struct TestExternalities<H: Hasher>
@@ -122,6 +123,13 @@ where
 		self.backend.insert(vec![(None, vec![(k, Some(v))])]);
 	}
 
+	/// Insert key/value into backend.
+	///
+	/// This only supports inserting keys in child tries.
+	pub fn insert_child(&mut self, c: sp_core::storage::ChildInfo, k: StorageKey, v: StorageValue) {
+		self.backend.insert(vec![(Some(c), vec![(k, Some(v))])]);
+	}
+
 	/// Registers the given extension for this instance.
 	pub fn register_extension<E: Any + Extension>(&mut self, ext: E) {
 		self.extensions.register(ext);
@@ -171,9 +179,29 @@ where
 		sp_externalities::set_and_run_with_externalities(&mut ext, execute)
 	}
 
+	/// Execute the given closure while `self`, with `proving_backend` as backend, is set as
+	/// externalities.
+	///
+	/// This implementation will wipe the proof recorded in between calls. Consecutive calls will
+	/// get their own proof from scratch.
+	pub fn execute_and_prove<'a, R>(&mut self, execute: impl FnOnce() -> R) -> (R, StorageProof) {
+		let proving_backend = InMemoryProvingBackend::new(&self.backend);
+		let mut proving_ext = Ext::new(
+			&mut self.overlay,
+			&mut self.storage_transaction_cache,
+			&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();
+
+		(outcome, proof)
+	}
+
 	/// Execute the given closure while `self` is set as externalities.
 	///
-	/// Returns the result of the given closure, if no panics occured.
+	/// Returns the result of the given closure, if no panics occurred.
 	/// Otherwise, returns `Err`.
 	pub fn execute_with_safe<R>(
 		&mut self,

substrate/primitives/storage/src/lib.rs

@@ -210,6 +210,14 @@ pub mod well_known_keys {
 	/// Prefix of the default child storage keys in the top trie.
 	pub const DEFAULT_CHILD_STORAGE_KEY_PREFIX: &'static [u8] = b":child_storage:default:";
 
+	/// Whether a key is a default child storage key.
+	///
+	/// This is convenience function which basically checks if the given `key` starts
+	/// with `DEFAULT_CHILD_STORAGE_KEY_PREFIX` and doesn't do anything apart from that.
+	pub fn is_default_child_storage_key(key: &[u8]) -> bool {
+		key.starts_with(DEFAULT_CHILD_STORAGE_KEY_PREFIX)
+	}
+
 	/// Whether a key is a child storage key.
 	///
 	/// This is convenience function which basically checks if the given `key` starts
@@ -231,7 +239,7 @@ pub mod well_known_keys {
 
 /// Information related to a child state.
 #[derive(Debug, Clone)]
-#[cfg_attr(feature = "std", derive(PartialEq, Eq, Hash, PartialOrd, Ord))]
+#[cfg_attr(feature = "std", derive(PartialEq, Eq, Hash, PartialOrd, Ord, Encode, Decode))]
 pub enum ChildInfo {
 	/// This is the one used by default.
 	ParentKeyId(ChildTrieParentKeyId),
@@ -370,16 +378,14 @@ impl ChildType {
 }
 
 /// A child trie of default type.
-/// It uses the same default implementation as the top trie,
-/// top trie being a child trie with no keyspace and no storage key.
-/// Its keyspace is the variable (unprefixed) part of its storage key.
-/// It shares its trie nodes backend storage with every other
-/// child trie, so its storage key needs to be a unique id
-/// that will be use only once.
-/// Those unique id also required to be long enough to avoid any
-/// unique id to be prefixed by an other unique id.
+///
+/// It uses the same default implementation as the top trie, top trie being a child trie with no
+/// keyspace and no storage key. Its keyspace is the variable (unprefixed) part of its storage key.
+/// It shares its trie nodes backend storage with every other child trie, so its storage key needs
+/// to be a unique id that will be use only once. Those unique id also required to be long enough to
+/// avoid any unique id to be prefixed by an other unique id.
 #[derive(Debug, Clone)]
-#[cfg_attr(feature = "std", derive(PartialEq, Eq, Hash, PartialOrd, Ord))]
+#[cfg_attr(feature = "std", derive(PartialEq, Eq, Hash, PartialOrd, Ord, Encode, Decode))]
 pub struct ChildTrieParentKeyId {
 	/// Data is the storage key without prefix.
 	data: Vec<u8>,

substrate/primitives/trie/src/storage_proof.rs

@@ -67,6 +67,7 @@ impl StorageProof {
 	pub fn into_nodes(self) -> Vec<Vec<u8>> {
 		self.trie_nodes
 	}
+
 	/// Creates a `MemoryDB` from `Self`.
 	pub fn into_memory_db<H: Hasher>(self) -> crate::MemoryDB<H> {
 		self.into()
@@ -100,8 +101,9 @@ impl StorageProof {
 
 	/// Returns the estimated encoded size of the compact proof.
 	///
-	/// Runing this operation is a slow operation (build the whole compact proof) and should only be
-	/// in non sensitive path.
+	/// Running this operation is a slow operation (build the whole compact proof) and should only
+	/// be in non sensitive path.
+	///
 	/// Return `None` on error.
 	pub fn encoded_compact_size<H: Hasher>(self, root: H::Out) -> Option<usize> {
 		let compact_proof = self.into_compact_proof::<H>(root);

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

@@ -34,7 +34,10 @@ use serde::de::DeserializeOwned;
 use sp_core::{
 	hashing::twox_128,
 	hexdisplay::HexDisplay,
-	storage::{StorageData, StorageKey},
+	storage::{
+		well_known_keys::{is_default_child_storage_key, DEFAULT_CHILD_STORAGE_KEY_PREFIX},
+		ChildInfo, ChildType, PrefixedStorageKey, StorageData, StorageKey,
+	},
 };
 pub use sp_io::TestExternalities;
 use sp_runtime::traits::Block as BlockT;
@@ -45,7 +48,9 @@ use std::{
 
 pub mod rpc_api;
 
-type KeyPair = (StorageKey, StorageData);
+type KeyValue = (StorageKey, StorageData);
+type TopKeyValues = Vec<KeyValue>;
+type ChildKeyValues = Vec<(ChildInfo, Vec<KeyValue>)>;
 
 const LOG_TARGET: &str = "remote-ext";
 const DEFAULT_TARGET: &str = "wss://rpc.polkadot.io:443";
@@ -53,6 +58,22 @@ const BATCH_SIZE: usize = 1000;
 
 #[rpc(client)]
 pub trait RpcApi<Hash> {
+	#[method(name = "childstate_getKeys")]
+	fn child_get_keys(
+		&self,
+		child_key: PrefixedStorageKey,
+		prefix: StorageKey,
+		hash: Option<Hash>,
+	) -> Result<Vec<StorageKey>, RpcError>;
+
+	#[method(name = "childstate_getStorage")]
+	fn child_get_storage(
+		&self,
+		child_key: PrefixedStorageKey,
+		prefix: StorageKey,
+		hash: Option<Hash>,
+	) -> Result<StorageData, RpcError>;
+
 	#[method(name = "state_getStorage")]
 	fn get_storage(&self, prefix: StorageKey, hash: Option<Hash>) -> Result<StorageData, RpcError>;
 
@@ -180,7 +201,7 @@ impl Default for SnapshotConfig {
 pub struct Builder<B: BlockT> {
 	/// Custom key-pairs to be injected into the externalities. The *hashed* keys and values must
 	/// be given.
-	hashed_key_values: Vec<KeyPair>,
+	hashed_key_values: Vec<KeyValue>,
 	/// Storage entry key prefixes to be injected into the externalities. The *hashed* prefix must
 	/// be given.
 	hashed_prefixes: Vec<Vec<u8>>,
@@ -234,21 +255,22 @@ impl<B: BlockT + DeserializeOwned> Builder<B> {
 	) -> Result<StorageData, &'static str> {
 		trace!(target: LOG_TARGET, "rpc: get_storage");
 		self.as_online().rpc_client().get_storage(key, maybe_at).await.map_err(|e| {
-			error!("Error = {:?}", e);
+			error!(target: LOG_TARGET, "Error = {:?}", e);
 			"rpc get_storage failed."
 		})
 	}
+
 	/// Get the latest finalized head.
 	async fn rpc_get_head(&self) -> Result<B::Hash, &'static str> {
 		trace!(target: LOG_TARGET, "rpc: finalized_head");
 		self.as_online().rpc_client().finalized_head().await.map_err(|e| {
-			error!("Error = {:?}", e);
+			error!(target: LOG_TARGET, "Error = {:?}", e);
 			"rpc finalized_head failed."
 		})
 	}
 
 	/// Get all the keys at `prefix` at `hash` using the paged, safe RPC methods.
-	async fn get_keys_paged(
+	async fn rpc_get_keys_paged(
 		&self,
 		prefix: StorageKey,
 		at: B::Hash,
@@ -277,7 +299,7 @@ impl<B: BlockT + DeserializeOwned> Builder<B> {
 					all_keys.last().expect("all_keys is populated; has .last(); qed");
 				log::debug!(
 					target: LOG_TARGET,
-					"new total = {}, full page received: {:?}",
+					"new total = {}, full page received: {}",
 					all_keys.len(),
 					HexDisplay::from(new_last_key)
 				);
@@ -296,12 +318,12 @@ impl<B: BlockT + DeserializeOwned> Builder<B> {
 		&self,
 		prefix: StorageKey,
 		at: B::Hash,
-	) -> Result<Vec<KeyPair>, &'static str> {
-		let keys = self.get_keys_paged(prefix, at).await?;
+	) -> Result<Vec<KeyValue>, &'static str> {
+		let keys = self.rpc_get_keys_paged(prefix, at).await?;
 		let keys_count = keys.len();
 		log::debug!(target: LOG_TARGET, "Querying a total of {} keys", keys.len());
 
-		let mut key_values: Vec<KeyPair> = vec![];
+		let mut key_values: Vec<KeyValue> = vec![];
 		let client = self.as_online().rpc_client();
 		for chunk_keys in keys.chunks(BATCH_SIZE) {
 			let batch = chunk_keys
@@ -318,7 +340,9 @@ impl<B: BlockT + DeserializeOwned> Builder<B> {
 				);
 				"batch failed."
 			})?;
+
 			assert_eq!(chunk_keys.len(), values.len());
+
 			for (idx, key) in chunk_keys.into_iter().enumerate() {
 				let maybe_value = values[idx].clone();
 				let value = maybe_value.unwrap_or_else(|| {
@@ -341,26 +365,216 @@ impl<B: BlockT + DeserializeOwned> Builder<B> {
 
 		Ok(key_values)
 	}
+
+	/// Get the values corresponding to `child_keys` at the given `prefixed_top_key`.
+	pub(crate) async fn rpc_child_get_storage_paged(
+		&self,
+		prefixed_top_key: &StorageKey,
+		child_keys: Vec<StorageKey>,
+		at: B::Hash,
+	) -> Result<Vec<KeyValue>, &'static str> {
+		let mut child_kv_inner = vec![];
+		for batch_child_key in child_keys.chunks(BATCH_SIZE) {
+			let batch_request = batch_child_key
+				.iter()
+				.cloned()
+				.map(|key| {
+					(
+						"childstate_getStorage",
+						rpc_params![
+							PrefixedStorageKey::new(prefixed_top_key.as_ref().to_vec()),
+							key,
+							at
+						],
+					)
+				})
+				.collect::<Vec<_>>();
+
+			let batch_response = self
+				.as_online()
+				.rpc_client()
+				.batch_request::<Option<StorageData>>(batch_request)
+				.await
+				.map_err(|e| {
+					log::error!(
+						target: LOG_TARGET,
+						"failed to execute batch: {:?}. Error: {:?}",
+						batch_child_key,
+						e
+					);
+					"batch failed."
+				})?;
+
+			assert_eq!(batch_child_key.len(), batch_response.len());
+
+			for (idx, key) in batch_child_key.into_iter().enumerate() {
+				let maybe_value = batch_response[idx].clone();
+				let value = maybe_value.unwrap_or_else(|| {
+					log::warn!(target: LOG_TARGET, "key {:?} had none corresponding value.", &key);
+					StorageData(vec![])
+				});
+				child_kv_inner.push((key.clone(), value));
+			}
+		}
+
+		Ok(child_kv_inner)
+	}
+
+	pub(crate) async fn rpc_child_get_keys(
+		&self,
+		prefixed_top_key: &StorageKey,
+		child_prefix: StorageKey,
+		at: B::Hash,
+	) -> Result<Vec<StorageKey>, &'static str> {
+		let child_keys = self
+			.as_online()
+			.rpc_client()
+			.child_get_keys(
+				PrefixedStorageKey::new(prefixed_top_key.as_ref().to_vec()),
+				child_prefix,
+				Some(at),
+			)
+			.await
+			.map_err(|e| {
+				error!(target: LOG_TARGET, "Error = {:?}", e);
+				"rpc child_get_keys failed."
+			})?;
+
+		debug!(
+			target: LOG_TARGET,
+			"scraped {} child-keys of the child-bearing top key: {}",
+			child_keys.len(),
+			HexDisplay::from(prefixed_top_key)
+		);
+
+		Ok(child_keys)
+	}
 }
 
 // Internal methods
 impl<B: BlockT + DeserializeOwned> Builder<B> {
-	/// Save the given data as state snapshot.
-	fn save_state_snapshot(&self, data: &[KeyPair], path: &Path) -> Result<(), &'static str> {
-		log::debug!(target: LOG_TARGET, "writing to state snapshot file {:?}", path);
-		fs::write(path, data.encode()).map_err(|_| "fs::write failed.")?;
+	/// Save the given data to the top keys snapshot.
+	fn save_top_snapshot(&self, data: &[KeyValue], path: &PathBuf) -> Result<(), &'static str> {
+		let mut path = path.clone();
+		let encoded = data.encode();
+		path.set_extension("top");
+		debug!(
+			target: LOG_TARGET,
+			"writing {} bytes to state snapshot file {:?}",
+			encoded.len(),
+			path
+		);
+		fs::write(path, encoded).map_err(|_| "fs::write failed.")?;
 		Ok(())
 	}
 
-	/// initialize `Self` from state snapshot. Panics if the file does not exist.
-	fn load_state_snapshot(&self, path: &Path) -> Result<Vec<KeyPair>, &'static str> {
-		log::info!(target: LOG_TARGET, "scraping key-pairs from state snapshot {:?}", path);
+	/// Save the given data to the child keys snapshot.
+	fn save_child_snapshot(
+		&self,
+		data: &ChildKeyValues,
+		path: &PathBuf,
+	) -> Result<(), &'static str> {
+		let mut path = path.clone();
+		path.set_extension("child");
+		let encoded = data.encode();
+		debug!(
+			target: LOG_TARGET,
+			"writing {} bytes to state snapshot file {:?}",
+			encoded.len(),
+			path
+		);
+		fs::write(path, encoded).map_err(|_| "fs::write failed.")?;
+		Ok(())
+	}
+
+	fn load_top_snapshot(&self, path: &PathBuf) -> Result<TopKeyValues, &'static str> {
+		let mut path = path.clone();
+		path.set_extension("top");
+		info!(target: LOG_TARGET, "loading top key-pairs from snapshot {:?}", path);
 		let bytes = fs::read(path).map_err(|_| "fs::read failed.")?;
-		Decode::decode(&mut &*bytes).map_err(|_| "decode failed")
+		Decode::decode(&mut &*bytes).map_err(|e| {
+			log::error!(target: LOG_TARGET, "{:?}", e);
+			"decode failed"
+		})
+	}
+
+	fn load_child_snapshot(&self, path: &PathBuf) -> Result<ChildKeyValues, &'static str> {
+		let mut path = path.clone();
+		path.set_extension("child");
+		info!(target: LOG_TARGET, "loading child key-pairs from snapshot {:?}", path);
+		let bytes = fs::read(path).map_err(|_| "fs::read failed.")?;
+		Decode::decode(&mut &*bytes).map_err(|e| {
+			log::error!(target: LOG_TARGET, "{:?}", e);
+			"decode failed"
+		})
+	}
+
+	/// Load all the `top` keys from the remote config, and maybe write then to cache.
+	async fn load_top_remote_and_maybe_save(&self) -> Result<TopKeyValues, &'static str> {
+		let top_kv = self.load_top_remote().await?;
+		if let Some(c) = &self.as_online().state_snapshot {
+			self.save_top_snapshot(&top_kv, &c.path)?;
+		}
+		Ok(top_kv)
+	}
+
+	/// Load all of the child keys from the remote config, given the already scraped list of top key
+	/// pairs.
+	///
+	/// Stores all values to cache as well, if provided.
+	async fn load_child_remote_and_maybe_save(
+		&self,
+		top_kv: &[KeyValue],
+	) -> Result<ChildKeyValues, &'static str> {
+		let child_kv = self.load_child_remote(&top_kv).await?;
+		if let Some(c) = &self.as_online().state_snapshot {
+			self.save_child_snapshot(&child_kv, &c.path)?;
+		}
+		Ok(child_kv)
+	}
+
+	/// Load all of the child keys from the remote config, given the already scraped list of top key
+	/// pairs.
+	///
+	/// `top_kv` need not be only child-bearing top keys. It should be all of the top keys that are
+	/// included thus far.
+	async fn load_child_remote(&self, top_kv: &[KeyValue]) -> Result<ChildKeyValues, &'static str> {
+		let child_roots = top_kv
+			.iter()
+			.filter_map(|(k, _)| is_default_child_storage_key(k.as_ref()).then(|| k))
+			.collect::<Vec<_>>();
+
+		info!(
+			target: LOG_TARGET,
+			"👩‍👦 scraping child-tree data from {} top keys",
+			child_roots.len()
+		);
+
+		let mut child_kv = vec![];
+		for prefixed_top_key in child_roots {
+			let at = self.as_online().at.expect("at must be initialized in online mode.");
+			let child_keys =
+				self.rpc_child_get_keys(prefixed_top_key, StorageKey(vec![]), at).await?;
+			let child_kv_inner =
+				self.rpc_child_get_storage_paged(prefixed_top_key, child_keys, at).await?;
+
+			let prefixed_top_key = PrefixedStorageKey::new(prefixed_top_key.clone().0);
+			let un_prefixed = match ChildType::from_prefixed_key(&prefixed_top_key) {
+				Some((ChildType::ParentKeyId, storage_key)) => storage_key,
+				None => {
+					log::error!(target: LOG_TARGET, "invalid key: {:?}", prefixed_top_key);
+					return Err("Invalid child key")
+				},
+			};
+
+			child_kv.push((ChildInfo::new_default(&un_prefixed), child_kv_inner));
+		}
+
+		Ok(child_kv)
 	}
 
 	/// Build `Self` from a network node denoted by `uri`.
-	async fn load_remote(&self) -> Result<Vec<KeyPair>, &'static str> {
+	async fn load_top_remote(&self) -> Result<TopKeyValues, &'static str> {
 		let config = self.as_online();
 		let at = self
 			.as_online()
@@ -371,17 +585,17 @@ impl<B: BlockT + DeserializeOwned> Builder<B> {
 
 		let mut keys_and_values = if config.pallets.len() > 0 {
 			let mut filtered_kv = vec![];
-			for f in config.pallets.iter() {
-				let hashed_prefix = StorageKey(twox_128(f.as_bytes()).to_vec());
-				let module_kv = self.rpc_get_pairs_paged(hashed_prefix.clone(), at).await?;
+			for p in config.pallets.iter() {
+				let hashed_prefix = StorageKey(twox_128(p.as_bytes()).to_vec());
+				let pallet_kv = self.rpc_get_pairs_paged(hashed_prefix.clone(), at).await?;
 				log::info!(
 					target: LOG_TARGET,
-					"downloaded data for module {} (count: {} / prefix: {:?}).",
-					f,
-					module_kv.len(),
+					"downloaded data for module {} (count: {} / prefix: {}).",
+					p,
+					pallet_kv.len(),
 					HexDisplay::from(&hashed_prefix),
 				);
-				filtered_kv.extend(module_kv);
+				filtered_kv.extend(pallet_kv);
 			}
 			filtered_kv
 		} else {
@@ -423,7 +637,10 @@ impl<B: BlockT + DeserializeOwned> Builder<B> {
 			.max_request_body_size(u32::MAX)
 			.build(&online.transport.uri)
 			.await
-			.map_err(|_| "failed to build ws client")?;
+			.map_err(|e| {
+				log::error!(target: LOG_TARGET, "error: {:?}", e);
+				"failed to build ws client"
+			})?;
 		online.transport.client = Some(ws_client);
 
 		// Then, if `at` is not set, set it.
@@ -435,27 +652,21 @@ impl<B: BlockT + DeserializeOwned> Builder<B> {
 		Ok(())
 	}
 
-	pub(crate) async fn pre_build(mut self) -> Result<Vec<KeyPair>, &'static str> {
-		let mut base_kv = match self.mode.clone() {
-			Mode::Offline(config) => self.load_state_snapshot(&config.state_snapshot.path)?,
-			Mode::Online(config) => {
+	pub(crate) async fn pre_build(
+		mut self,
+	) -> Result<(TopKeyValues, ChildKeyValues), &'static str> {
+		let mut top_kv = match self.mode.clone() {
+			Mode::Offline(config) => self.load_top_snapshot(&config.state_snapshot.path)?,
+			Mode::Online(_) => {
 				self.init_remote_client().await?;
-				let kp = self.load_remote().await?;
-				if let Some(c) = config.state_snapshot {
-					self.save_state_snapshot(&kp, &c.path)?;
-				}
-				kp
+				self.load_top_remote_and_maybe_save().await?
 			},
-			Mode::OfflineOrElseOnline(offline_config, online_config) => {
-				if let Ok(kv) = self.load_state_snapshot(&offline_config.state_snapshot.path) {
+			Mode::OfflineOrElseOnline(offline_config, _) => {
+				if let Ok(kv) = self.load_top_snapshot(&offline_config.state_snapshot.path) {
 					kv
 				} else {
 					self.init_remote_client().await?;
-					let kp = self.load_remote().await?;
-					if let Some(c) = online_config.state_snapshot {
-						self.save_state_snapshot(&kp, &c.path)?;
-					}
-					kp
+					self.load_top_remote_and_maybe_save().await?
 				}
 			},
 		};
@@ -467,7 +678,7 @@ impl<B: BlockT + DeserializeOwned> Builder<B> {
 				"extending externalities with {} manually injected key-values",
 				self.hashed_key_values.len()
 			);
-			base_kv.extend(self.hashed_key_values.clone());
+			top_kv.extend(self.hashed_key_values.clone());
 		}
 
 		// exclude manual key values.
@@ -477,10 +688,30 @@ impl<B: BlockT + DeserializeOwned> Builder<B> {
 				"excluding externalities from {} keys",
 				self.hashed_blacklist.len()
 			);
-			base_kv.retain(|(k, _)| !self.hashed_blacklist.contains(&k.0))
+			top_kv.retain(|(k, _)| !self.hashed_blacklist.contains(&k.0))
 		}
 
-		Ok(base_kv)
+		let child_kv = match self.mode.clone() {
+			Mode::Online(_) => self.load_child_remote_and_maybe_save(&top_kv).await?,
+			Mode::OfflineOrElseOnline(offline_config, _) =>
+				if let Ok(kv) = self.load_child_snapshot(&offline_config.state_snapshot.path) {
+					kv
+				} else {
+					self.load_child_remote_and_maybe_save(&top_kv).await?
+				},
+			Mode::Offline(ref config) => self
+				.load_child_snapshot(&config.state_snapshot.path)
+				.map_err(|why| {
+					log::warn!(
+						target: LOG_TARGET,
+						"failed to load child-key file due to {:?}.",
+						why
+					)
+				})
+				.unwrap_or_default(),
+		};
+
+		Ok((top_kv, child_kv))
 	}
 }
 
@@ -492,12 +723,13 @@ impl<B: BlockT + DeserializeOwned> Builder<B> {
 	}
 
 	/// Inject a manual list of key and values to the storage.
-	pub fn inject_hashed_key_value(mut self, injections: &[KeyPair]) -> Self {
+	pub fn inject_hashed_key_value(mut self, injections: &[KeyValue]) -> Self {
 		for i in injections {
 			self.hashed_key_values.push(i.clone());
 		}
 		self
 	}
+
 	/// Inject a hashed prefix. This is treated as-is, and should be pre-hashed.
 	///
 	/// This should be used to inject a "PREFIX", like a storage (double) map.
@@ -506,6 +738,22 @@ impl<B: BlockT + DeserializeOwned> Builder<B> {
 		self
 	}
 
+	/// Just a utility wrapper of [`inject_hashed_prefix`] that injects
+	/// [`DEFAULT_CHILD_STORAGE_KEY_PREFIX`] as a prefix.
+	///
+	/// If set, this will guarantee that the child-tree data of ALL pallets will be downloaded.
+	///
+	/// This is not needed if the entire state is being downloaded.
+	///
+	/// Otherwise, the only other way to make sure a child-tree is manually included is to inject
+	/// its root (`DEFAULT_CHILD_STORAGE_KEY_PREFIX`, plus some other postfix) into
+	/// [`inject_hashed_key`]. Unfortunately, there's no federated way of managing child tree roots
+	/// as of now and each pallet does its own thing. Therefore, it is not possible for this library
+	/// to automatically include child trees of pallet X, when its top keys are included.
+	pub fn inject_default_child_tree_prefix(self) -> Self {
+		self.inject_hashed_prefix(DEFAULT_CHILD_STORAGE_KEY_PREFIX)
+	}
+
 	/// Inject a hashed key to scrape. This is treated as-is, and should be pre-hashed.
 	///
 	/// This should be used to inject a "KEY", like a storage value.
@@ -540,16 +788,37 @@ impl<B: BlockT + DeserializeOwned> Builder<B> {
 
 	/// Build the test externalities.
 	pub async fn build(self) -> Result<TestExternalities, &'static str> {
-		let kv = self.pre_build().await?;
-		let mut ext = TestExternalities::new_empty();
+		let (top_kv, child_kv) = self.pre_build().await?;
+		let mut ext = TestExternalities::new_with_code(Default::default(), Default::default());
 
-		log::info!(target: LOG_TARGET, "injecting a total of {} keys", kv.len());
-		for (k, v) in kv {
-			let (k, v) = (k.0, v.0);
-			// Insert the key,value pair into the test trie backend
-			ext.insert(k, v);
+		info!(target: LOG_TARGET, "injecting a total of {} top keys", top_kv.len());
+		for (k, v) in top_kv {
+			// skip writing the child root data.
+			if is_default_child_storage_key(k.as_ref()) {
+				continue
+			}
+			ext.insert(k.0, v.0);
 		}
 
+		info!(
+			target: LOG_TARGET,
+			"injecting a total of {} child keys",
+			child_kv.iter().map(|(_, kv)| kv).flatten().count()
+		);
+
+		for (info, key_values) in child_kv {
+			for (k, v) in key_values {
+				ext.insert_child(info.clone(), k.0, v.0);
+			}
+		}
+
+		ext.commit_all().unwrap();
+		info!(
+			target: LOG_TARGET,
+			"initialized state externalities with storage root {:?}",
+			ext.as_backend().root()
+		);
+
 		Ok(ext)
 	}
 }
@@ -621,7 +890,6 @@ mod tests {
 #[cfg(all(test, feature = "remote-test"))]
 mod remote_tests {
 	use super::test_prelude::*;
-	use pallet_elections_phragmen::Members;
 	const REMOTE_INACCESSIBLE: &'static str = "Can't reach the remote node. Is it running?";
 
 	#[tokio::test]
@@ -631,11 +899,11 @@ mod remote_tests {
 		Builder::<Block>::new()
 			.mode(Mode::OfflineOrElseOnline(
 				OfflineConfig {
-					state_snapshot: SnapshotConfig::new("test_snapshot_to_remove.bin"),
+					state_snapshot: SnapshotConfig::new("offline_else_online_works_data"),
 				},
 				OnlineConfig {
 					pallets: vec!["Proxy".to_owned()],
-					state_snapshot: Some(SnapshotConfig::new("test_snapshot_to_remove.bin")),
+					state_snapshot: Some(SnapshotConfig::new("offline_else_online_works_data")),
 					..Default::default()
 				},
 			))
@@ -648,11 +916,11 @@ mod remote_tests {
 		Builder::<Block>::new()
 			.mode(Mode::OfflineOrElseOnline(
 				OfflineConfig {
-					state_snapshot: SnapshotConfig::new("test_snapshot_to_remove.bin"),
+					state_snapshot: SnapshotConfig::new("offline_else_online_works_data"),
 				},
 				OnlineConfig {
 					pallets: vec!["Proxy".to_owned()],
-					state_snapshot: Some(SnapshotConfig::new("test_snapshot_to_remove.bin")),
+					state_snapshot: Some(SnapshotConfig::new("offline_else_online_works_data")),
 					transport: "ws://non-existent:666".to_owned().into(),
 					..Default::default()
 				},
@@ -661,14 +929,56 @@ mod remote_tests {
 			.await
 			.expect(REMOTE_INACCESSIBLE)
 			.execute_with(|| {});
+
+		let to_delete = std::fs::read_dir(Path::new("."))
+			.unwrap()
+			.into_iter()
+			.map(|d| d.unwrap())
+			.filter(|p| {
+				p.path().file_name().unwrap_or_default() == "offline_else_online_works_data" ||
+					p.path().extension().unwrap_or_default() == "top" ||
+					p.path().extension().unwrap_or_default() == "child"
+			})
+			.collect::<Vec<_>>();
+		assert!(to_delete.len() > 0);
+		for d in to_delete {
+			std::fs::remove_file(d.path()).unwrap();
+		}
 	}
 
 	#[tokio::test]
-	async fn can_build_one_pallet() {
+	#[ignore = "too slow"]
+	async fn can_build_one_big_pallet() {
 		init_logger();
 		Builder::<Block>::new()
 			.mode(Mode::Online(OnlineConfig {
-				pallets: vec!["Proxy".to_owned()],
+				pallets: vec!["System".to_owned()],
+				..Default::default()
+			}))
+			.build()
+			.await
+			.expect(REMOTE_INACCESSIBLE)
+			.execute_with(|| {});
+	}
+
+	#[tokio::test]
+	async fn can_build_one_small_pallet() {
+		init_logger();
+		Builder::<Block>::new()
+			.mode(Mode::Online(OnlineConfig {
+				transport: "wss://kusama-rpc.polkadot.io:443".to_owned().into(),
+				pallets: vec!["Council".to_owned()],
+				..Default::default()
+			}))
+			.build()
+			.await
+			.expect(REMOTE_INACCESSIBLE)
+			.execute_with(|| {});
+
+		Builder::<Block>::new()
+			.mode(Mode::Online(OnlineConfig {
+				transport: "wss://rpc.polkadot.io:443".to_owned().into(),
+				pallets: vec!["Council".to_owned()],
 				..Default::default()
 			}))
 			.build()
@@ -682,6 +992,18 @@ mod remote_tests {
 		init_logger();
 		Builder::<Block>::new()
 			.mode(Mode::Online(OnlineConfig {
+				transport: "wss://kusama-rpc.polkadot.io:443".to_owned().into(),
+				pallets: vec!["Proxy".to_owned(), "Multisig".to_owned()],
+				..Default::default()
+			}))
+			.build()
+			.await
+			.expect(REMOTE_INACCESSIBLE)
+			.execute_with(|| {});
+
+		Builder::<Block>::new()
+			.mode(Mode::Online(OnlineConfig {
+				transport: "wss://rpc.polkadot.io:443".to_owned().into(),
 				pallets: vec!["Proxy".to_owned(), "Multisig".to_owned()],
 				..Default::default()
 			}))
@@ -692,46 +1014,11 @@ mod remote_tests {
 	}
 
 	#[tokio::test]
-	async fn sanity_check_decoding() {
-		use sp_core::crypto::Ss58Codec;
-
-		type AccountId = sp_runtime::AccountId32;
-		type Balance = u128;
-		frame_support::generate_storage_alias!(
-			PhragmenElection,
-			Members =>
-			Value<Vec<SeatHolder<AccountId, Balance>>>
-		);
-
+	async fn can_create_top_snapshot() {
 		init_logger();
 		Builder::<Block>::new()
 			.mode(Mode::Online(OnlineConfig {
-				pallets: vec!["PhragmenElection".to_owned()],
-				..Default::default()
-			}))
-			.build()
-			.await
-			.expect(REMOTE_INACCESSIBLE)
-			.execute_with(|| {
-				// Gav's polkadot account. 99% this will be in the council.
-				let gav_polkadot =
-					AccountId::from_ss58check("13RDY9nrJpyTDBSUdBw12dGwhk19sGwsrVZ2bxkzYHBSagP2")
-						.unwrap();
-				let members = Members::get();
-				assert!(members
-					.iter()
-					.map(|s| s.who.clone())
-					.find(|a| a == &gav_polkadot)
-					.is_some());
-			});
-	}
-
-	#[tokio::test]
-	async fn can_create_state_snapshot() {
-		init_logger();
-		Builder::<Block>::new()
-			.mode(Mode::Online(OnlineConfig {
-				state_snapshot: Some(SnapshotConfig::new("test_snapshot_to_remove.bin")),
+				state_snapshot: Some(SnapshotConfig::new("can_create_top_snapshot_data")),
 				pallets: vec!["Proxy".to_owned()],
 				..Default::default()
 			}))
@@ -744,24 +1031,118 @@ mod remote_tests {
 			.unwrap()
 			.into_iter()
 			.map(|d| d.unwrap())
-			.filter(|p| p.path().extension().unwrap_or_default() == "bin")
+			.filter(|p| {
+				p.path().file_name().unwrap_or_default() == "can_create_top_snapshot_data" ||
+					p.path().extension().unwrap_or_default() == "top" ||
+					p.path().extension().unwrap_or_default() == "child"
+			})
 			.collect::<Vec<_>>();
 
 		assert!(to_delete.len() > 0);
 
 		for d in to_delete {
+			use std::os::unix::fs::MetadataExt;
+			if d.path().extension().unwrap_or_default() == "top" {
+				// if this is the top snapshot it must not be empty.
+				assert!(std::fs::metadata(d.path()).unwrap().size() > 1);
+			} else {
+				// the child is empty for this pallet.
+				assert!(std::fs::metadata(d.path()).unwrap().size() == 1);
+			}
 			std::fs::remove_file(d.path()).unwrap();
 		}
 	}
 
 	#[tokio::test]
-	#[ignore = "takes too much time on average."]
-	async fn can_fetch_all() {
+	async fn can_build_child_tree() {
 		init_logger();
 		Builder::<Block>::new()
+			.mode(Mode::Online(OnlineConfig {
+				transport: "wss://rpc.polkadot.io:443".to_owned().into(),
+				pallets: vec!["Crowdloan".to_owned()],
+				..Default::default()
+			}))
 			.build()
 			.await
 			.expect(REMOTE_INACCESSIBLE)
 			.execute_with(|| {});
 	}
+
+	#[tokio::test]
+	async fn can_create_child_snapshot() {
+		init_logger();
+		Builder::<Block>::new()
+			.mode(Mode::Online(OnlineConfig {
+				state_snapshot: Some(SnapshotConfig::new("can_create_child_snapshot_data")),
+				pallets: vec!["Crowdloan".to_owned()],
+				..Default::default()
+			}))
+			.inject_default_child_tree_prefix()
+			.build()
+			.await
+			.expect(REMOTE_INACCESSIBLE)
+			.execute_with(|| {});
+
+		let to_delete = std::fs::read_dir(Path::new("."))
+			.unwrap()
+			.into_iter()
+			.map(|d| d.unwrap())
+			.filter(|p| {
+				p.path().file_name().unwrap_or_default() == "can_create_child_snapshot_data" ||
+					p.path().extension().unwrap_or_default() == "top" ||
+					p.path().extension().unwrap_or_default() == "child"
+			})
+			.collect::<Vec<_>>();
+
+		assert!(to_delete.len() > 0);
+
+		for d in to_delete {
+			use std::os::unix::fs::MetadataExt;
+			// if this is the top snapshot it must not be empty
+			if d.path().extension().unwrap_or_default() == "child" {
+				assert!(std::fs::metadata(d.path()).unwrap().size() > 1);
+			} else {
+				assert!(std::fs::metadata(d.path()).unwrap().size() > 1);
+			}
+			std::fs::remove_file(d.path()).unwrap();
+		}
+	}
+
+	#[tokio::test]
+	async fn can_fetch_all() {
+		init_logger();
+		Builder::<Block>::new()
+			.mode(Mode::Online(OnlineConfig {
+				state_snapshot: Some(SnapshotConfig::new("can_fetch_all_data")),
+				..Default::default()
+			}))
+			.build()
+			.await
+			.expect(REMOTE_INACCESSIBLE)
+			.execute_with(|| {});
+
+		let to_delete = std::fs::read_dir(Path::new("."))
+			.unwrap()
+			.into_iter()
+			.map(|d| d.unwrap())
+			.filter(|p| {
+				p.path().file_name().unwrap_or_default() == "can_fetch_all_data" ||
+					p.path().extension().unwrap_or_default() == "top" ||
+					p.path().extension().unwrap_or_default() == "child"
+			})
+			.collect::<Vec<_>>();
+
+		assert!(to_delete.len() > 0);
+
+		for d in to_delete {
+			use std::os::unix::fs::MetadataExt;
+			// if we download everything, child tree must also be filled.
+			if d.path().extension().unwrap_or_default() == "child" {
+				assert!(std::fs::metadata(d.path()).unwrap().size() > 1);
+			} else {
+				assert!(std::fs::metadata(d.path()).unwrap().size() > 1);
+			}
+			std::fs::remove_file(d.path()).unwrap();
+		}
+	}
 }

substrate/utils/frame/try-runtime/cli/Cargo.toml

@@ -32,3 +32,5 @@ sp-version = { version = "4.0.0-dev", path = "../../../../primitives/version" }
 
 remote-externalities = { version = "0.10.0-dev", path = "../../remote-externalities" }
 jsonrpsee = { version = "0.4.1", default-features = false, features = ["ws-client"]}
+
+zstd = "0.9.0"

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

@@ -17,7 +17,7 @@
 
 use crate::{
 	build_executor, ensure_matching_spec, extract_code, full_extensions, hash_of, local_spec,
-	state_machine_call, SharedParams, State, LOG_TARGET,
+	state_machine_call_with_proof, SharedParams, State, LOG_TARGET,
 };
 use remote_externalities::rpc_api;
 use sc_service::{Configuration, NativeExecutionDispatch};
@@ -167,7 +167,7 @@ where
 	)
 	.await;
 
-	let _ = state_machine_call::<Block, ExecDispatch>(
+	let _ = state_machine_call_with_proof::<Block, ExecDispatch>(
 		&ext,
 		&executor,
 		execution,

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

@@ -17,7 +17,7 @@
 
 use crate::{
 	build_executor, ensure_matching_spec, extract_code, full_extensions, local_spec, parse,
-	state_machine_call, SharedParams, LOG_TARGET,
+	state_machine_call_with_proof, SharedParams, LOG_TARGET,
 };
 use jsonrpsee::{
 	types::{traits::SubscriptionClient, Subscription},
@@ -139,7 +139,7 @@ where
 		let state_ext =
 			maybe_state_ext.as_mut().expect("state_ext either existed or was just created");
 
-		let (mut changes, encoded_result) = state_machine_call::<Block, ExecDispatch>(
+		let (mut changes, encoded_result) = state_machine_call_with_proof::<Block, ExecDispatch>(
 			&state_ext,
 			&executor,
 			execution,

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

@@ -23,7 +23,7 @@ use sc_service::Configuration;
 use sp_runtime::traits::{Block as BlockT, NumberFor};
 
 use crate::{
-	build_executor, ensure_matching_spec, extract_code, local_spec, state_machine_call,
+	build_executor, ensure_matching_spec, extract_code, local_spec, state_machine_call_with_proof,
 	SharedParams, State, LOG_TARGET,
 };
 
@@ -69,7 +69,7 @@ where
 		.await;
 	}
 
-	let (_, encoded_result) = state_machine_call::<Block, ExecDispatch>(
+	let (_, encoded_result) = state_machine_call_with_proof::<Block, ExecDispatch>(
 		&ext,
 		&executor,
 		execution,

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

@@ -285,7 +285,7 @@ use sp_runtime::{
 	traits::{Block as BlockT, NumberFor},
 	DeserializeOwned,
 };
-use sp_state_machine::{OverlayedChanges, StateMachine};
+use sp_state_machine::{InMemoryProvingBackend, OverlayedChanges, StateMachine};
 use std::{fmt::Debug, path::PathBuf, str::FromStr};
 
 mod commands;
@@ -462,6 +462,14 @@ pub enum State {
 		/// The pallets to scrape. If empty, entire chain state will be scraped.
 		#[structopt(short, long, require_delimiter = true)]
 		pallets: Option<Vec<String>>,
+
+		/// Fetch the child-keys as well.
+		///
+		/// Default is `false`, if specific `pallets` are specified, true otherwise. In other
+		/// words, if you scrape the whole state the child tree data is included out of the box.
+		/// Otherwise, it must be enabled explicitly using this flag.
+		#[structopt(long, require_delimiter = true)]
+		child_tree: bool,
 	},
 }
 
@@ -477,21 +485,26 @@ impl State {
 				Builder::<Block>::new().mode(Mode::Offline(OfflineConfig {
 					state_snapshot: SnapshotConfig::new(snapshot_path),
 				})),
-			State::Live { snapshot_path, pallets, uri, at } => {
+			State::Live { snapshot_path, pallets, uri, at, child_tree } => {
 				let at = match at {
 					Some(at_str) => Some(hash_of::<Block>(at_str)?),
 					None => None,
 				};
-				Builder::<Block>::new()
+				let mut builder = Builder::<Block>::new()
 					.mode(Mode::Online(OnlineConfig {
 						transport: uri.to_owned().into(),
 						state_snapshot: snapshot_path.as_ref().map(SnapshotConfig::new),
-						pallets: pallets.to_owned().unwrap_or_default(),
+						pallets: pallets.clone().unwrap_or_default(),
 						at,
+						..Default::default()
 					}))
 					.inject_hashed_key(
 						&[twox_128(b"System"), twox_128(b"LastRuntimeUpgrade")].concat(),
-					)
+					);
+				if *child_tree {
+					builder = builder.inject_default_child_tree_prefix();
+				}
+				builder
 			},
 		})
 	}
@@ -697,6 +710,85 @@ pub(crate) fn state_machine_call<Block: BlockT, D: NativeExecutionDispatch + 'st
 	Ok((changes, encoded_results))
 }
 
+/// Same as [`state_machine_call`], but it also computes and prints the storage proof in different
+/// size and formats.
+///
+/// Make sure [`LOG_TARGET`] is enabled in logging.
+pub(crate) fn state_machine_call_with_proof<Block: BlockT, D: NativeExecutionDispatch + 'static>(
+	ext: &TestExternalities,
+	executor: &NativeElseWasmExecutor<D>,
+	execution: sc_cli::ExecutionStrategy,
+	method: &'static str,
+	data: &[u8],
+	extensions: Extensions,
+) -> sc_cli::Result<(OverlayedChanges, Vec<u8>)> {
+	use parity_scale_codec::Encode;
+	use sp_core::hexdisplay::HexDisplay;
+
+	let mut changes = Default::default();
+	let backend = ext.backend.clone();
+	let proving_backend = InMemoryProvingBackend::new(&backend);
+
+	let runtime_code_backend = sp_state_machine::backend::BackendRuntimeCode::new(&proving_backend);
+	let runtime_code = runtime_code_backend.runtime_code()?;
+
+	let pre_root = backend.root().clone();
+
+	let encoded_results = StateMachine::new(
+		&proving_backend,
+		&mut changes,
+		executor,
+		method,
+		data,
+		extensions,
+		&runtime_code,
+		sp_core::testing::TaskExecutor::new(),
+	)
+	.execute(execution.into())
+	.map_err(|e| format!("failed to execute {}: {:?}", method, e))
+	.map_err::<sc_cli::Error, _>(Into::into)?;
+
+	let proof = proving_backend.extract_proof();
+	let proof_size = proof.encoded_size();
+	let compact_proof = proof
+		.clone()
+		.into_compact_proof::<sp_runtime::traits::BlakeTwo256>(pre_root)
+		.map_err(|e| format!("failed to generate compact proof {}: {:?}", method, e))?;
+
+	let compact_proof_size = compact_proof.encoded_size();
+	let compressed_proof = zstd::stream::encode_all(&compact_proof.encode()[..], 0)
+		.map_err(|e| format!("failed to generate compact proof {}: {:?}", method, e))?;
+
+	let proof_nodes = proof.into_nodes();
+
+	let humanize = |s| {
+		if s < 1024 * 1024 {
+			format!("{:.2} KB ({} bytes)", s as f64 / 1024f64, s)
+		} else {
+			format!(
+				"{:.2} MB ({} KB) ({} bytes)",
+				s as f64 / (1024f64 * 1024f64),
+				s as f64 / 1024f64,
+				s
+			)
+		}
+	};
+	log::info!(
+		target: LOG_TARGET,
+		"proof: {} / {} nodes",
+		HexDisplay::from(&proof_nodes.iter().flatten().cloned().collect::<Vec<_>>()),
+		proof_nodes.len()
+	);
+	log::info!(target: LOG_TARGET, "proof size: {}", humanize(proof_size));
+	log::info!(target: LOG_TARGET, "compact proof size: {}", humanize(compact_proof_size),);
+	log::info!(
+		target: LOG_TARGET,
+		"zstd-compressed compact proof {}",
+		humanize(compressed_proof.len()),
+	);
+	Ok((changes, encoded_results))
+}
+
 /// Get the spec `(name, version)` from the local runtime.
 pub(crate) fn local_spec<Block: BlockT, D: NativeExecutionDispatch + 'static>(
 	ext: &TestExternalities,