create proof size to hashes trait

substrate/primitives/runtime/src/proving_trie/base16.rs

@@ -24,7 +24,7 @@
 //! Proofs are created with latest substrate trie format (`LayoutV1`), and are not compatible with
 //! proofs using `LayoutV0`.
 
-use super::{ProvingTrie, TrieError};
+use super::{ProofSizeToHashes, ProvingTrie, TrieError};
 use crate::{Decode, DispatchError, Encode};
 use codec::MaxEncodedLen;
 use sp_std::vec::Vec;
@@ -72,7 +72,6 @@ where
 impl<Hashing, Key, Value> ProvingTrie<Hashing, Key, Value> for BasicProvingTrie<Hashing, Key, Value>
 where
 	Hashing: sp_core::Hasher,
-	Hashing::Out: MaxEncodedLen,
 	Key: Encode,
 	Value: Encode + Decode,
 {
@@ -134,7 +133,13 @@ where
 	) -> Result<(), DispatchError> {
 		verify_proof::<Hashing, Key, Value>(root, proof, key, value)
 	}
+}
 
+impl<Hashing, Key, Value> ProofSizeToHashes for BasicProvingTrie<Hashing, Key, Value>
+where
+	Hashing: sp_core::Hasher,
+	Hashing::Out: MaxEncodedLen,
+{
 	fn proof_size_to_hashes(proof_size: &u32) -> u32 {
 		let hash_len = Hashing::Out::max_encoded_len() as u32;
 		// A base 16 trie is expected to include the data for 15 hashes per layer.

substrate/primitives/runtime/src/proving_trie/base2.rs

@@ -20,7 +20,7 @@
 //! this library is designed to work more easily with runtime native types, which simply need to
 //! implement `Encode`/`Decode`.
 
-use super::{ProvingTrie, TrieError};
+use super::{ProofSizeToHashes, ProvingTrie, TrieError};
 use crate::{Decode, DispatchError, Encode};
 use binary_merkle_tree::{merkle_proof, merkle_root, MerkleProof};
 use codec::MaxEncodedLen;
@@ -41,7 +41,7 @@ where
 impl<Hashing, Key, Value> ProvingTrie<Hashing, Key, Value> for BasicProvingTrie<Hashing, Key, Value>
 where
 	Hashing: sp_core::Hasher,
-	Hashing::Out: Encode + Decode + MaxEncodedLen,
+	Hashing::Out: Encode + Decode,
 	Key: Encode + Decode + Ord,
 	Value: Encode + Decode + Clone,
 {
@@ -100,11 +100,19 @@ where
 	) -> Result<(), DispatchError> {
 		verify_proof::<Hashing, Key, Value>(root, proof, key, value)
 	}
+}
 
-	/// A base 2 trie is expected to include the data for 1 hash per layer.
+impl<Hashing, Key, Value> ProofSizeToHashes for BasicProvingTrie<Hashing, Key, Value>
+where
+	Hashing: sp_core::Hasher,
+	Hashing::Out: MaxEncodedLen,
+{
 	fn proof_size_to_hashes(proof_size: &u32) -> u32 {
 		let hash_len = Hashing::Out::max_encoded_len() as u32;
+		// A base 2 trie is expected to include the data for 1 hash per layer.
 		let layer_len = 1 * hash_len;
+		// The proof includes `number_of_leaves: u32` and `leaf_index: u32`.
+		let proof_size = proof_size.saturating_sub(8);
 		// The implementation of this trie also includes the `key` and `value` encoded within the
 		// proof, but since we cannot know the "minimum" size of those items, we count it toward
 		// the number of hashes for a worst case scenario.
@@ -243,4 +251,24 @@ mod tests {
 			Err(TrieError::IncompleteProof.into())
 		);
 	}
+
+	// We make assumptions about the structure of the merkle proof in order to provide the
+	// `proof_size_to_hashes` function. This test keeps those assumptions checked.
+	#[test]
+	fn assert_structure_of_merkle_proof() {
+		let balance_trie = create_balance_trie();
+		let root = *balance_trie.root();
+		// Create a proof for a valid key.
+		let proof = balance_trie.create_proof(&6u32).unwrap();
+		let decoded_proof: MerkleProof<H256, Vec<u8>> = Decode::decode(&mut &proof[..]).unwrap();
+
+		let constructed_proof = MerkleProof::<H256, Vec<u8>> {
+			root,
+			proof: decoded_proof.proof.clone(),
+			number_of_leaves: 100,
+			leaf_index: 6,
+			leaf: (6u32, 6u128).encode(),
+		};
+		assert_eq!(constructed_proof, decoded_proof);
+	}
 }

substrate/primitives/runtime/src/proving_trie/mod.rs

@@ -137,6 +137,13 @@ where
 		key: &Key,
 		value: &Value,
 	) -> Result<(), DispatchError>;
+}
+
+/// This trait is one strategy that can be used to benchmark a trie proof verification for the
+/// runtime. This strategy assumes that the majority complexity of verifying a merkle proof comes
+/// from computing hashes to recreate the merkle root. This trait converts the size of the proof, in
+/// bytes, to the number of hashes we expect to execute.
+pub trait ProofSizeToHashes {
 	/// This function returns the number of hashes we expect to calculate based on the
 	/// size of the proof. This is used for benchmarking, so for worst case scenario, we should
 	/// round up.
@@ -181,8 +188,6 @@ mod tests {
 
 	#[test]
 	fn proof_size_to_hashes() {
-		use crate::{testing::H256, traits::BlakeTwo256};
-
 		// We can be off by up to 2 hashes... should be trivial.
 		let tolerance = 2;