State machine call proof backend (#3945)

* drafting a proof extraction at the hashdb level (to include everything
for the 'call' proof case).

* use full proof by default (replace previous proof recorder).

* fix warnings.

* Cache value not found in proof recorder.

* Remove need or internal backend struct.

* fix type.

* doc

substrate/core/client/src/call_executor.rs

@@ -14,7 +14,7 @@
 // You should have received a copy of the GNU General Public License
 // along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
 
-use std::{sync::Arc, cmp::Ord, panic::UnwindSafe, result, cell::RefCell, rc::Rc};
+use std::{sync::Arc, cmp::Ord, panic::UnwindSafe, result, cell::RefCell};
 use codec::{Encode, Decode};
 use sr_primitives::{
 	generic::BlockId, traits::Block as BlockT, traits::NumberFor,
@@ -81,7 +81,7 @@ where
 		execution_manager: ExecutionManager<EM>,
 		native_call: Option<NC>,
 		side_effects_handler: Option<OffchainExt>,
-		proof_recorder: &Option<Rc<RefCell<ProofRecorder<B>>>>,
+		proof_recorder: &Option<ProofRecorder<B>>,
 		enable_keystore: bool,
 	) -> error::Result<NativeOrEncoded<R>> where ExecutionManager<EM>: Clone;
 
@@ -241,7 +241,7 @@ where
 		execution_manager: ExecutionManager<EM>,
 		native_call: Option<NC>,
 		side_effects_handler: Option<OffchainExt>,
-		recorder: &Option<Rc<RefCell<ProofRecorder<Block>>>>,
+		recorder: &Option<ProofRecorder<Block>>,
 		enable_keystore: bool,
 	) -> Result<NativeOrEncoded<R>, error::Error> where ExecutionManager<EM>: Clone {
 		match initialize_block {

substrate/core/client/src/client.rs

@@ -18,7 +18,7 @@
 
 use std::{
 	marker::PhantomData, collections::{HashSet, BTreeMap, HashMap}, sync::Arc,
-	panic::UnwindSafe, result, cell::RefCell, rc::Rc,
+	panic::UnwindSafe, result, cell::RefCell,
 };
 use log::{info, trace, warn};
 use futures03::channel::mpsc;
@@ -1494,7 +1494,7 @@ impl<B, E, Block, RA> CallRuntimeAt<Block> for Client<B, E, Block, RA> where
 		initialize_block: InitializeBlock<'a, Block>,
 		native_call: Option<NC>,
 		context: ExecutionContext,
-		recorder: &Option<Rc<RefCell<ProofRecorder<Block>>>>,
+		recorder: &Option<ProofRecorder<Block>>,
 	) -> error::Result<NativeOrEncoded<R>> {
 		let manager = match context {
 			ExecutionContext::BlockConstruction =>

substrate/core/client/src/light/call_executor.rs

@@ -16,7 +16,9 @@
 
 //! Methods that light client could use to execute runtime calls.
 
-use std::{sync::Arc, panic::UnwindSafe, result, cell::RefCell, rc::Rc};
+use std::{
+	sync::Arc, panic::UnwindSafe, result, cell::RefCell,
+};
 
 use codec::{Encode, Decode};
 use primitives::{
@@ -108,7 +110,7 @@ impl<Block, B, Local> CallExecutor<Block, Blake2Hasher> for
 		_manager: ExecutionManager<EM>,
 		native_call: Option<NC>,
 		side_effects_handler: Option<OffchainExt>,
-		recorder: &Option<Rc<RefCell<ProofRecorder<Block>>>>,
+		recorder: &Option<ProofRecorder<Block>>,
 		enable_keystore: bool,
 	) -> ClientResult<NativeOrEncoded<R>> where ExecutionManager<EM>: Clone {
 		// there's no actual way/need to specify native/wasm execution strategy on light node
@@ -334,7 +336,7 @@ mod tests {
 			_execution_manager: ExecutionManager<EM>,
 			_native_call: Option<NC>,
 			_side_effects_handler: Option<OffchainExt>,
-			_proof_recorder: &Option<Rc<RefCell<ProofRecorder<Block>>>>,
+			_proof_recorder: &Option<ProofRecorder<Block>>,
 			_enable_keystore: bool,
 		) -> ClientResult<NativeOrEncoded<R>> where ExecutionManager<EM>: Clone {
 			unreachable!()

substrate/core/sr-api/proc-macro/src/decl_runtime_apis.rs

@@ -410,7 +410,7 @@ fn generate_call_api_at_calls(decl: &ItemTrait) -> Result<TokenStream> {
 				initialized_block: &std::cell::RefCell<Option<#crate_::BlockId<Block>>>,
 				native_call: Option<NC>,
 				context: #crate_::ExecutionContext,
-				recorder: &Option<std::rc::Rc<std::cell::RefCell<#crate_::ProofRecorder<Block>>>>,
+				recorder: &Option<#crate_::ProofRecorder<Block>>,
 			) -> std::result::Result<#crate_::NativeOrEncoded<R>, T::Error> {
 				let version = call_runtime_at.runtime_version_at(at)?;
 				use #crate_::InitializeBlock;

substrate/core/sr-api/proc-macro/src/impl_runtime_apis.rs

@@ -257,7 +257,7 @@ fn generate_runtime_api_base_structures(impls: &[ItemImpl]) -> Result<TokenStrea
 			commit_on_success: std::cell::RefCell<bool>,
 			initialized_block: std::cell::RefCell<Option<#block_id>>,
 			changes: std::cell::RefCell<#crate_::OverlayedChanges>,
-			recorder: Option<std::rc::Rc<std::cell::RefCell<#crate_::ProofRecorder<#block>>>>,
+			recorder: Option<#crate_::ProofRecorder<#block>>,
 		}
 
 		// `RuntimeApi` itself is not threadsafe. However, an instance is only available in a
@@ -300,11 +300,9 @@ fn generate_runtime_api_base_structures(impls: &[ItemImpl]) -> Result<TokenStrea
 				self.recorder
 					.take()
 					.map(|recorder| {
-						let trie_nodes = recorder
-							.borrow_mut()
-							.drain()
-							.into_iter()
-							.map(|record| record.data)
+						let trie_nodes = recorder.read()
+							.iter()
+							.filter_map(|(_k, v)| v.as_ref().map(|v| v.to_vec()))
 							.collect();
 						#crate_::StorageProof::new(trie_nodes)
 					})
@@ -339,7 +337,7 @@ fn generate_runtime_api_base_structures(impls: &[ItemImpl]) -> Result<TokenStrea
 					&Self,
 					&std::cell::RefCell<#crate_::OverlayedChanges>,
 					&std::cell::RefCell<Option<#crate_::BlockId<#block>>>,
-					&Option<std::rc::Rc<std::cell::RefCell<#crate_::ProofRecorder<#block>>>>,
+					&Option<#crate_::ProofRecorder<#block>>,
 				) -> std::result::Result<#crate_::NativeOrEncoded<R>, E>,
 				E,
 			>(

substrate/core/sr-api/src/lib.rs

@@ -59,16 +59,14 @@ use rstd::result;
 pub use codec::{Encode, Decode};
 use primitives::OpaqueMetadata;
 #[cfg(feature = "std")]
-use std::{panic::UnwindSafe, cell::RefCell, rc::Rc};
-#[cfg(feature = "std")]
-use primitives::Hasher as HasherT;
+use std::{panic::UnwindSafe, cell::RefCell};
 
 pub use sr_api_proc_macro::{decl_runtime_apis, impl_runtime_apis};
 
 #[cfg(feature = "std")]
 /// A type that records all accessed trie nodes and generates a proof out of it.
 pub type ProofRecorder<B> = state_machine::ProofRecorder<
-	<<<<B as BlockT>::Header as HeaderT>::Hashing as HashT>::Hasher as HasherT>::Out
+	<<<B as BlockT>::Header as HeaderT>::Hashing as HashT>::Hasher
 >;
 
 /// Something that can be constructed to a runtime api.
@@ -168,7 +166,7 @@ pub trait CallRuntimeAt<Block: BlockT> {
 		initialize_block: InitializeBlock<'a, Block>,
 		native_call: Option<NC>,
 		context: ExecutionContext,
-		recorder: &Option<Rc<RefCell<ProofRecorder<Block>>>>,
+		recorder: &Option<ProofRecorder<Block>>,
 	) -> Result<NativeOrEncoded<R>, Self::Error>;
 
 	/// Returns the runtime version at the given block.

substrate/core/state-machine/src/changes_trie/changes_iterator.rs

@@ -28,7 +28,7 @@ use crate::changes_trie::input::{DigestIndex, ExtrinsicIndex, DigestIndexValue,
 use crate::changes_trie::storage::{TrieBackendAdapter, InMemoryStorage};
 use crate::changes_trie::input::ChildIndex;
 use crate::changes_trie::surface_iterator::{surface_iterator, SurfaceIterator};
-use crate::proving_backend::ProvingBackendEssence;
+use crate::proving_backend::ProvingBackendRecorder;
 use crate::trie_backend_essence::{TrieBackendEssence};
 
 /// Return changes of given key at given blocks range.
@@ -366,7 +366,7 @@ impl<'a, H, Number> Iterator for ProvingDrilldownIterator<'a, H, Number>
 		let proof_recorder = &mut *self.proof_recorder.try_borrow_mut()
 			.expect("only fails when already borrowed; storage() is non-reentrant; qed");
 		self.essence.next(|storage, root, key|
-			ProvingBackendEssence::<_, H> {
+			ProvingBackendRecorder::<_, H> {
 				backend: &TrieBackendEssence::new(TrieBackendAdapter::new(storage), root),
 				proof_recorder,
 			}.storage(key))

substrate/core/state-machine/src/changes_trie/prune.rs

@@ -20,7 +20,7 @@ use hash_db::Hasher;
 use trie::Recorder;
 use log::warn;
 use num_traits::{One, Zero};
-use crate::proving_backend::ProvingBackendEssence;
+use crate::proving_backend::ProvingBackendRecorder;
 use crate::trie_backend_essence::TrieBackendEssence;
 use crate::changes_trie::{AnchorBlockId, Configuration, Storage, BlockNumber};
 use crate::changes_trie::storage::TrieBackendAdapter;
@@ -122,7 +122,7 @@ fn prune_trie<S: Storage<H, Number>, H: Hasher, Number: BlockNumber, F: FnMut(H:
 	// (effectively - all changes trie nodes)
 	let mut proof_recorder: Recorder<H::Out> = Default::default();
 	{
-		let mut trie = ProvingBackendEssence::<_, H> {
+		let mut trie = ProvingBackendRecorder::<_, H> {
 			backend: &TrieBackendEssence::new(TrieBackendAdapter::new(storage), root),
 			proof_recorder: &mut proof_recorder,
 		};

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

@@ -60,7 +60,7 @@ pub use changes_trie::{
 pub use overlayed_changes::OverlayedChanges;
 pub use proving_backend::{
 	create_proof_check_backend, create_proof_check_backend_storage, merge_storage_proofs,
-	Recorder as ProofRecorder, ProvingBackend, StorageProof,
+	ProofRecorder, ProvingBackend, ProvingBackendRecorder, StorageProof,
 };
 pub use trie_backend_essence::{TrieBackendStorage, Storage};
 pub use trie_backend::TrieBackend;

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

@@ -16,10 +16,11 @@
 
 //! Proving state machine backend.
 
-use std::{cell::RefCell, collections::HashSet, rc::Rc};
+use std::sync::Arc;
+use parking_lot::RwLock;
 use codec::{Decode, Encode};
 use log::debug;
-use hash_db::{Hasher, HashDB, EMPTY_PREFIX};
+use hash_db::{Hasher, HashDB, EMPTY_PREFIX, Prefix};
 use trie::{
 	MemoryDB, PrefixedMemoryDB, default_child_trie_root,
 	read_trie_value_with, read_child_trie_value_with, record_all_keys
@@ -29,6 +30,14 @@ pub use trie::trie_types::{Layout, TrieError};
 use crate::trie_backend::TrieBackend;
 use crate::trie_backend_essence::{Ephemeral, TrieBackendEssence, TrieBackendStorage};
 use crate::{Error, ExecutionError, Backend};
+use std::collections::{HashMap, HashSet};
+use crate::DBValue;
+
+/// Patricia trie-based backend specialized in get value proofs.
+pub struct ProvingBackendRecorder<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> {
+	pub(crate) backend: &'a TrieBackendEssence<S, H>,
+	pub(crate) proof_recorder: &'a mut Recorder<H::Out>,
+}
 
 /// A proof that some set of key-value pairs are included in the storage trie. The proof contains
 /// the storage values so that the partial storage backend can be reconstructed by a verifier that
@@ -106,18 +115,12 @@ pub fn merge_storage_proofs<I>(proofs: I) -> StorageProof
 	StorageProof { trie_nodes }
 }
 
-/// Patricia trie-based backend essence which also tracks all touched storage trie values.
-/// These can be sent to remote node and used as a proof of execution.
-pub struct ProvingBackendEssence<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> {
-	pub(crate) backend: &'a TrieBackendEssence<S, H>,
-	pub(crate) proof_recorder: &'a mut Recorder<H::Out>,
-}
-
-impl<'a, S, H> ProvingBackendEssence<'a, S, H>
+impl<'a, S, H> ProvingBackendRecorder<'a, S, H>
 	where
 		S: TrieBackendStorage<H>,
 		H: Hasher,
 {
+	/// Produce proof for a key query.
 	pub fn storage(&mut self, key: &[u8]) -> Result<Option<Vec<u8>>, String> {
 		let mut read_overlay = S::Overlay::default();
 		let eph = Ephemeral::new(
@@ -135,6 +138,7 @@ impl<'a, S, H> ProvingBackendEssence<'a, S, H>
 		).map_err(map_e)
 	}
 
+	/// Produce proof for a child key query.
 	pub fn child_storage(
 		&mut self,
 		storage_key: &[u8],
@@ -160,6 +164,7 @@ impl<'a, S, H> ProvingBackendEssence<'a, S, H>
 		).map_err(map_e)
 	}
 
+	/// Produce proof for the whole backend.
 	pub fn record_all_keys(&mut self) {
 		let mut read_overlay = S::Overlay::default();
 		let eph = Ephemeral::new(
@@ -178,45 +183,67 @@ impl<'a, S, H> ProvingBackendEssence<'a, S, H>
 	}
 }
 
+/// Global proof recorder, act as a layer over a hash db for recording queried
+/// data.
+pub type ProofRecorder<H> = Arc<RwLock<HashMap<<H as Hasher>::Out, Option<DBValue>>>>;
+
 /// Patricia trie-based backend which also tracks all touched storage trie values.
 /// These can be sent to remote node and used as a proof of execution.
-pub struct ProvingBackend<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> {
-	backend: &'a TrieBackend<S, H>,
-	proof_recorder: Rc<RefCell<Recorder<H::Out>>>,
+pub struct ProvingBackend<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> (
+	TrieBackend<ProofRecorderBackend<'a, S, H>, H>,
+);
+
+/// Trie backend storage with its proof recorder.
+pub struct ProofRecorderBackend<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> {
+	backend: &'a S,
+	proof_recorder: ProofRecorder<H>,
 }
 
 impl<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> ProvingBackend<'a, S, H> {
 	/// Create new proving backend.
 	pub fn new(backend: &'a TrieBackend<S, H>) -> Self {
-		ProvingBackend {
-			backend,
-			proof_recorder: Rc::new(RefCell::new(Recorder::new())),
-		}
+		let proof_recorder = Default::default();
+		Self::new_with_recorder(backend, proof_recorder)
 	}
 
 	/// Create new proving backend with the given recorder.
 	pub fn new_with_recorder(
 		backend: &'a TrieBackend<S, H>,
-		proof_recorder: Rc<RefCell<Recorder<H::Out>>>,
+		proof_recorder: ProofRecorder<H>,
 	) -> Self {
-		ProvingBackend {
-			backend,
-			proof_recorder,
-		}
+		let essence = backend.essence();
+		let root = essence.root().clone();
+		let recorder = ProofRecorderBackend {
+			backend: essence.backend_storage(),
+			proof_recorder: proof_recorder,
+		};
+		ProvingBackend(TrieBackend::new(recorder, root))
 	}
 
-	/// Consume the backend, extracting the gathered proof in lexicographical order by value.
+	/// Extracting the gathered unordered proof.
 	pub fn extract_proof(&self) -> StorageProof {
-		let trie_nodes = self.proof_recorder
-			.borrow_mut()
-			.drain()
-			.into_iter()
-			.map(|record| record.data)
+		let trie_nodes = self.0.essence().backend_storage().proof_recorder
+			.read()
+			.iter()
+			.filter_map(|(_k, v)| v.as_ref().map(|v| v.to_vec()))
 			.collect();
 		StorageProof::new(trie_nodes)
 	}
 }
 
+impl<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> TrieBackendStorage<H> for ProofRecorderBackend<'a, S, H> {
+	type Overlay = S::Overlay;
+
+	fn get(&self, key: &H::Out, prefix: Prefix) -> Result<Option<DBValue>, String> {
+		if let Some(v) = self.proof_recorder.read().get(key) {
+			return Ok(v.clone());
+		}
+		let backend_value =  self.backend.get(key, prefix)?;
+		self.proof_recorder.write().insert(key.clone(), backend_value.clone());
+		Ok(backend_value)
+	}
+}
+
 impl<'a, S: 'a + TrieBackendStorage<H>, H: 'a + Hasher> std::fmt::Debug for ProvingBackend<'a, S, H> {
 	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
 		write!(f, "ProvingBackend")
@@ -234,53 +261,45 @@ impl<'a, S, H> Backend<H> for ProvingBackend<'a, S, H>
 	type TrieBackendStorage = PrefixedMemoryDB<H>;
 
 	fn storage(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
-		ProvingBackendEssence {
-			backend: self.backend.essence(),
-			proof_recorder: &mut *self.proof_recorder.try_borrow_mut()
-				.expect("only fails when already borrowed; storage() is non-reentrant; qed"),
-		}.storage(key)
+		self.0.storage(key)
 	}
 
 	fn child_storage(&self, storage_key: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
-		ProvingBackendEssence {
-			backend: self.backend.essence(),
-			proof_recorder: &mut *self.proof_recorder.try_borrow_mut()
-				.expect("only fails when already borrowed; child_storage() is non-reentrant; qed"),
-		}.child_storage(storage_key, key)
+		self.0.child_storage(storage_key, key)
 	}
 
 	fn for_keys_in_child_storage<F: FnMut(&[u8])>(&self, storage_key: &[u8], f: F) {
-		self.backend.for_keys_in_child_storage(storage_key, f)
+		self.0.for_keys_in_child_storage(storage_key, f)
 	}
 
 	fn for_keys_with_prefix<F: FnMut(&[u8])>(&self, prefix: &[u8], f: F) {
-		self.backend.for_keys_with_prefix(prefix, f)
+		self.0.for_keys_with_prefix(prefix, f)
 	}
 
 	fn for_key_values_with_prefix<F: FnMut(&[u8], &[u8])>(&self, prefix: &[u8], f: F) {
-		self.backend.for_key_values_with_prefix(prefix, f)
+		self.0.for_key_values_with_prefix(prefix, f)
 	}
 
 	fn for_child_keys_with_prefix<F: FnMut(&[u8])>(&self, storage_key: &[u8], prefix: &[u8], f: F) {
-		self.backend.for_child_keys_with_prefix(storage_key, prefix, f)
+		self.0.for_child_keys_with_prefix(storage_key, prefix, f)
 	}
 
 	fn pairs(&self) -> Vec<(Vec<u8>, Vec<u8>)> {
-		self.backend.pairs()
+		self.0.pairs()
 	}
 
 	fn keys(&self, prefix: &[u8]) -> Vec<Vec<u8>> {
-		self.backend.keys(prefix)
+		self.0.keys(prefix)
 	}
 
 	fn child_keys(&self, child_storage_key: &[u8], prefix: &[u8]) -> Vec<Vec<u8>> {
-		self.backend.child_keys(child_storage_key, prefix)
+		self.0.child_keys(child_storage_key, prefix)
 	}
 
 	fn storage_root<I>(&self, delta: I) -> (H::Out, Self::Transaction)
 		where I: IntoIterator<Item=(Vec<u8>, Option<Vec<u8>>)>
 	{
-		self.backend.storage_root(delta)
+		self.0.storage_root(delta)
 	}
 
 	fn child_storage_root<I>(&self, storage_key: &[u8], delta: I) -> (Vec<u8>, bool, Self::Transaction)
@@ -288,7 +307,7 @@ impl<'a, S, H> Backend<H> for ProvingBackend<'a, S, H>
 		I: IntoIterator<Item=(Vec<u8>, Option<Vec<u8>>)>,
 		H::Out: Ord
 	{
-		self.backend.child_storage_root(storage_key, delta)
+		self.0.child_storage_root(storage_key, delta)
 	}
 }
 
@@ -329,6 +348,7 @@ mod tests {
 	use crate::trie_backend::tests::test_trie;
 	use super::*;
 	use primitives::{Blake2Hasher, storage::ChildStorageKey};
+	use crate::proving_backend::create_proof_check_backend;
 
 	fn test_proving<'a>(
 		trie_backend: &'a TrieBackend<PrefixedMemoryDB<Blake2Hasher>,Blake2Hasher>,