Ancestry proof benchmarks

substrate/frame/beefy-mmr/src/benchmarking.rs

@@ -20,15 +20,92 @@
 #![cfg(feature = "runtime-benchmarks")]
 
 use super::*;
+use crate::Pallet as BeefyMmr;
+use codec::Encode;
 use frame_benchmarking::v2::*;
+use frame_support::traits::Hooks;
+use frame_system::{Config as SystemConfig, Pallet as System};
+use pallet_mmr::{Nodes, Pallet as Mmr};
+use sp_consensus_beefy::Payload;
+use sp_runtime::traits::Zero;
+
+pub trait Config:
+	pallet_mmr::Config<Hashing = sp_consensus_beefy::MmrHashing> + crate::Config
+{
+}
+
+impl<T> Config for T where
+	T: pallet_mmr::Config<Hashing = sp_consensus_beefy::MmrHashing> + crate::Config
+{
+}
+
+fn init_block<T: Config>(block_num: u32) {
+	let block_num = block_num.into();
+	System::<T>::initialize(&block_num, &<T as SystemConfig>::Hash::default(), &Default::default());
+	Mmr::<T>::on_initialize(block_num);
+}
 
 #[benchmarks]
 mod benchmarks {
 	use super::*;
 
 	#[benchmark]
-	fn n_items_proof_is_non_canonical(n: Linear<1, 64>) {
+	fn extract_validation_context() {
+		init_block::<T>(0);
+		let header = System::<T>::finalize();
+		frame_system::BlockHash::<T>::insert(BlockNumberFor::<T>::zero(), header.hash());
+
+		let validation_context;
+		#[block]
+		{
+			validation_context =
+				<BeefyMmr<T> as AncestryHelper<HeaderFor<T>>>::extract_validation_context(header);
+		}
+
+		assert!(validation_context.is_some());
+	}
+
+	#[benchmark]
+	fn read_peak() {
+		init_block::<T>(0);
+
+		let peak;
+		#[block]
+		{
+			peak = Nodes::<T>::get(0)
+		}
+
+		assert!(peak.is_some());
+	}
+
+	/// Generate ancestry proofs with `n` nodes and benchmark the verification logic.
+	/// These proofs are inflated, containing all the leafs, so we won't read any peak during
+	/// the verification. We need to account for the peaks separately.
+	#[benchmark]
+	fn n_items_proof_is_non_canonical(n: Linear<2, 512>) {
+		for block_num in 1..=n {
+			init_block::<T>(block_num);
+		}
+		let proof = Mmr::<T>::generate_mock_ancestry_proof().unwrap();
+		assert_eq!(proof.items.len(), n as usize);
+
+		let is_non_canonical;
 		#[block]
-		{}
+		{
+			is_non_canonical = <BeefyMmr<T> as AncestryHelper<HeaderFor<T>>>::is_non_canonical(
+				&Commitment {
+					payload: Payload::from_single_entry(
+						known_payloads::MMR_ROOT_ID,
+						MerkleRootOf::<T>::default().encode(),
+					),
+					block_number: n.into(),
+					validator_set_id: 0,
+				},
+				proof,
+				Mmr::<T>::mmr_root(),
+			);
+		};
+
+		assert_eq!(is_non_canonical, true);
 	}
 }

substrate/frame/beefy-mmr/src/lib.rs

@@ -38,22 +38,20 @@ extern crate alloc;
 use sp_runtime::{
 	generic::OpaqueDigestItemId,
 	traits::{Convert, Header, Member},
+	SaturatedConversion,
 };
 
 use alloc::vec::Vec;
 use codec::Decode;
-use pallet_mmr::{primitives::AncestryProof, LeafDataProvider, ParentNumberAndHash};
+use pallet_mmr::{primitives::AncestryProof, LeafDataProvider, NodesUtils, ParentNumberAndHash};
 use sp_consensus_beefy::{
 	known_payloads,
 	mmr::{BeefyAuthoritySet, BeefyDataProvider, BeefyNextAuthoritySet, MmrLeaf, MmrLeafVersion},
 	AncestryHelper, AncestryHelperWeightInfo, Commitment, ConsensusLog,
 	ValidatorSet as BeefyValidatorSet,
 };
 
-use frame_support::{
-	crypto::ecdsa::ECDSAExt, pallet_prelude::Weight, traits::Get,
-	weights::constants::RocksDbWeight as DbWeight,
-};
+use frame_support::{crypto::ecdsa::ECDSAExt, pallet_prelude::Weight, traits::Get};
 use frame_system::pallet_prelude::{BlockNumberFor, HeaderFor};
 
 pub use pallet::*;
@@ -278,14 +276,23 @@ impl<T: Config> AncestryHelperWeightInfo<HeaderFor<T>> for Pallet<T>
 where
 	T: pallet_mmr::Config<Hashing = sp_consensus_beefy::MmrHashing>,
 {
-	fn extract_validation_context(_header: &HeaderFor<T>) -> Weight {
-		// We're only reading from `BlockHash` and then searching in a small vec,
-		// which should be insignificant.
-		DbWeight::get().reads(1)
+	fn extract_validation_context() -> Weight {
+		<T as Config>::WeightInfo::extract_validation_context()
 	}
 
 	fn is_non_canonical(proof: &<Self as AncestryHelper<HeaderFor<T>>>::Proof) -> Weight {
-		<T as Config>::WeightInfo::n_items_proof_is_non_canonical(proof.items.len() as u32)
+		let mmr_utils = NodesUtils::new(proof.leaf_count);
+		let num_peaks = mmr_utils.number_of_peaks();
+
+		// The approximated cost of verifying an ancestry proof with `n` nodes.
+		// We add the previous peaks to the total number of nodes,
+		// since they have to be processed as well.
+		<T as Config>::WeightInfo::n_items_proof_is_non_canonical(
+			proof.items.len().saturating_add(proof.prev_peaks.len()).saturated_into(),
+		)
+		// `n_items_proof_is_non_canonical()` uses inflated proofs that contain all the leafs,
+		// where no peak needs to be read. So we need to also add the cost of reading the peaks.
+		.saturating_add(<T as Config>::WeightInfo::read_peak().saturating_mul(num_peaks))
 	}
 }
 

substrate/frame/beefy-mmr/src/tests.rs

@@ -40,8 +40,6 @@ fn init_block(block: u64, maybe_parent_hash: Option<H256>) {
 	System::initialize(&block, &parent_hash, &Default::default());
 	Session::on_initialize(block);
 	Mmr::on_initialize(block);
-	Beefy::on_initialize(block);
-	BeefyMmr::on_initialize(block);
 }
 
 pub fn beefy_log(log: ConsensusLog<BeefyId>) -> DigestItem {
@@ -268,7 +266,7 @@ fn is_non_canonical_should_work_correctly() {
 	ext.register_extension(OffchainWorkerExt::new(offchain));
 
 	ext.execute_with(|| {
-		let valid_proof = Mmr::generate_ancestry_proof(250, None).unwrap();
+		let valid_proof = BeefyMmr::generate_proof(250, None).unwrap();
 		let mut invalid_proof = valid_proof.clone();
 		invalid_proof.items.push((300, Default::default()));
 
@@ -343,7 +341,7 @@ fn is_non_canonical_should_work_correctly() {
 		// - should return false, if the commitment is targeting the canonical chain
 		// - should return true if the commitment is NOT targeting the canonical chain
 		for prev_block_number in 1usize..=500 {
-			let proof = Mmr::generate_ancestry_proof(prev_block_number as u64, None).unwrap();
+			let proof = BeefyMmr::generate_proof(prev_block_number as u64, None).unwrap();
 
 			assert_eq!(
 				BeefyMmr::is_non_canonical(

substrate/frame/beefy/src/mock.rs

@@ -133,7 +133,7 @@ impl<Header: HeaderT> AncestryHelper<Header> for MockAncestryHelper {
 }
 
 impl<Header: HeaderT> AncestryHelperWeightInfo<Header> for MockAncestryHelper {
-	fn extract_validation_context(_header: &Header) -> Weight {
+	fn extract_validation_context() -> Weight {
 		unimplemented!()
 	}
 

substrate/frame/merkle-mountain-range/src/lib.rs

@@ -439,6 +439,14 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		mmr.generate_ancestry_proof(prev_leaf_count)
 	}
 
+	#[cfg(feature = "runtime-benchmarks")]
+	pub fn generate_mock_ancestry_proof() -> Result<primitives::AncestryProof<HashOf<T, I>>, Error>
+	{
+		let leaf_count = Self::block_num_to_leaf_count(<frame_system::Pallet<T>>::block_number())?;
+		let mmr: ModuleMmr<mmr::storage::OffchainStorage, T, I> = mmr::Mmr::new(leaf_count);
+		mmr.generate_mock_ancestry_proof()
+	}
+
 	pub fn verify_ancestry_proof(
 		root: HashOf<T, I>,
 		ancestry_proof: primitives::AncestryProof<HashOf<T, I>>,

substrate/frame/merkle-mountain-range/src/mmr/mmr.rs

@@ -156,7 +156,7 @@ where
 	}
 
 	/// Return the internal size of the MMR (number of nodes).
-	#[cfg(test)]
+	#[cfg(any(test, feature = "runtime-benchmarks"))]
 	pub fn size(&self) -> NodeIndex {
 		self.mmr.mmr_size()
 	}
@@ -252,4 +252,48 @@ where
 				.collect(),
 		})
 	}
+
+	/// Generate an inflated ancestry proof for the latest leaf in the MMR.
+	///
+	/// The generated proof contains all the leafs in the MMR, so this way we can generate a proof
+	/// with exactly `leaf_count` items.
+	#[cfg(feature = "runtime-benchmarks")]
+	pub fn generate_mock_ancestry_proof(
+		&self,
+	) -> Result<sp_mmr_primitives::AncestryProof<HashOf<T, I>>, Error> {
+		use crate::ModuleMmr;
+		use alloc::vec;
+		use sp_mmr_primitives::mmr_lib::helper;
+
+		let mmr: ModuleMmr<OffchainStorage, T, I> = Mmr::new(self.leaves);
+		let store = <Storage<OffchainStorage, T, I, L>>::default();
+
+		let mut prev_peaks = vec![];
+		for peak_pos in helper::get_peaks(mmr.size()) {
+			let peak = store
+				.get_elem(peak_pos)
+				.map_err(|_| Error::GenerateProof)?
+				.ok_or(Error::GenerateProof)?
+				.hash();
+			prev_peaks.push(peak);
+		}
+
+		let mut proof_items = vec![];
+		for leaf_idx in 0..self.leaves {
+			let leaf_pos = NodesUtils::leaf_index_to_leaf_node_index(leaf_idx);
+			let leaf = store
+				.get_elem(leaf_pos)
+				.map_err(|_| Error::GenerateProof)?
+				.ok_or(Error::GenerateProof)?
+				.hash();
+			proof_items.push((leaf_pos, leaf));
+		}
+
+		Ok(sp_mmr_primitives::AncestryProof {
+			prev_peaks,
+			prev_leaf_count: self.leaves,
+			leaf_count: self.leaves,
+			items: proof_items,
+		})
+	}
 }