allow try-runtime and TestExternalities to report PoV size (parityt…

…ech#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 <[email protected]>

* last touches

* fix

Co-authored-by: Bastian Köcher <[email protected]>

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 = []

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]);
+		});
+	}
 }

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 {

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,

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,

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>,

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);