update api to use references everywhere

paritytech · Sep 21, 2024 · 3ab4d94 · 3ab4d94
1 parent 99135e4
commit 3ab4d94
Show file tree

Hide file tree

Showing 4 changed files with 70 additions and 74 deletions.
diff --git a/substrate/frame/support/src/traits/proving.rs b/substrate/frame/support/src/traits/proving.rs
@@ -107,7 +107,7 @@ impl VerifyExistenceProof for () {
 #[cfg(test)]
 mod tests {
 	use super::*;
-	use sp_runtime::{proving_trie::BasicProvingTrie, traits::BlakeTwo256};
+	use sp_runtime::{proving_trie::base16::BasicProvingTrie, traits::BlakeTwo256};
 
 	#[test]
 	fn verify_binary_merkle_tree_prover_works() {

diff --git a/substrate/primitives/runtime/src/proving_trie/base16.rs b/substrate/primitives/runtime/src/proving_trie/base16.rs
@@ -101,7 +101,7 @@ where
 
 	/// Query a value contained within the current trie. Returns `None` if the
 	/// nodes within the current `MemoryDB` are insufficient to query the item.
-	fn query(&self, key: Key) -> Option<Value> {
+	fn query(&self, key: &Key) -> Option<Value> {
 		let trie = TrieDBBuilder::new(&self.db, &self.root).build();
 		key.using_encoded(|s| trie.get(s))
 			.ok()?
@@ -112,18 +112,24 @@ where
 	///
 	/// This function makes a proof with latest substrate trie format (`LayoutV1`), and is not
 	/// compatible with `LayoutV0`.
-	fn create_proof(&self, key: Key) -> Result<Vec<u8>, DispatchError> {
-		self.create_multi_value_proof(&[key])
+	fn create_proof(&self, key: &Key) -> Result<Vec<u8>, DispatchError> {
+		sp_trie::generate_trie_proof::<LayoutV1<Hashing>, _, _, _>(
+			&self.db,
+			self.root,
+			&[key.encode()],
+		)
+		.map_err(|err| TrieError::from(*err).into())
+		.map(|structured_proof| structured_proof.encode())
 	}
 
 	/// Verify the existence of `key` and `value` in a given trie root and proof.
 	///
 	/// Proofs must be created with latest substrate trie format (`LayoutV1`).
 	fn verify_proof(
-		root: Hashing::Out,
+		root: &Hashing::Out,
 		proof: &[u8],
-		key: Key,
-		value: Value,
+		key: &Key,
+		value: &Value,
 	) -> Result<(), DispatchError> {
 		verify_proof::<Hashing, Key, Value>(root, proof, key, value)
 	}
@@ -133,10 +139,10 @@ where
 ///
 /// Proofs must be created with latest substrate trie format (`LayoutV1`).
 pub fn verify_proof<Hashing, Key, Value>(
-	root: Hashing::Out,
+	root: &Hashing::Out,
 	proof: &[u8],
-	key: Key,
-	value: Value,
+	key: &Key,
+	value: &Value,
 ) -> Result<(), DispatchError>
 where
 	Hashing: sp_core::Hasher,
@@ -157,7 +163,7 @@ where
 ///
 /// Proofs must be created with latest substrate trie format (`LayoutV1`).
 pub fn verify_multi_value_proof<Hashing, Key, Value>(
-	root: Hashing::Out,
+	root: &Hashing::Out,
 	proof: &[u8],
 	items: &[(Key, Value)],
 ) -> Result<(), DispatchError>
@@ -211,11 +217,11 @@ mod tests {
 		assert!(root != empty_root());
 
 		// Assert valid keys are queryable.
-		assert_eq!(balance_trie.query(6u32), Some(6u128));
-		assert_eq!(balance_trie.query(9u32), Some(9u128));
-		assert_eq!(balance_trie.query(69u32), Some(69u128));
+		assert_eq!(balance_trie.query(&6u32), Some(6u128));
+		assert_eq!(balance_trie.query(&9u32), Some(9u128));
+		assert_eq!(balance_trie.query(&69u32), Some(69u128));
 		// Invalid key returns none.
-		assert_eq!(balance_trie.query(6969u32), None);
+		assert_eq!(balance_trie.query(&6969u32), None);
 
 		balance_trie
 	}
@@ -232,22 +238,24 @@ mod tests {
 		let root = *balance_trie.root();
 
 		// Create a proof for a valid key.
-		let proof = balance_trie.create_proof(6u32).unwrap();
+		let proof = balance_trie.create_proof(&6u32).unwrap();
 
 		// Assert key is provable, all other keys are invalid.
 		for i in 0..200u32 {
 			if i == 6 {
 				assert_eq!(
-					verify_proof::<BlakeTwo256, _, _>(root, &proof, i, u128::from(i)),
+					verify_proof::<BlakeTwo256, _, _>(&root, &proof, &i, &u128::from(i)),
 					Ok(())
 				);
 				// Wrong value is invalid.
 				assert_eq!(
-					verify_proof::<BlakeTwo256, _, _>(root, &proof, i, u128::from(i + 1)),
+					verify_proof::<BlakeTwo256, _, _>(&root, &proof, &i, &u128::from(i + 1)),
 					Err(TrieError::RootMismatch.into())
 				);
 			} else {
-				assert!(verify_proof::<BlakeTwo256, _, _>(root, &proof, i, u128::from(i)).is_err());
+				assert!(
+					verify_proof::<BlakeTwo256, _, _>(&root, &proof, &i, &u128::from(i)).is_err()
+				);
 			}
 		}
 	}
@@ -261,7 +269,7 @@ mod tests {
 		let proof = balance_trie.create_multi_value_proof(&[6u32, 9u32, 69u32]).unwrap();
 		let items = [(6u32, 6u128), (9u32, 9u128), (69u32, 69u128)];
 
-		assert_eq!(verify_multi_value_proof::<BlakeTwo256, _, _>(root, &proof, &items), Ok(()));
+		assert_eq!(verify_multi_value_proof::<BlakeTwo256, _, _>(&root, &proof, &items), Ok(()));
 	}
 
 	#[test]
@@ -270,23 +278,23 @@ mod tests {
 		let root = *balance_trie.root();
 
 		// Create a proof for a valid key.
-		let proof = balance_trie.create_proof(6u32).unwrap();
+		let proof = balance_trie.create_proof(&6u32).unwrap();
 
 		// Correct data verifies successfully
-		assert_eq!(verify_proof::<BlakeTwo256, _, _>(root, &proof, 6u32, 6u128), Ok(()));
+		assert_eq!(verify_proof::<BlakeTwo256, _, _>(&root, &proof, &6u32, &6u128), Ok(()));
 
 		// Fail to verify proof with wrong root
 		assert_eq!(
-			verify_proof::<BlakeTwo256, _, _>(Default::default(), &proof, 6u32, 6u128),
+			verify_proof::<BlakeTwo256, _, _>(&Default::default(), &proof, &6u32, &6u128),
 			Err(TrieError::RootMismatch.into())
 		);
 
 		// Crete a bad proof.
-		let bad_proof = balance_trie.create_proof(99u32).unwrap();
+		let bad_proof = balance_trie.create_proof(&99u32).unwrap();
 
 		// Fail to verify data with the wrong proof
 		assert_eq!(
-			verify_proof::<BlakeTwo256, _, _>(root, &bad_proof, 6u32, 6u128),
+			verify_proof::<BlakeTwo256, _, _>(&root, &bad_proof, &6u32, &6u128),
 			Err(TrieError::ExtraneousHashReference.into())
 		);
 	}

diff --git a/substrate/primitives/runtime/src/proving_trie/base2.rs b/substrate/primitives/runtime/src/proving_trie/base2.rs
@@ -64,21 +64,21 @@ where
 
 	/// Query a value contained within the current trie. Returns `None` if the
 	/// nodes within the current `db` are insufficient to query the item.
-	fn query(&self, key: Key) -> Option<Value> {
+	fn query(&self, key: &Key) -> Option<Value> {
 		self.db.get(&key).cloned()
 	}
 
 	/// Create a compact merkle proof needed to prove a single key and its value are in the trie.
 	/// Returns `None` if the nodes within the current `db` are insufficient to create a
 	/// proof.
-	fn create_proof(&self, key: Key) -> Result<Vec<u8>, DispatchError> {
+	fn create_proof(&self, key: &Key) -> Result<Vec<u8>, DispatchError> {
 		let mut encoded = Vec::with_capacity(self.db.len());
 		let mut found_index = None;
 
 		// Find the index of our key, and encode the (key, value) pair.
 		for (i, (k, v)) in self.db.iter().enumerate() {
 			// If we found the key we are looking for, save it.
-			if *k == key {
+			if k == key {
 				found_index = Some(i);
 			}
 
@@ -92,21 +92,21 @@ where
 
 	/// Verify the existence of `key` and `value` in a given trie root and proof.
 	fn verify_proof(
-		root: Hashing::Out,
+		root: &Hashing::Out,
 		proof: &[u8],
-		key: Key,
-		value: Value,
+		key: &Key,
+		value: &Value,
 	) -> Result<(), DispatchError> {
 		verify_proof::<Hashing, Key, Value>(root, proof, key, value)
 	}
 }
 
 /// Verify the existence of `key` and `value` in a given trie root and proof.
 pub fn verify_proof<Hashing, Key, Value>(
-	root: Hashing::Out,
+	root: &Hashing::Out,
 	proof: &[u8],
-	key: Key,
-	value: Value,
+	key: &Key,
+	value: &Value,
 ) -> Result<(), DispatchError>
 where
 	Hashing: sp_core::Hasher,
@@ -116,11 +116,11 @@ where
 {
 	let decoded_proof: MerkleProof<Hashing::Out, Vec<u8>> =
 		Decode::decode(&mut &proof[..]).map_err(|_| TrieError::IncompleteProof)?;
-	if root != decoded_proof.root {
+	if *root != decoded_proof.root {
 		return Err(TrieError::RootMismatch.into());
 	}
 
-	if (&key, &value).encode() != decoded_proof.leaf {
+	if (key, value).encode() != decoded_proof.leaf {
 		return Err(TrieError::ValueMismatch.into());
 	}
 
@@ -168,9 +168,9 @@ mod tests {
 		assert!(root != empty_root());
 
 		// Assert valid keys are queryable.
-		assert_eq!(balance_trie.query(6u32), Some(6u128));
-		assert_eq!(balance_trie.query(9u32), Some(9u128));
-		assert_eq!(balance_trie.query(69u32), Some(69u128));
+		assert_eq!(balance_trie.query(&6u32), Some(6u128));
+		assert_eq!(balance_trie.query(&9u32), Some(9u128));
+		assert_eq!(balance_trie.query(&69u32), Some(69u128));
 
 		balance_trie
 	}
@@ -187,33 +187,24 @@ mod tests {
 		let root = *balance_trie.root();
 
 		// Create a proof for a valid key.
-		let proof = balance_trie.create_proof(6u32).unwrap();
+		let proof = balance_trie.create_proof(&6u32).unwrap();
 
 		// Assert key is provable, all other keys are invalid.
 		for i in 0..200u32 {
 			if i == 6 {
 				assert_eq!(
-					verify_proof::<BlakeTwo256, _, _>(root, &proof, i, u128::from(i)),
+					verify_proof::<BlakeTwo256, _, _>(&root, &proof, &i, &u128::from(i)),
 					Ok(())
 				);
 				// Wrong value is invalid.
 				assert_eq!(
-					verify_proof::<BlakeTwo256, _, _>(
-						root,
-						&proof,
-						i,
-						u128::from(i + 1)
-					),
+					verify_proof::<BlakeTwo256, _, _>(&root, &proof, &i, &u128::from(i + 1)),
 					Err(TrieError::ValueMismatch.into())
 				);
 			} else {
-				assert!(verify_proof::<BlakeTwo256, _, _>(
-					root,
-					&proof,
-					i,
-					u128::from(i)
-				)
-				.is_err());
+				assert!(
+					verify_proof::<BlakeTwo256, _, _>(&root, &proof, &i, &u128::from(i)).is_err()
+				);
 			}
 		}
 	}
@@ -224,23 +215,20 @@ mod tests {
 		let root = *balance_trie.root();
 
 		// Create a proof for a valid key.
-		let proof = balance_trie.create_proof(6u32).unwrap();
+		let proof = balance_trie.create_proof(&6u32).unwrap();
 
 		// Correct data verifies successfully
-		assert_eq!(
-			verify_proof::<BlakeTwo256, _, _>(root, &proof, 6u32, 6u128),
-			Ok(())
-		);
+		assert_eq!(verify_proof::<BlakeTwo256, _, _>(&root, &proof, &6u32, &6u128), Ok(()));
 
 		// Fail to verify proof with wrong root
 		assert_eq!(
-			verify_proof::<BlakeTwo256, _, _>(Default::default(), &proof, 6u32, 6u128),
+			verify_proof::<BlakeTwo256, _, _>(&Default::default(), &proof, &6u32, &6u128),
 			Err(TrieError::RootMismatch.into())
 		);
 
 		// Fail to verify proof with wrong data
 		assert_eq!(
-			verify_proof::<BlakeTwo256, _, _>(root, &[], 6u32, 6u128),
+			verify_proof::<BlakeTwo256, _, _>(&root, &[], &6u32, &6u128),
 			Err(TrieError::IncompleteProof.into())
 		);
 	}

diff --git a/substrate/primitives/runtime/src/proving_trie/mod.rs b/substrate/primitives/runtime/src/proving_trie/mod.rs
@@ -127,15 +127,15 @@ where
 	fn root(&self) -> &Hashing::Out;
 	/// Query a value contained within the current trie. Returns `None` if the
 	/// the value does not exist in the trie.
-	fn query(&self, key: Key) -> Option<Value>;
+	fn query(&self, key: &Key) -> Option<Value>;
 	/// Create a proof that can be used to verify a key and its value are in the trie.
-	fn create_proof(&self, key: Key) -> Result<Vec<u8>, DispatchError>;
+	fn create_proof(&self, key: &Key) -> Result<Vec<u8>, DispatchError>;
 	/// Verify the existence of `key` and `value` in a given trie root and proof.
 	fn verify_proof(
-		root: Hashing::Out,
+		root: &Hashing::Out,
 		proof: &[u8],
-		key: Key,
-		value: Value,
+		key: &Key,
+		value: &Value,
 	) -> Result<(), DispatchError>;
 }
 
@@ -154,19 +154,19 @@ mod tests {
 	fn basic_api_usage_base_2() {
 		let balance_trie = BalanceTrie2::generate_for((0..100u32).map(|i| (i, i.into()))).unwrap();
 		let root = *balance_trie.root();
-		assert_eq!(balance_trie.query(69), Some(69));
-		assert_eq!(balance_trie.query(6969), None);
-		let proof = balance_trie.create_proof(69u32).unwrap();
-		assert_eq!(BalanceTrie2::verify_proof(root, &proof, 69u32, 69u128), Ok(()));
+		assert_eq!(balance_trie.query(&69), Some(69));
+		assert_eq!(balance_trie.query(&6969), None);
+		let proof = balance_trie.create_proof(&69u32).unwrap();
+		assert_eq!(BalanceTrie2::verify_proof(&root, &proof, &69u32, &69u128), Ok(()));
 	}
 
 	#[test]
 	fn basic_api_usage_base_16() {
 		let balance_trie = BalanceTrie16::generate_for((0..100u32).map(|i| (i, i.into()))).unwrap();
 		let root = *balance_trie.root();
-		assert_eq!(balance_trie.query(69), Some(69));
-		assert_eq!(balance_trie.query(6969), None);
-		let proof = balance_trie.create_proof(69u32).unwrap();
-		assert_eq!(BalanceTrie16::verify_proof(root, &proof, 69u32, 69u128), Ok(()));
+		assert_eq!(balance_trie.query(&69), Some(69));
+		assert_eq!(balance_trie.query(&6969), None);
+		let proof = balance_trie.create_proof(&69u32).unwrap();
+		assert_eq!(BalanceTrie16::verify_proof(&root, &proof, &69u32, &69u128), Ok(()));
 	}
 }