Improve APIs for Tries in Runtime (#5756)

This is a refactor and improvement from:
#3881

- `sp_runtime::proving_trie` now exposes a `BasicProvingTrie` for both
`base2` and `base16`.
- APIs for `base16` are more focused on single value proofs, also
aligning their APIs with the `base2` trie
- A `ProvingTrie` trait is included which wraps both the `base2` and
`base16` trie, and exposes all APIs needed for an end to end scenario.
- A `ProofToHashes` trait is exposed which can allow us to write proper
benchmarks for the merkle trie.

---------

Co-authored-by: Ankan <10196091+Ank4n@users.noreply.github.com>
Co-authored-by: Adrian Catangiu <adrian@parity.io>

prdoc/pr_5756.prdoc

@@ -0,0 +1,16 @@
+# Schema: Polkadot SDK PRDoc Schema (prdoc) v1.0.0
+# See doc at https://raw.githubusercontent.com/paritytech/polkadot-sdk/master/prdoc/schema_user.json
+
+title: Improve APIs for Tries in Runtime
+
+doc:
+  - audience: Runtime Dev
+    description: |
+      This PR introduces a trait `ProvingTrie` which has all the function you need to use tries in the runtime.
+      This trait includes the ability to create, query, and prove data in a trie. Another trait `ProofToHashes`
+      allows developers to express the computational complexity of proof verification using the proof data.
+crates:
+  - name: sp-runtime
+    bump: major
+  - name: frame-support
+    bump: major

substrate/frame/support/src/traits/proving.rs

@@ -20,6 +20,10 @@
 use alloc::vec::Vec;
 use codec::{Decode, Encode};
 use sp_core::Hasher;
+use sp_runtime::DispatchError;
+
+// Re-export the `proving_trie` types and traits.
+pub use sp_runtime::proving_trie::*;
 
 /// Something that can verify the existence of some data in a given proof.
 pub trait VerifyExistenceProof {
@@ -31,7 +35,7 @@ pub trait VerifyExistenceProof {
 	/// Verify the given `proof`.
 	///
 	/// Ensures that the `proof` was build for `root` and returns the proved data.
-	fn verify_proof(proof: Self::Proof, root: &Self::Hash) -> Result<Vec<u8>, ()>;
+	fn verify_proof(proof: Self::Proof, root: &Self::Hash) -> Result<Vec<u8>, DispatchError>;
 }
 
 /// Implements [`VerifyExistenceProof`] using a binary merkle tree.
@@ -44,9 +48,9 @@ where
 	type Proof = binary_merkle_tree::MerkleProof<H::Out, Vec<u8>>;
 	type Hash = H::Out;
 
-	fn verify_proof(proof: Self::Proof, root: &Self::Hash) -> Result<Vec<u8>, ()> {
+	fn verify_proof(proof: Self::Proof, root: &Self::Hash) -> Result<Vec<u8>, DispatchError> {
 		if proof.root != *root {
-			return Err(());
+			return Err(TrieError::RootMismatch.into());
 		}
 
 		if binary_merkle_tree::verify_proof::<H, _, _>(
@@ -58,13 +62,25 @@ where
 		) {
 			Ok(proof.leaf)
 		} else {
-			Err(())
+			Err(TrieError::IncompleteProof.into())
 		}
 	}
 }
 
+impl<H: Hasher> ProofToHashes for BinaryMerkleTreeProver<H> {
+	type Proof = binary_merkle_tree::MerkleProof<H::Out, Vec<u8>>;
+
+	// This base 2 merkle trie includes a `proof` field which is a `Vec<Hash>`.
+	// The length of this vector tells us the depth of the proof, and how many
+	// hashes we need to calculate.
+	fn proof_to_hashes(proof: &Self::Proof) -> Result<u32, DispatchError> {
+		let depth = proof.proof.len();
+		Ok(depth as u32)
+	}
+}
+
 /// Proof used by [`SixteenPatriciaMerkleTreeProver`] for [`VerifyExistenceProof`].
-#[derive(Encode, Decode)]
+#[derive(Encode, Decode, Clone)]
 pub struct SixteenPatriciaMerkleTreeExistenceProof {
 	/// The key of the value to prove.
 	pub key: Vec<u8>,
@@ -81,21 +97,35 @@ impl<H: Hasher> VerifyExistenceProof for SixteenPatriciaMerkleTreeProver<H> {
 	type Proof = SixteenPatriciaMerkleTreeExistenceProof;
 	type Hash = H::Out;
 
-	fn verify_proof(proof: Self::Proof, root: &Self::Hash) -> Result<Vec<u8>, ()> {
+	fn verify_proof(proof: Self::Proof, root: &Self::Hash) -> Result<Vec<u8>, DispatchError> {
 		sp_trie::verify_trie_proof::<sp_trie::LayoutV1<H>, _, _, _>(
 			&root,
 			&proof.proof,
 			[&(&proof.key, Some(&proof.value))],
 		)
-		.map_err(drop)
+		.map_err(|err| TrieError::from(err).into())
 		.map(|_| proof.value)
 	}
 }
 
+impl<H: Hasher> ProofToHashes for SixteenPatriciaMerkleTreeProver<H> {
+	type Proof = SixteenPatriciaMerkleTreeExistenceProof;
+
+	// This base 16 trie uses a raw proof of `Vec<Vec<u8>`, where the length of the first `Vec`
+	// is the depth of the trie. We can use this to predict the number of hashes.
+	fn proof_to_hashes(proof: &Self::Proof) -> Result<u32, DispatchError> {
+		let depth = proof.proof.len();
+		Ok(depth as u32)
+	}
+}
+
 #[cfg(test)]
 mod tests {
 	use super::*;
-	use sp_runtime::{proving_trie::BasicProvingTrie, traits::BlakeTwo256};
+	use sp_runtime::{
+		proving_trie::{base16::BasicProvingTrie, ProvingTrie},
+		traits::BlakeTwo256,
+	};
 
 	#[test]
 	fn verify_binary_merkle_tree_prover_works() {
@@ -113,23 +143,87 @@ mod tests {
 
 	#[test]
 	fn verify_sixteen_patricia_merkle_tree_prover_works() {
-		let trie = BasicProvingTrie::<BlakeTwo256, u32, &[u8]>::generate_for(vec![
-			(0u32, &b"hey"[..]),
-			(1u32, &b"yes"[..]),
+		let trie = BasicProvingTrie::<BlakeTwo256, u32, _>::generate_for(vec![
+			(0u32, String::from("hey")),
+			(1u32, String::from("yes")),
 		])
 		.unwrap();
-		let proof = trie.create_single_value_proof(1u32).unwrap();
+		let proof = trie.create_proof(&1u32).unwrap();
+		let structured_proof: Vec<Vec<u8>> = Decode::decode(&mut &proof[..]).unwrap();
 		let root = *trie.root();
 
 		let proof = SixteenPatriciaMerkleTreeExistenceProof {
 			key: 1u32.encode(),
-			value: b"yes"[..].encode(),
-			proof,
+			value: String::from("yes").encode(),
+			proof: structured_proof,
 		};
 
 		assert_eq!(
 			SixteenPatriciaMerkleTreeProver::<BlakeTwo256>::verify_proof(proof, &root).unwrap(),
-			b"yes"[..].encode()
+			String::from("yes").encode()
 		);
 	}
+
+	#[test]
+	fn proof_to_hashes_sixteen() {
+		let mut i: u32 = 1;
+
+		// Compute log base 16 and round up
+		let log16 = |x: u32| -> u32 {
+			let x_f64 = x as f64;
+			let log16_x = (x_f64.ln() / 16_f64.ln()).ceil();
+			log16_x as u32
+		};
+
+		while i < 10_000_000 {
+			let trie = BasicProvingTrie::<BlakeTwo256, u32, _>::generate_for(
+				(0..i).map(|i| (i, u128::from(i))),
+			)
+			.unwrap();
+			let proof = trie.create_proof(&0).unwrap();
+			let structured_proof: Vec<Vec<u8>> = Decode::decode(&mut &proof[..]).unwrap();
+			let root = *trie.root();
+
+			let proof = SixteenPatriciaMerkleTreeExistenceProof {
+				key: 0u32.encode(),
+				value: 0u128.encode(),
+				proof: structured_proof,
+			};
+			let hashes =
+				SixteenPatriciaMerkleTreeProver::<BlakeTwo256>::proof_to_hashes(&proof).unwrap();
+			let log16 = log16(i).max(1);
+			assert_eq!(hashes, log16);
+
+			assert_eq!(
+				SixteenPatriciaMerkleTreeProver::<BlakeTwo256>::verify_proof(proof.clone(), &root)
+					.unwrap(),
+				proof.value
+			);
+
+			i = i * 10;
+		}
+	}
+
+	#[test]
+	fn proof_to_hashes_binary() {
+		let mut i: u32 = 1;
+		while i < 10_000_000 {
+			let proof = binary_merkle_tree::merkle_proof::<BlakeTwo256, _, _>(
+				(0..i).map(|i| u128::from(i).encode()),
+				0,
+			);
+			let root = proof.root;
+
+			let hashes = BinaryMerkleTreeProver::<BlakeTwo256>::proof_to_hashes(&proof).unwrap();
+			let log2 = (i as f64).log2().ceil() as u32;
+			assert_eq!(hashes, log2);
+
+			assert_eq!(
+				BinaryMerkleTreeProver::<BlakeTwo256>::verify_proof(proof, &root).unwrap(),
+				0u128.encode()
+			);
+
+			i = i * 10;
+		}
+	}
 }

substrate/primitives/runtime/Cargo.toml

@@ -36,6 +36,7 @@ sp-trie = { workspace = true }
 sp-weights = { workspace = true }
 docify = { workspace = true }
 tracing = { workspace = true, features = ["log"], default-features = false }
+binary-merkle-tree = { workspace = true }
 
 simple-mermaid = { version = "0.1.1", optional = true }
 
@@ -53,6 +54,7 @@ runtime-benchmarks = []
 try-runtime = []
 default = ["std"]
 std = [
+	"binary-merkle-tree/std",
 	"codec/std",
 	"either/use_std",
 	"hash256-std-hasher/std",