allow try-runtime and TestExternalities to report PoV size

primitives/state-machine/Cargo.toml

@@ -34,8 +34,10 @@ tracing = { version = "0.1.29", optional = true }
 [dev-dependencies]
 hex-literal = "0.3.4"
 sp-runtime = { version = "4.0.0-dev", path = "../runtime" }
+sp-io = { version = "4.0.0-dev", path = "../io" }
 pretty_assertions = "1.0.0"
 rand = "0.7.2"
+zstd = "0.9"
 
 [features]
 default = ["std"]

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>
@@ -69,6 +70,21 @@ where
 		)
 	}
 
+	/// Get an externalities implementation, using the given `proving_backend`.
+	///
+	/// This will be capable of computing the PoV. See [`execute_and_get_proof`].
+	pub fn proving_ext<'a>(
+		&'a mut self,
+		proving_backend: &'a InMemoryProvingBackend<'a, H>,
+	) -> Ext<H, InMemoryProvingBackend<'a, H>> {
+		Ext::new(
+			&mut self.overlay,
+			&mut self.storage_transaction_cache,
+			&proving_backend,
+			Some(&mut self.extensions),
+		)
+	}
+
 	/// Create a new instance of `TestExternalities` with storage.
 	pub fn new(storage: Storage) -> Self {
 		Self::new_with_code(&[], storage)
@@ -122,6 +138,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 +194,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,
@@ -275,6 +318,7 @@ where
 #[cfg(test)]
 mod tests {
 	use super::*;
+	use crate::create_proof_check_backend;
 	use hex_literal::hex;
 	use sp_core::{storage::ChildInfo, traits::Externalities, H256};
 	use sp_runtime::traits::BlakeTwo256;
@@ -348,4 +392,49 @@ mod tests {
 		ext.commit_all().unwrap();
 		assert!(ext.backend.eq(&backend), "Both backend should be equal.");
 	}
+
+	#[test]
+	fn execute_and_generate_proof_works() {
+		use codec::Encode;
+		let mut ext = TestExternalities::<BlakeTwo256>::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::<sp_runtime::traits::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/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);