feat: generalize some functions in sp-trie (#12376)

* feat: add to_memory_db to StorageProof

* feat: add iter method and generalize iter_nodes

* fmt

* feat: generalize `encode_compact` like `decode_compact`, add to_compact_proof to StorageProof

* fix to_compact_proof

* improve by suggestions

* improve by suggestions

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

client/rpc/src/state/state_full.rs

@@ -346,7 +346,7 @@ where
 			.and_then(|block| {
 				self.client
 					.read_proof(&block, &mut keys.iter().map(|key| key.0.as_ref()))
-					.map(|proof| proof.iter_nodes().map(|node| node.into()).collect())
+					.map(|proof| proof.into_iter_nodes().map(|node| node.into()).collect())
 					.map(|proof| ReadProof { at: block, proof })
 			})
 			.map_err(client_err)
@@ -498,7 +498,7 @@ where
 						&child_info,
 						&mut keys.iter().map(|key| key.0.as_ref()),
 					)
-					.map(|proof| proof.iter_nodes().map(|node| node.into()).collect())
+					.map(|proof| proof.into_iter_nodes().map(|node| node.into()).collect())
 					.map(|proof| ReadProof { at: block, proof })
 			})
 			.map_err(client_err)

client/service/src/client/client.rs

@@ -1190,8 +1190,8 @@ where
 		let (proof, count) = prove_range_read_with_child_with_size::<_, HashFor<Block>>(
 			state, size_limit, start_key,
 		)?;
-		// This is read proof only, we can use either LayoutV0 or LayoutV1.
-		let proof = sp_trie::encode_compact::<sp_trie::LayoutV0<HashFor<Block>>>(proof, root)
+		let proof = proof
+			.into_compact_proof::<HashFor<Block>>(root)
 			.map_err(|e| sp_blockchain::Error::from_state(Box::new(e)))?;
 		Ok((proof, count))
 	}

primitives/state-machine/src/lib.rs

@@ -1939,13 +1939,13 @@ mod tests {
 		let (proof, count) =
 			prove_range_read_with_size(remote_backend, None, None, 0, None).unwrap();
 		// Always contains at least some nodes.
-		assert_eq!(proof.into_memory_db::<BlakeTwo256>().drain().len(), 3);
+		assert_eq!(proof.to_memory_db::<BlakeTwo256>().drain().len(), 3);
 		assert_eq!(count, 1);
 
 		let remote_backend = trie_backend::tests::test_trie(state_version, None, None);
 		let (proof, count) =
 			prove_range_read_with_size(remote_backend, None, None, 800, Some(&[])).unwrap();
-		assert_eq!(proof.clone().into_memory_db::<BlakeTwo256>().drain().len(), 9);
+		assert_eq!(proof.to_memory_db::<BlakeTwo256>().drain().len(), 9);
 		assert_eq!(count, 85);
 		let (results, completed) = read_range_proof_check::<BlakeTwo256>(
 			remote_root,
@@ -1968,7 +1968,7 @@ mod tests {
 		let remote_backend = trie_backend::tests::test_trie(state_version, None, None);
 		let (proof, count) =
 			prove_range_read_with_size(remote_backend, None, None, 50000, Some(&[])).unwrap();
-		assert_eq!(proof.clone().into_memory_db::<BlakeTwo256>().drain().len(), 11);
+		assert_eq!(proof.to_memory_db::<BlakeTwo256>().drain().len(), 11);
 		assert_eq!(count, 132);
 		let (results, completed) =
 			read_range_proof_check::<BlakeTwo256>(remote_root, proof, None, None, None, None)
@@ -2053,7 +2053,7 @@ mod tests {
 			)
 			.unwrap();
 			// Always contains at least some nodes.
-			assert!(proof.clone().into_memory_db::<BlakeTwo256>().drain().len() > 0);
+			assert!(proof.to_memory_db::<BlakeTwo256>().drain().len() > 0);
 			assert!(count < 3); // when doing child we include parent and first child key.
 
 			let (result, completed_depth) = read_range_proof_check_with_child::<BlakeTwo256>(

primitives/state-machine/src/trie_backend.rs

@@ -980,7 +980,7 @@ pub mod tests {
 			let proof = backend.extract_proof().unwrap();
 
 			let mut nodes = Vec::new();
-			for node in proof.iter_nodes() {
+			for node in proof.into_iter_nodes() {
 				let hash = BlakeTwo256::hash(&node);
 				// Only insert the node/value that contains the important data.
 				if hash != value_hash {

primitives/trie/src/lib.rs

@@ -57,10 +57,10 @@ pub use trie_db::{
 pub use trie_stream::TrieStream;
 
 /// substrate trie layout
-pub struct LayoutV0<H>(sp_std::marker::PhantomData<H>);
+pub struct LayoutV0<H>(PhantomData<H>);
 
 /// substrate trie layout, with external value nodes.
-pub struct LayoutV1<H>(sp_std::marker::PhantomData<H>);
+pub struct LayoutV1<H>(PhantomData<H>);
 
 impl<H> TrieLayout for LayoutV0<H>
 where

primitives/trie/src/storage_proof.rs

@@ -18,7 +18,11 @@
 use codec::{Decode, Encode};
 use hash_db::{HashDB, Hasher};
 use scale_info::TypeInfo;
-use sp_std::{collections::btree_set::BTreeSet, iter::IntoIterator, vec::Vec};
+use sp_std::{
+	collections::btree_set::BTreeSet,
+	iter::{DoubleEndedIterator, IntoIterator},
+	vec::Vec,
+};
 // Note that `LayoutV1` usage here (proof compaction) is compatible
 // with `LayoutV0`.
 use crate::LayoutV1 as Layout;
@@ -54,10 +58,16 @@ impl StorageProof {
 		self.trie_nodes.is_empty()
 	}
 
+	/// Convert into an iterator over encoded trie nodes in lexicographical order constructed
+	/// from the proof.
+	pub fn into_iter_nodes(self) -> impl Sized + DoubleEndedIterator<Item = Vec<u8>> {
+		self.trie_nodes.into_iter()
+	}
+
 	/// Create an iterator over encoded trie nodes in lexicographical order constructed
 	/// from the proof.
-	pub fn iter_nodes(self) -> StorageProofNodeIterator {
-		StorageProofNodeIterator::new(self)
+	pub fn iter_nodes(&self) -> impl Sized + DoubleEndedIterator<Item = &Vec<u8>> {
+		self.trie_nodes.iter()
 	}
 
 	/// Convert into plain node vector.
@@ -70,14 +80,19 @@ impl StorageProof {
 		self.into()
 	}
 
+	/// Creates a [`MemoryDB`](crate::MemoryDB) from `Self` reference.
+	pub fn to_memory_db<H: Hasher>(&self) -> crate::MemoryDB<H> {
+		self.into()
+	}
+
 	/// Merges multiple storage proofs covering potentially different sets of keys into one proof
 	/// covering all keys. The merged proof output may be smaller than the aggregate size of the
 	/// input proofs due to deduplication of trie nodes.
 	pub fn merge(proofs: impl IntoIterator<Item = Self>) -> Self {
 		let trie_nodes = proofs
 			.into_iter()
-			.flat_map(|proof| proof.iter_nodes())
-			.collect::<sp_std::collections::btree_set::BTreeSet<_>>()
+			.flat_map(|proof| proof.into_iter_nodes())
+			.collect::<BTreeSet<_>>()
 			.into_iter()
 			.collect();
 
@@ -89,7 +104,17 @@ impl StorageProof {
 		self,
 		root: H::Out,
 	) -> Result<CompactProof, crate::CompactProofError<H::Out, crate::Error<H::Out>>> {
-		crate::encode_compact::<Layout<H>>(self, root)
+		let db = self.into_memory_db();
+		crate::encode_compact::<Layout<H>, crate::MemoryDB<H>>(&db, &root)
+	}
+
+	/// Encode as a compact proof with default trie layout.
+	pub fn to_compact_proof<H: Hasher>(
+		&self,
+		root: H::Out,
+	) -> Result<CompactProof, crate::CompactProofError<H::Out, crate::Error<H::Out>>> {
+		let db = self.to_memory_db();
+		crate::encode_compact::<Layout<H>, crate::MemoryDB<H>>(&db, &root)
 	}
 
 	/// Returns the estimated encoded size of the compact proof.
@@ -106,6 +131,12 @@ impl StorageProof {
 
 impl<H: Hasher> From<StorageProof> for crate::MemoryDB<H> {
 	fn from(proof: StorageProof) -> Self {
+		From::from(&proof)
+	}
+}
+
+impl<H: Hasher> From<&StorageProof> for crate::MemoryDB<H> {
+	fn from(proof: &StorageProof) -> Self {
 		let mut db = crate::MemoryDB::default();
 		proof.iter_nodes().for_each(|n| {
 			db.insert(crate::EMPTY_PREFIX, &n);
@@ -169,23 +200,3 @@ impl CompactProof {
 		Ok((db, root))
 	}
 }
-
-/// 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 {
-	inner: <BTreeSet<Vec<u8>> as IntoIterator>::IntoIter,
-}
-
-impl StorageProofNodeIterator {
-	fn new(proof: StorageProof) -> Self {
-		StorageProofNodeIterator { inner: proof.trie_nodes.into_iter() }
-	}
-}
-
-impl Iterator for StorageProofNodeIterator {
-	type Item = Vec<u8>;
-
-	fn next(&mut self) -> Option<Self::Item> {
-		self.inner.next()
-	}
-}

primitives/trie/src/trie_codec.rs

@@ -20,7 +20,7 @@
 //! This uses compact proof from trie crate and extends
 //! it to substrate specific layout and child trie system.
 
-use crate::{CompactProof, HashDBT, StorageProof, TrieConfiguration, TrieHash, EMPTY_PREFIX};
+use crate::{CompactProof, HashDBT, TrieConfiguration, TrieHash, EMPTY_PREFIX};
 use sp_std::{boxed::Box, vec::Vec};
 use trie_db::{CError, Trie};
 
@@ -149,17 +149,17 @@ where
 /// Then parse all child trie root and compress main trie content first
 /// then all child trie contents.
 /// Child trie are ordered by the order of their roots in the top trie.
-pub fn encode_compact<L>(
-	proof: StorageProof,
-	root: TrieHash<L>,
+pub fn encode_compact<L, DB>(
+	partial_db: &DB,
+	root: &TrieHash<L>,
 ) -> Result<CompactProof, Error<TrieHash<L>, CError<L>>>
 where
 	L: TrieConfiguration,
+	DB: HashDBT<L::Hash, trie_db::DBValue> + hash_db::HashDBRef<L::Hash, trie_db::DBValue>,
 {
 	let mut child_tries = Vec::new();
-	let partial_db = proof.into_memory_db();
 	let mut compact_proof = {
-		let trie = crate::TrieDBBuilder::<L>::new(&partial_db, &root).build();
+		let trie = crate::TrieDBBuilder::<L>::new(partial_db, root).build();
 
 		let mut iter = trie.iter()?;
 
@@ -191,13 +191,13 @@ where
 	};
 
 	for child_root in child_tries {
-		if !HashDBT::<L::Hash, _>::contains(&partial_db, &child_root, EMPTY_PREFIX) {
+		if !HashDBT::<L::Hash, _>::contains(partial_db, &child_root, EMPTY_PREFIX) {
 			// child proof are allowed to be missing (unused root can be included
 			// due to trie structure modification).
 			continue
 		}
 
-		let trie = crate::TrieDBBuilder::<L>::new(&partial_db, &child_root).build();
+		let trie = crate::TrieDBBuilder::<L>::new(partial_db, &child_root).build();
 		let child_proof = trie_db::encode_compact::<L>(&trie)?;
 
 		compact_proof.extend(child_proof);