Refactor Beefy MMR and remove parachain specific implementations (#10664)

* refactor beefy mmr

* use plain vector of bytes for leaf extra

* update comment

* update comments

* remove unused vars

* Use sp_std::vec::Vec

Co-authored-by: Adrian Catangiu <[email protected]>

* make extra data generic

* fix tests

* refactor beefy-mmr

* Update frame/beefy-mmr/src/lib.rs

* minor fix

* fmt

* Update frame/beefy-mmr/src/lib.rs

Co-authored-by: Adrian Catangiu <[email protected]>

Author: Wizdave97

frame/beefy-mmr/src/lib.rs

@@ -29,17 +29,16 @@
 //! The MMR leaf contains:
 //! 1. Block number and parent block hash.
 //! 2. Merkle Tree Root Hash of next BEEFY validator set.
-//! 3. Merkle Tree Root Hash of current parachain heads state.
+//! 3. Arbitrary extra leaf data to be used by downstream pallets to include custom data.
 //!
 //! and thanks to versioning can be easily updated in the future.
 
-use sp_runtime::traits::{Convert, Hash};
+use sp_runtime::traits::{Convert, Hash, Member};
 use sp_std::prelude::*;
 
-use beefy_primitives::mmr::{BeefyNextAuthoritySet, MmrLeaf, MmrLeafVersion};
+use beefy_primitives::mmr::{BeefyDataProvider, BeefyNextAuthoritySet, MmrLeaf, MmrLeafVersion};
 use pallet_mmr::primitives::LeafDataProvider;
 
-use codec::Encode;
 use frame_support::traits::Get;
 
 pub use pallet::*;
@@ -90,23 +89,6 @@ impl Convert<beefy_primitives::crypto::AuthorityId, Vec<u8>> for BeefyEcdsaToEth
 }
 
 type MerkleRootOf<T> = <T as pallet_mmr::Config>::Hash;
-type ParaId = u32;
-type ParaHead = Vec<u8>;
-
-/// A type that is able to return current list of parachain heads that end up in the MMR leaf.
-pub trait ParachainHeadsProvider {
-	/// Return a list of tuples containing a `ParaId` and Parachain Header data (ParaHead).
-	///
-	/// The returned data does not have to be sorted.
-	fn parachain_heads() -> Vec<(ParaId, ParaHead)>;
-}
-
-/// A default implementation for runtimes without parachains.
-impl ParachainHeadsProvider for () {
-	fn parachain_heads() -> Vec<(ParaId, ParaHead)> {
-		Default::default()
-	}
-}
 
 #[frame_support::pallet]
 pub mod pallet {
@@ -138,12 +120,11 @@ pub mod pallet {
 		/// efficiency reasons.
 		type BeefyAuthorityToMerkleLeaf: Convert<<Self as pallet_beefy::Config>::BeefyId, Vec<u8>>;
 
-		/// Retrieve a list of current parachain heads.
-		///
-		/// The trait is implemented for `paras` module, but since not all chains might have
-		/// parachains, and we want to keep the MMR leaf structure uniform, it's possible to use
-		/// `()` as well to simply put dummy data to the leaf.
-		type ParachainHeads: ParachainHeadsProvider;
+		/// The type expected for the leaf extra data
+		type LeafExtra: Member + codec::FullCodec;
+
+		/// Retrieve arbitrary data that should be added to the mmr leaf
+		type BeefyDataProvider: BeefyDataProvider<Self::LeafExtra>;
 	}
 
 	/// Details of next BEEFY authority set.
@@ -163,13 +144,14 @@ where
 		<T as frame_system::Config>::BlockNumber,
 		<T as frame_system::Config>::Hash,
 		MerkleRootOf<T>,
+		T::LeafExtra,
 	>;
 
 	fn leaf_data() -> Self::LeafData {
 		MmrLeaf {
 			version: T::LeafVersion::get(),
 			parent_number_and_hash: frame_system::Pallet::<T>::leaf_data(),
-			parachain_heads: Pallet::<T>::parachain_heads_merkle_root(),
+			leaf_extra: T::BeefyDataProvider::extra_data(),
 			beefy_next_authority_set: Pallet::<T>::update_beefy_next_authority_set(),
 		}
 	}
@@ -188,23 +170,6 @@ impl<T: Config> Pallet<T>
 where
 	MerkleRootOf<T>: From<beefy_merkle_tree::Hash> + Into<beefy_merkle_tree::Hash>,
 {
-	/// Returns latest root hash of a merkle tree constructed from all active parachain headers.
-	///
-	/// The leafs are sorted by `ParaId` to allow more efficient lookups and non-existence proofs.
-	///
-	/// NOTE this does not include parathreads - only parachains are part of the merkle tree.
-	///
-	/// NOTE This is an initial and inefficient implementation, which re-constructs
-	/// the merkle tree every block. Instead we should update the merkle root in
-	/// [Self::on_initialize] call of this pallet and update the merkle tree efficiently (use
-	/// on-chain storage to persist inner nodes).
-	fn parachain_heads_merkle_root() -> MerkleRootOf<T> {
-		let mut para_heads = T::ParachainHeads::parachain_heads();
-		para_heads.sort();
-		let para_heads = para_heads.into_iter().map(|pair| pair.encode());
-		beefy_merkle_tree::merkle_root::<Self, _, _>(para_heads).into()
-	}
-
 	/// Returns details of the next BEEFY authority set.
 	///
 	/// Details contain authority set id, authority set length and a merkle root,

frame/beefy-mmr/src/mock.rs

@@ -18,6 +18,7 @@
 use std::vec;
 
 use beefy_primitives::mmr::MmrLeafVersion;
+use codec::Encode;
 use frame_support::{
 	construct_runtime, parameter_types,
 	sp_io::TestExternalities,
@@ -34,7 +35,9 @@ use sp_runtime::{
 
 use crate as pallet_beefy_mmr;
 
-pub use beefy_primitives::{crypto::AuthorityId as BeefyId, ConsensusLog, BEEFY_ENGINE_ID};
+pub use beefy_primitives::{
+	crypto::AuthorityId as BeefyId, mmr::BeefyDataProvider, ConsensusLog, BEEFY_ENGINE_ID,
+};
 
 impl_opaque_keys! {
 	pub struct MockSessionKeys {
@@ -102,6 +105,7 @@ pub type MmrLeaf = beefy_primitives::mmr::MmrLeaf<
 	<Test as frame_system::Config>::BlockNumber,
 	<Test as frame_system::Config>::Hash,
 	<Test as pallet_mmr::Config>::Hash,
+	Vec<u8>,
 >;
 
 impl pallet_mmr::Config for Test {
@@ -131,13 +135,20 @@ impl pallet_beefy_mmr::Config for Test {
 
 	type BeefyAuthorityToMerkleLeaf = pallet_beefy_mmr::BeefyEcdsaToEthereum;
 
-	type ParachainHeads = DummyParaHeads;
+	type LeafExtra = Vec<u8>;
+
+	type BeefyDataProvider = DummyDataProvider;
 }
 
-pub struct DummyParaHeads;
-impl pallet_beefy_mmr::ParachainHeadsProvider for DummyParaHeads {
-	fn parachain_heads() -> Vec<(pallet_beefy_mmr::ParaId, pallet_beefy_mmr::ParaHead)> {
-		vec![(15, vec![1, 2, 3]), (5, vec![4, 5, 6])]
+pub struct DummyDataProvider;
+impl BeefyDataProvider<Vec<u8>> for DummyDataProvider {
+	fn extra_data() -> Vec<u8> {
+		let mut col = vec![(15, vec![1, 2, 3]), (5, vec![4, 5, 6])];
+		col.sort();
+		beefy_merkle_tree::merkle_root::<crate::Pallet<Test>, _, _>(
+			col.into_iter().map(|pair| pair.encode()),
+		)
+		.to_vec()
 	}
 }
 

frame/beefy-mmr/src/tests.rs

@@ -71,7 +71,7 @@ fn should_contain_mmr_digest() {
 		assert_eq!(
 			System::digest().logs,
 			vec![beefy_log(ConsensusLog::MmrRoot(
-				hex!("969d516e5279540ef38e4a710fb0645cab4c3b01e528be7285b85ec9c5fb55c8").into()
+				hex!("fa0275b19b2565089f7e2377ee73b9050e8d53bce108ef722a3251fd9d371d4b").into()
 			))]
 		);
 
@@ -82,13 +82,13 @@ fn should_contain_mmr_digest() {
 			System::digest().logs,
 			vec![
 				beefy_log(ConsensusLog::MmrRoot(
-					hex!("969d516e5279540ef38e4a710fb0645cab4c3b01e528be7285b85ec9c5fb55c8").into()
+					hex!("fa0275b19b2565089f7e2377ee73b9050e8d53bce108ef722a3251fd9d371d4b").into()
 				)),
 				beefy_log(ConsensusLog::AuthoritiesChange(
 					ValidatorSet::new(vec![mock_beefy_id(3), mock_beefy_id(4),], 1,).unwrap()
 				)),
 				beefy_log(ConsensusLog::MmrRoot(
-					hex!("8c42b7b040d262f7f2e26abeb61ab0c3c448f60c7f2f19e6ca0035d9bb3ae7e2").into()
+					hex!("85554fa7d4e863cce3cdce668c1ae82c0174ad37f8d1399284018bec9f9971c3").into()
 				)),
 			]
 		);
@@ -114,10 +114,8 @@ fn should_contain_valid_leaf_data() {
 				root: hex!("176e73f1bf656478b728e28dd1a7733c98621b8acf830bff585949763dca7a96")
 					.into(),
 			},
-			parachain_heads: hex!(
-				"ed893c8f8cc87195a5d4d2805b011506322036bcace79642aa3e94ab431e442e"
-			)
-			.into(),
+			leaf_extra: hex!("55b8e9e1cc9f0db7776fac0ca66318ef8acfb8ec26db11e373120583e07ee648")
+				.to_vec(),
 		}
 	);
 
@@ -138,10 +136,8 @@ fn should_contain_valid_leaf_data() {
 				root: hex!("9c6b2c1b0d0b25a008e6c882cc7b415f309965c72ad2b944ac0931048ca31cd5")
 					.into(),
 			},
-			parachain_heads: hex!(
-				"ed893c8f8cc87195a5d4d2805b011506322036bcace79642aa3e94ab431e442e"
-			)
-			.into(),
+			leaf_extra: hex!("55b8e9e1cc9f0db7776fac0ca66318ef8acfb8ec26db11e373120583e07ee648")
+				.to_vec()
 		}
 	);
 }

primitives/beefy/src/mmr.rs

@@ -26,12 +26,26 @@
 //! but we imagine they will be useful for other chains that either want to bridge with Polkadot
 //! or are completely standalone, but heavily inspired by Polkadot.
 
+use crate::Vec;
 use codec::{Decode, Encode, MaxEncodedLen};
 use scale_info::TypeInfo;
 
+/// A provider for extra data that gets added to the Mmr leaf
+pub trait BeefyDataProvider<ExtraData> {
+	/// Return a vector of bytes, ideally should be a merkle root hash
+	fn extra_data() -> ExtraData;
+}
+
+/// A default implementation for runtimes.
+impl BeefyDataProvider<Vec<u8>> for () {
+	fn extra_data() -> Vec<u8> {
+		Vec::new()
+	}
+}
+
 /// A standard leaf that gets added every block to the MMR constructed by Substrate's `pallet_mmr`.
 #[derive(Debug, PartialEq, Eq, Clone, Encode, Decode)]
-pub struct MmrLeaf<BlockNumber, Hash, MerkleRoot> {
+pub struct MmrLeaf<BlockNumber, Hash, MerkleRoot, ExtraData> {
 	/// Version of the leaf format.
 	///
 	/// Can be used to enable future format migrations and compatibility.
@@ -41,8 +55,9 @@ pub struct MmrLeaf<BlockNumber, Hash, MerkleRoot> {
 	pub parent_number_and_hash: (BlockNumber, Hash),
 	/// A merkle root of the next BEEFY authority set.
 	pub beefy_next_authority_set: BeefyNextAuthoritySet<MerkleRoot>,
-	/// A merkle root of all registered parachain heads.
-	pub parachain_heads: MerkleRoot,
+	/// Arbitrary extra leaf data to be used by downstream pallets to include custom data in the
+	/// [`MmrLeaf`]
+	pub leaf_extra: ExtraData,
 }
 
 /// A MMR leaf versioning scheme.