Add test method for growing storage proof size (#2224)

Author: serban300

bridges/modules/messages/Cargo.toml

@@ -32,6 +32,7 @@ sp-std = { path = "../../../substrate/primitives/std", default-features = false
 sp-trie = { path = "../../../substrate/primitives/trie", default-features = false, optional = true }
 
 [dev-dependencies]
+bp-runtime = { path = "../../primitives/runtime", features = ["test-helpers"] }
 bp-test-utils = { path = "../../primitives/test-utils" }
 pallet-bridge-grandpa = { path = "../grandpa" }
 pallet-balances = { path = "../../../substrate/frame/balances" }
@@ -71,4 +72,5 @@ try-runtime = [
 ]
 test-helpers = [
 	"sp-trie",
+	"bp-runtime/test-helpers"
 ]

bridges/primitives/runtime/Cargo.toml

@@ -53,3 +53,4 @@ std = [
 	"sp-trie/std",
 	"trie-db/std",
 ]
+test-helpers = []

bridges/primitives/runtime/src/lib.rs

@@ -40,10 +40,9 @@ pub use chain::{
 };
 pub use frame_support::storage::storage_prefix as storage_value_final_key;
 use num_traits::{CheckedAdd, CheckedSub, One, SaturatingAdd, Zero};
-pub use storage_proof::{
-	grow_storage_value, StorageProofError, StorageProofSize, UnverifiedStorageProof,
-	VerifiedStorageProof,
-};
+#[cfg(feature = "test-helpers")]
+pub use storage_proof::{grow_storage_proof, grow_storage_value, StorageProofSize};
+pub use storage_proof::{StorageProofError, UnverifiedStorageProof, VerifiedStorageProof};
 pub use storage_types::BoundedStorageValue;
 
 pub mod extensions;

bridges/primitives/runtime/src/storage_proof.rs

@@ -29,6 +29,8 @@ use codec::{Decode, Encode};
 use hash_db::Hasher;
 use scale_info::TypeInfo;
 use trie_db::{DBValue, Recorder, Trie};
+#[cfg(feature = "test-helpers")]
+use trie_db::{TrieConfiguration, TrieDBMut};
 
 use crate::Size;
 
@@ -95,7 +97,7 @@ impl UnverifiedStorageProof {
 	/// Creates a new instance of `UnverifiedStorageProof` from the provided entries.
 	///
 	/// **This function is used only in tests and benchmarks.**
-	#[cfg(feature = "std")]
+	#[cfg(any(all(feature = "std", feature = "test-helpers"), test))]
 	pub fn try_from_entries<H: Hasher>(
 		state_version: StateVersion,
 		entries: &[(RawStorageKey, Option<DBValue>)],
@@ -116,6 +118,7 @@ impl UnverifiedStorageProof {
 	/// Creates a new instance of `UnverifiedStorageProof` from the provided db.
 	///
 	/// **This function is used only in tests and benchmarks.**
+	#[cfg(any(feature = "test-helpers", test))]
 	pub fn try_from_db<H: Hasher, DB>(
 		db: &DB,
 		root: H::Out,
@@ -249,6 +252,7 @@ impl VerifiedStorageProof {
 /// Storage proof size requirements.
 ///
 /// This is currently used by benchmarks when generating storage proofs.
+#[cfg(feature = "test-helpers")]
 #[derive(Clone, Copy, Debug)]
 pub enum StorageProofSize {
 	/// The storage proof is expected to be minimal. If value size may be changed, then it is
@@ -260,6 +264,7 @@ pub enum StorageProofSize {
 }
 
 /// Add extra data to the storage value so that it'll be of given size.
+#[cfg(feature = "test-helpers")]
 pub fn grow_storage_value(mut value: Vec<u8>, size: StorageProofSize) -> Vec<u8> {
 	match size {
 		StorageProofSize::Minimal(_) => (),
@@ -271,6 +276,57 @@ pub fn grow_storage_value(mut value: Vec<u8>, size: StorageProofSize) -> Vec<u8>
 	value
 }
 
+/// Insert values in the provided trie at common-prefix keys in order to inflate the resulting
+/// storage proof.
+///
+/// This function can add at most 15 common-prefix keys per prefix nibble (4 bits).
+/// Each such key adds about 33 bytes (a node) to the proof.
+#[cfg(feature = "test-helpers")]
+pub fn grow_storage_proof<L: TrieConfiguration>(
+	trie: &mut TrieDBMut<L>,
+	prefix: Vec<u8>,
+	num_extra_nodes: usize,
+) {
+	use sp_trie::TrieMut;
+
+	let mut added_nodes = 0;
+	for i in 0..prefix.len() {
+		let mut prefix = prefix[0..=i].to_vec();
+		// 1 byte has 2 nibbles (4 bits each)
+		let first_nibble = (prefix[i] & 0xf0) >> 4;
+		let second_nibble = prefix[i] & 0x0f;
+
+		// create branches at the 1st nibble
+		for branch in 1..=15 {
+			if added_nodes >= num_extra_nodes {
+				return
+			}
+
+			// create branches at the 1st nibble
+			prefix[i] = (first_nibble.wrapping_add(branch) % 16) << 4;
+			trie.insert(&prefix, &[0; 32])
+				.map_err(|_| "TrieMut::insert has failed")
+				.expect("TrieMut::insert should not fail in benchmarks");
+			added_nodes += 1;
+		}
+
+		// create branches at the 2nd nibble
+		for branch in 1..=15 {
+			if added_nodes >= num_extra_nodes {
+				return
+			}
+
+			prefix[i] = (first_nibble << 4) | (second_nibble.wrapping_add(branch) % 16);
+			trie.insert(&prefix, &[0; 32])
+				.map_err(|_| "TrieMut::insert has failed")
+				.expect("TrieMut::insert should not fail in benchmarks");
+			added_nodes += 1;
+		}
+	}
+
+	assert_eq!(added_nodes, num_extra_nodes)
+}
+
 #[cfg(test)]
 mod tests {
 	use super::*;