Compact proof utilities in sp_trie.

client/db/src/bench.rs

@@ -23,7 +23,7 @@ use std::cell::{Cell, RefCell};
 use std::collections::HashMap;
 
 use hash_db::{Prefix, Hasher};
-use sp_trie::{MemoryDB, prefixed_key};
+use sp_trie::{MemoryDB, prefixed_key, CompactProof};
 use sp_core::{
 	storage::{ChildInfo, TrackedStorageKey},
 	hexdisplay::HexDisplay
@@ -117,6 +117,7 @@ pub struct BenchmarkingState<B: BlockT> {
 	read_write_tracker: RefCell<ReadWriteTracker>,
 	whitelist: RefCell<Vec<TrackedStorageKey>>,
 	proof_recorder: Option<ProofRecorder<B::Hash>>,
+	proof_recorder_root: Cell<B::Hash>,
 }
 
 impl<B: BlockT> BenchmarkingState<B> {
@@ -129,7 +130,7 @@ impl<B: BlockT> BenchmarkingState<B> {
 		let mut state = BenchmarkingState {
 			state: RefCell::new(None),
 			db: Cell::new(None),
-			root: Cell::new(root),
+			root: Cell::new(root.clone()),
 			genesis: Default::default(),
 			genesis_root: Default::default(),
 			record: Default::default(),
@@ -139,6 +140,7 @@ impl<B: BlockT> BenchmarkingState<B> {
 			read_write_tracker: Default::default(),
 			whitelist: Default::default(),
 			proof_recorder: record_proof.then(Default::default),
+			proof_recorder_root: Cell::new(root.clone()),
 		};
 
 		state.add_whitelist_to_tracker();
@@ -166,7 +168,10 @@ impl<B: BlockT> BenchmarkingState<B> {
 			None => Arc::new(kvdb_memorydb::create(1)),
 		};
 		self.db.set(Some(db.clone()));
-		self.proof_recorder.as_ref().map(|r| r.reset());
+		if let Some(recorder) = &self.proof_recorder {
+			recorder.reset();
+			self.proof_recorder_root.set(self.root.get());
+		}
 		let storage_db = Arc::new(StorageDb::<B> {
 			db,
 			proof_recorder: self.proof_recorder.clone(),
@@ -515,8 +520,33 @@ impl<B: BlockT> StateBackend<HashFor<B>> for BenchmarkingState<B> {
 		self.state.borrow().as_ref().map_or(sp_state_machine::UsageInfo::empty(), |s| s.usage_info())
 	}
 
-	fn proof_size(&self) -> Option<u32> {
-		self.proof_recorder.as_ref().map(|recorder| recorder.estimate_encoded_size() as u32)
+	fn proof_size(&self) -> Option<(u32, u32)> {
+		self.proof_recorder.as_ref().map(|recorder| {
+			let proof_size = recorder.estimate_encoded_size() as u32;
+			let proof = recorder.to_storage_proof();
+			let proof_recorder_root = self.proof_recorder_root.get();
+			let compact_size = if proof_recorder_root == Default::default() || proof_size == 1 {
+				// empty trie
+				proof_size
+			} else {
+				if let Some(size) = CompactProof::encoded_compact_size::<HashFor<B>>(
+					proof,
+					proof_recorder_root,
+				) {
+					size as u32
+				} else {
+					panic!(
+						"proof rec root {:?}, root {:?}, genesis {:?}, rec_len {:?}",
+						self.proof_recorder_root.get(),
+						self.root.get(),
+						self.genesis_root,
+						proof_size,
+					);
+				}
+			};
+
+			(proof_size, compact_size)
+		})
 	}
 }
 

frame/benchmarking/src/analysis.rs

@@ -39,6 +39,7 @@ pub enum BenchmarkSelector {
 	Reads,
 	Writes,
 	ProofSize,
+	CompactProofSize,
 }
 
 #[derive(Debug)]
@@ -88,6 +89,7 @@ impl Analysis {
 				BenchmarkSelector::Reads => result.reads.into(),
 				BenchmarkSelector::Writes => result.writes.into(),
 				BenchmarkSelector::ProofSize => result.proof_size.into(),
+				BenchmarkSelector::CompactProofSize => result.compact_proof_size.into(),
 			}
 		).collect();
 
@@ -129,6 +131,7 @@ impl Analysis {
 						BenchmarkSelector::Reads => result.reads.into(),
 						BenchmarkSelector::Writes => result.writes.into(),
 						BenchmarkSelector::ProofSize => result.proof_size.into(),
+						BenchmarkSelector::CompactProofSize => result.compact_proof_size.into(),
 					};
 					(result.components[i].1, data)
 				})
@@ -194,6 +197,7 @@ impl Analysis {
 					BenchmarkSelector::Reads => result.reads.into(),
 					BenchmarkSelector::Writes => result.writes.into(),
 					BenchmarkSelector::ProofSize => result.proof_size.into(),
+					BenchmarkSelector::CompactProofSize => result.compact_proof_size.into(),
 				})
 		}
 
@@ -375,6 +379,7 @@ mod tests {
 			writes,
 			repeat_writes: 0,
 			proof_size: 0,
+			compact_proof_size: 0,
 		}
 	}
 

frame/benchmarking/src/lib.rs

@@ -784,8 +784,11 @@ macro_rules! impl_benchmark {
 
 							// Calculate the diff caused by the benchmark.
 							let elapsed_extrinsic = finish_extrinsic.saturating_sub(start_extrinsic);
-							let diff_pov = match (start_pov, end_pov) {
-								(Some(start), Some(end)) => end.saturating_sub(start),
+							let (diff_pov, compact_diff_pov) = match (start_pov, end_pov) {
+								(Some((start, compact_start)), Some((end, compact_end))) => (
+									end.saturating_sub(start),
+									compact_end.saturating_sub(compact_start),
+								),
 								_ => Default::default(),
 							};
 
@@ -816,6 +819,7 @@ macro_rules! impl_benchmark {
 								writes: read_write_count.2,
 								repeat_writes: read_write_count.3,
 								proof_size: diff_pov,
+								compact_proof_size: compact_diff_pov,
 							});
 						}
 

frame/benchmarking/src/utils.rs

@@ -63,6 +63,7 @@ pub struct BenchmarkResults {
 	pub writes: u32,
 	pub repeat_writes: u32,
 	pub proof_size: u32,
+	pub compact_proof_size: u32,
 }
 
 /// Configuration used to setup and run runtime benchmarks.
@@ -165,7 +166,7 @@ pub trait Benchmarking {
 	}
 
 	/// Get current estimated proof size.
-	fn proof_size(&self) -> Option<u32> {
+	fn proof_size(&self) -> Option<(u32, u32)> {
 		self.proof_size()
 	}
 }

primitives/externalities/src/lib.rs

@@ -288,7 +288,7 @@ pub trait Externalities: ExtensionStore {
 	///
 	/// Returns estimated proof size for the state queries so far.
 	/// Proof is reset on commit and wipe.
-	fn proof_size(&self) -> Option<u32> {
+	fn proof_size(&self) -> Option<(u32, u32)> {
 		None
 	}
 }

primitives/state-machine/src/backend.rs

@@ -247,7 +247,7 @@ pub trait Backend<H: Hasher>: sp_std::fmt::Debug {
 	fn set_whitelist(&self, _: Vec<TrackedStorageKey>) {}
 
 	/// Estimate proof size
-	fn proof_size(&self) -> Option<u32> {
+	fn proof_size(&self) -> Option<(u32, u32)> {
 		unimplemented!()
 	}
 }

primitives/state-machine/src/ext.rs

@@ -734,7 +734,7 @@ where
 		self.backend.set_whitelist(new)
 	}
 
-	fn proof_size(&self) -> Option<u32> {
+	fn proof_size(&self) -> Option<(u32, u32)> {
 		self.backend.proof_size()
 	}
 }

primitives/state-machine/src/lib.rs

@@ -1404,12 +1404,33 @@ mod tests {
 
 	#[test]
 	fn prove_read_and_proof_check_works() {
+		fn test_compact(remote_proof: StorageProof, remote_root: &sp_core::H256) -> StorageProof {
+			type Layout = sp_trie::Layout<BlakeTwo256>;
+			let compact_remote_proof = sp_trie::encode_compact::<Layout>(
+				remote_proof,
+				remote_root.clone(),
+			).unwrap();
+			let mut db = sp_trie::MemoryDB::<BlakeTwo256>::new(&[]);
+			sp_trie::decode_compact::<Layout, _, _>(
+				&mut db,
+				compact_remote_proof.iter_compact_encoded_nodes(),
+				Some(remote_root),
+			).unwrap();
+			StorageProof::new(db.drain().into_iter().filter_map(|kv|
+				if (kv.1).1 > 0 {
+					Some((kv.1).0)
+				} else {
+					None
+				}
+			).collect())
+		}
 		let child_info = ChildInfo::new_default(b"sub1");
 		let child_info = &child_info;
 		// fetch read proof from 'remote' full node
 		let remote_backend = trie_backend::tests::test_trie();
-		let remote_root = remote_backend.storage_root(::std::iter::empty()).0;
+		let remote_root = remote_backend.storage_root(std::iter::empty()).0;
 		let remote_proof = prove_read(remote_backend, &[b"value2"]).unwrap();
+		let remote_proof = test_compact(remote_proof, &remote_root);
  		// check proof locally
 		let local_result1 = read_proof_check::<BlakeTwo256, _>(
 			remote_root,
@@ -1429,12 +1450,13 @@ mod tests {
 		assert_eq!(local_result2, false);
 		// on child trie
 		let remote_backend = trie_backend::tests::test_trie();
-		let remote_root = remote_backend.storage_root(::std::iter::empty()).0;
+		let remote_root = remote_backend.storage_root(std::iter::empty()).0;
 		let remote_proof = prove_child_read(
 			remote_backend,
 			child_info,
 			&[b"value3"],
 		).unwrap();
+		let remote_proof = test_compact(remote_proof, &remote_root);
 		let local_result1 = read_child_proof_check::<BlakeTwo256, _>(
 			remote_root,
 			remote_proof.clone(),

primitives/trie/Cargo.toml

@@ -21,7 +21,7 @@ harness = false
 codec = { package = "parity-scale-codec", version = "2.0.0", default-features = false }
 sp-std = { version = "3.0.0", default-features = false, path = "../std" }
 hash-db = { version = "0.15.2", default-features = false }
-trie-db = { version = "0.22.2", default-features = false }
+trie-db = { version = "0.22.3", default-features = false }
 trie-root = { version = "0.16.0", default-features = false }
 memory-db = { version = "0.26.0", default-features = false }
 sp-core = { version = "3.0.0", default-features = false, path = "../core" }

primitives/trie/src/lib.rs

@@ -23,6 +23,7 @@ mod error;
 mod node_header;
 mod node_codec;
 mod storage_proof;
+mod trie_codec;
 mod trie_stream;
 
 use sp_std::{boxed::Box, marker::PhantomData, vec::Vec, borrow::Borrow};
@@ -35,7 +36,7 @@ pub use error::Error;
 pub use trie_stream::TrieStream;
 /// The Substrate format implementation of `NodeCodec`.
 pub use node_codec::NodeCodec;
-pub use storage_proof::StorageProof;
+pub use storage_proof::{StorageProof, CompactProof};
 /// Various re-exports from the `trie-db` crate.
 pub use trie_db::{
 	Trie, TrieMut, DBValue, Recorder, CError, Query, TrieLayout, TrieConfiguration, nibble_ops, TrieDBIterator,
@@ -45,6 +46,9 @@ pub use memory_db::KeyFunction;
 pub use memory_db::prefixed_key;
 /// Various re-exports from the `hash-db` crate.
 pub use hash_db::{HashDB as HashDBT, EMPTY_PREFIX};
+/// Trie codec reexport, mainly child trie support
+/// for trie compact proof.
+pub use trie_codec::{decode_compact, encode_compact};
 
 #[derive(Default)]
 /// substrate trie layout

primitives/trie/src/storage_proof.rs

@@ -31,6 +31,12 @@ pub struct StorageProof {
 	trie_nodes: Vec<Vec<u8>>,
 }
 
+/// Storage proof in compact form.
+#[derive(Debug, PartialEq, Eq, Clone, Encode, Decode)]
+pub struct CompactProof {
+	pub encoded_nodes: Vec<Vec<u8>>,
+}
+
 impl StorageProof {
 	/// Constructs a storage proof from a subset of encoded trie nodes in a storage backend.
 	pub fn new(trie_nodes: Vec<Vec<u8>>) -> Self {
@@ -77,6 +83,33 @@ impl StorageProof {
 	}
 }
 
+impl CompactProof {
+	/// Return an iterator on the compact encoded nodes.
+	pub fn iter_compact_encoded_nodes(&self) -> impl Iterator<Item = &[u8]> {
+		self.encoded_nodes.iter().map(Vec::as_slice)
+	}
+
+	/// 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.
+	/// Return `None` on error.
+	pub fn encoded_compact_size<H: Hasher>(
+		proof: StorageProof,
+		root: H::Out,
+	) -> Option<usize> {
+		let compact_proof = crate::encode_compact::<crate::Layout<H>>(
+			proof,
+			root,
+		);
+		if let Ok(proof) = compact_proof {
+			Some(proof.encoded_size())
+		} else {
+			None
+		}
+	}
+}
+
 /// An iterator over trie nodes constructed from a storage proof. The nodes are not guaranteed to
 /// be traversed in any particular order.
 pub struct StorageProofNodeIterator {