Support MMR Pruning (#9700)

* Use `0.3.2` * Replace `u64` with `NodeIndex` * Fix Typo * Add Pruning Logic * Fix Some Tests * Remove Comment * Log Only Under STD * Return while No Element to Append * Optimize Pruning Algorithm * Update Doc * Update Doc * Zero Copy Algorithm * Import Missing Type * Fix Merge Mistake * Import Missing Item * Make `verify` Off-Chain * `cargo fmt` * Avoid using NodeIndex in incorrect places. * Simplify pruning. * Format Co-authored-by: Tomasz Drwięga <tomasz@parity.io>
2026-05-30 00:01:03 +00:00 · 2021-11-23 21:04:05 +08:00
parent 76737891bb
commit f9122a870e
11 changed files with 165 additions and 61 deletions
@@ -14,7 +14,7 @@ targets = ["x86_64-unknown-linux-gnu"]
 [dependencies]
 codec = { package = "parity-scale-codec", version = "2.0.0", default-features = false }
 scale-info = { version = "1.0", default-features = false, features = ["derive"] }
-mmr-lib = { package = "ckb-merkle-mountain-range", default-features = false, version = "0.3.1" }
+mmr-lib = { package = "ckb-merkle-mountain-range", default-features = false, version = "0.3.2" }

 sp-core = { version = "4.0.0-dev", default-features = false, path = "../../primitives/core" }
 sp-io = { version = "4.0.0-dev", default-features = false, path = "../../primitives/io" }
@@ -26,6 +26,16 @@ use sp_std::fmt;
 #[cfg(not(feature = "std"))]
 use sp_std::prelude::Vec;

+/// A type to describe node position in the MMR (node index).
+pub type NodeIndex = u64;
+
+/// A type to describe leaf position in the MMR.
+///
+/// Note this is different from [`NodeIndex`], which can be applied to
+/// both leafs and inner nodes. Leafs will always have consecutive `LeafIndex`,
+/// but might be actually at different positions in the MMR `NodeIndex`.
+pub type LeafIndex = u64;
+
 /// A provider of the MMR's leaf data.
 pub trait LeafDataProvider {
 	/// A type that should end up in the leaf of MMR.
@@ -275,9 +285,9 @@ impl_leaf_data_for_tuple!(A:0, B:1, C:2, D:3, E:4);
 #[derive(codec::Encode, codec::Decode, RuntimeDebug, Clone, PartialEq, Eq)]
 pub struct Proof<Hash> {
 	/// The index of the leaf the proof is for.
-	pub leaf_index: u64,
+	pub leaf_index: LeafIndex,
 	/// Number of leaves in MMR, when the proof was generated.
-	pub leaf_count: u64,
+	pub leaf_count: NodeIndex,
 	/// Proof elements (hashes of siblings of inner nodes on the path to the leaf).
 	pub items: Vec<Hash>,
 }
@@ -402,7 +412,7 @@ sp_api::decl_runtime_apis! {
 	/// API to interact with MMR pallet.
 	pub trait MmrApi<Hash: codec::Codec> {
 		/// Generate MMR proof for a leaf under given index.
-		fn generate_proof(leaf_index: u64) -> Result<(EncodableOpaqueLeaf, Proof<Hash>), Error>;
+		fn generate_proof(leaf_index: LeafIndex) -> Result<(EncodableOpaqueLeaf, Proof<Hash>), Error>;

 		/// Verify MMR proof against on-chain MMR.
 		///
@@ -32,7 +32,7 @@ use sp_blockchain::HeaderBackend;
 use sp_core::Bytes;
 use sp_runtime::{generic::BlockId, traits::Block as BlockT};

-pub use pallet_mmr_primitives::MmrApi as MmrRuntimeApi;
+pub use pallet_mmr_primitives::{LeafIndex, MmrApi as MmrRuntimeApi};

 /// Retrieved MMR leaf and its proof.
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
@@ -71,7 +71,7 @@ pub trait MmrApi<BlockHash> {
 	#[rpc(name = "mmr_generateProof")]
 	fn generate_proof(
 		&self,
-		leaf_index: u64,
+		leaf_index: LeafIndex,
 		at: Option<BlockHash>,
 	) -> Result<LeafProof<BlockHash>>;
 }
@@ -98,7 +98,7 @@ where
 {
 	fn generate_proof(
 		&self,
-		leaf_index: u64,
+		leaf_index: LeafIndex,
 		at: Option<<Block as BlockT>::Hash>,
 	) -> Result<LeafProof<<Block as BlockT>::Hash>> {
 		let api = self.client.runtime_api();
@@ -25,7 +25,7 @@ benchmarks_instance_pallet! {
 	on_initialize {
 		let x in 1 .. 1_000;

-		let leaves = x as u64;
+		let leaves = x as NodeIndex;
 	}: {
 		for b in 0..leaves {
 			Pallet::<T, I>::on_initialize((b as u32).into());
@@ -70,10 +70,10 @@ mod mock;
 mod tests;

 pub use pallet::*;
-pub use pallet_mmr_primitives as primitives;
+pub use pallet_mmr_primitives::{self as primitives, NodeIndex};

 pub trait WeightInfo {
-	fn on_initialize(peaks: u64) -> Weight;
+	fn on_initialize(peaks: NodeIndex) -> Weight;
 }

 #[frame_support::pallet]
@@ -160,7 +160,7 @@ pub mod pallet {
 	/// Current size of the MMR (number of leaves).
 	#[pallet::storage]
 	#[pallet::getter(fn mmr_leaves)]
-	pub type NumberOfLeaves<T, I = ()> = StorageValue<_, u64, ValueQuery>;
+	pub type NumberOfLeaves<T, I = ()> = StorageValue<_, NodeIndex, ValueQuery>;

 	/// Hashes of the nodes in the MMR.
 	///
@@ -169,7 +169,7 @@ pub mod pallet {
 	#[pallet::storage]
 	#[pallet::getter(fn mmr_peak)]
 	pub type Nodes<T: Config<I>, I: 'static = ()> =
-		StorageMap<_, Identity, u64, <T as Config<I>>::Hash, OptionQuery>;
+		StorageMap<_, Identity, NodeIndex, <T as Config<I>>::Hash, OptionQuery>;

 	#[pallet::hooks]
 	impl<T: Config<I>, I: 'static> Hooks<BlockNumberFor<T>> for Pallet<T, I> {
@@ -228,7 +228,7 @@ where
 }

 impl<T: Config<I>, I: 'static> Pallet<T, I> {
-	fn offchain_key(pos: u64) -> sp_std::prelude::Vec<u8> {
+	fn offchain_key(pos: NodeIndex) -> sp_std::prelude::Vec<u8> {
 		(T::INDEXING_PREFIX, pos).encode()
 	}

@@ -239,7 +239,7 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 	/// all the leaves to be present.
 	/// It may return an error or panic if used incorrectly.
 	pub fn generate_proof(
-		leaf_index: u64,
+		leaf_index: NodeIndex,
 	) -> Result<(LeafOf<T, I>, primitives::Proof<<T as Config<I>>::Hash>), primitives::Error> {
 		let mmr: ModuleMmr<mmr::storage::OffchainStorage, T, I> = mmr::Mmr::new(Self::mmr_leaves());
 		mmr.generate_proof(leaf_index)
@@ -263,7 +263,7 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 				.log_debug("The proof has incorrect number of leaves or proof items."))
 		}

-		let mmr: ModuleMmr<mmr::storage::RuntimeStorage, T, I> = mmr::Mmr::new(proof.leaf_count);
+		let mmr: ModuleMmr<mmr::storage::OffchainStorage, T, I> = mmr::Mmr::new(proof.leaf_count);
 		let is_valid = mmr.verify_leaf_proof(leaf, proof)?;
 		if is_valid {
 			Ok(())
@@ -21,7 +21,7 @@ use crate::{
 		utils::NodesUtils,
 		Hasher, Node, NodeOf,
 	},
-	primitives::{self, Error},
+	primitives::{self, Error, NodeIndex},
 	Config, HashingOf,
 };
 #[cfg(not(feature = "std"))]
@@ -60,7 +60,7 @@ where
 	Storage<StorageType, T, I, L>: mmr_lib::MMRStore<NodeOf<T, I, L>>,
 {
 	mmr: mmr_lib::MMR<NodeOf<T, I, L>, Hasher<HashingOf<T, I>, L>, Storage<StorageType, T, I, L>>,
-	leaves: u64,
+	leaves: NodeIndex,
 }

 impl<StorageType, T, I, L> Mmr<StorageType, T, I, L>
@@ -71,7 +71,7 @@ where
 	Storage<StorageType, T, I, L>: mmr_lib::MMRStore<NodeOf<T, I, L>>,
 {
 	/// Create a pointer to an existing MMR with given number of leaves.
-	pub fn new(leaves: u64) -> Self {
+	pub fn new(leaves: NodeIndex) -> Self {
 		let size = NodesUtils::new(leaves).size();
 		Self { mmr: mmr_lib::MMR::new(size, Default::default()), leaves }
 	}
@@ -94,7 +94,7 @@ where

 	/// Return the internal size of the MMR (number of nodes).
 	#[cfg(test)]
-	pub fn size(&self) -> u64 {
+	pub fn size(&self) -> NodeIndex {
 		self.mmr.mmr_size()
 	}
 }
@@ -109,7 +109,7 @@ where
 	/// Push another item to the MMR.
 	///
 	/// Returns element position (index) in the MMR.
-	pub fn push(&mut self, leaf: L) -> Option<u64> {
+	pub fn push(&mut self, leaf: L) -> Option<NodeIndex> {
 		let position =
 			self.mmr.push(Node::Data(leaf)).map_err(|e| Error::Push.log_error(e)).ok()?;

@@ -120,7 +120,7 @@ where

 	/// Commit the changes to underlying storage, return current number of leaves and
 	/// calculate the new MMR's root hash.
-	pub fn finalize(self) -> Result<(u64, <T as Config<I>>::Hash), Error> {
+	pub fn finalize(self) -> Result<(NodeIndex, <T as Config<I>>::Hash), Error> {
 		let root = self.mmr.get_root().map_err(|e| Error::GetRoot.log_error(e))?;
 		self.mmr.commit().map_err(|e| Error::Commit.log_error(e))?;
 		Ok((self.leaves, root.hash()))
@@ -140,7 +140,7 @@ where
 	/// (i.e. you can't run the function in the pruned storage).
 	pub fn generate_proof(
 		&self,
-		leaf_index: u64,
+		leaf_index: NodeIndex,
 	) -> Result<(L, primitives::Proof<<T as Config<I>>::Hash>), Error> {
 		let position = mmr_lib::leaf_index_to_pos(leaf_index);
 		let store = <Storage<OffchainStorage, T, I, L>>::default();
@@ -18,19 +18,24 @@
 //! A MMR storage implementations.

 use codec::Encode;
+use frame_support::log;
+use mmr_lib::helper;
+use sp_io::offchain_index;
+use sp_std::iter::Peekable;
 #[cfg(not(feature = "std"))]
-use sp_std::prelude::Vec;
+use sp_std::prelude::*;

 use crate::{
-	mmr::{Node, NodeOf},
-	primitives, Config, Nodes, NumberOfLeaves, Pallet,
+	mmr::{utils::NodesUtils, Node, NodeOf},
+	primitives::{self, NodeIndex},
+	Config, Nodes, NumberOfLeaves, Pallet,
 };

 /// A marker type for runtime-specific storage implementation.
 ///
 /// Allows appending new items to the MMR and proof verification.
 /// MMR nodes are appended to two different storages:
-/// 1. We add nodes (leaves) hashes to the on-chain storge (see [crate::Nodes]).
+/// 1. We add nodes (leaves) hashes to the on-chain storage (see [crate::Nodes]).
 /// 2. We add full leaves (and all inner nodes as well) into the `IndexingAPI` during block
 ///    processing, so the values end up in the Offchain DB if indexing is enabled.
 pub struct RuntimeStorage;
@@ -60,14 +65,14 @@ where
 	I: 'static,
 	L: primitives::FullLeaf + codec::Decode,
 {
-	fn get_elem(&self, pos: u64) -> mmr_lib::Result<Option<NodeOf<T, I, L>>> {
+	fn get_elem(&self, pos: NodeIndex) -> mmr_lib::Result<Option<NodeOf<T, I, L>>> {
 		let key = Pallet::<T, I>::offchain_key(pos);
 		// Retrieve the element from Off-chain DB.
 		Ok(sp_io::offchain::local_storage_get(sp_core::offchain::StorageKind::PERSISTENT, &key)
 			.and_then(|v| codec::Decode::decode(&mut &*v).ok()))
 	}

-	fn append(&mut self, _: u64, _: Vec<NodeOf<T, I, L>>) -> mmr_lib::Result<()> {
+	fn append(&mut self, _: NodeIndex, _: Vec<NodeOf<T, I, L>>) -> mmr_lib::Result<()> {
 		panic!("MMR must not be altered in the off-chain context.")
 	}
 }
@@ -78,32 +83,90 @@ where
 	I: 'static,
 	L: primitives::FullLeaf,
 {
-	fn get_elem(&self, pos: u64) -> mmr_lib::Result<Option<NodeOf<T, I, L>>> {
+	fn get_elem(&self, pos: NodeIndex) -> mmr_lib::Result<Option<NodeOf<T, I, L>>> {
 		Ok(<Nodes<T, I>>::get(pos).map(Node::Hash))
 	}

-	fn append(&mut self, pos: u64, elems: Vec<NodeOf<T, I, L>>) -> mmr_lib::Result<()> {
-		let mut leaves = crate::NumberOfLeaves::<T, I>::get();
-		let mut size = crate::mmr::utils::NodesUtils::new(leaves).size();
+	fn append(&mut self, pos: NodeIndex, elems: Vec<NodeOf<T, I, L>>) -> mmr_lib::Result<()> {
+		if elems.is_empty() {
+			return Ok(())
+		}
+
+		sp_std::if_std! {
+			log::trace!("elems: {:?}", elems.iter().map(|elem| elem.hash()).collect::<Vec<_>>());
+		}
+
+		let leaves = NumberOfLeaves::<T, I>::get();
+		let size = NodesUtils::new(leaves).size();
+
 		if pos != size {
 			return Err(mmr_lib::Error::InconsistentStore)
 		}

-		for elem in elems {
-			// on-chain we only store the hash (even if it's a leaf)
-			<Nodes<T, I>>::insert(size, elem.hash());
-			// Indexing API is used to store the full leaf content.
-			let key = Pallet::<T, I>::offchain_key(size);
-			elem.using_encoded(|elem| sp_io::offchain_index::set(&key, elem));
-			size += 1;
+		let new_size = size + elems.len() as NodeIndex;

-			if let Node::Data(..) = elem {
-				leaves += 1;
+		// A sorted (ascending) iterator over peak indices to prune and persist.
+		let (peaks_to_prune, mut peaks_to_store) = peaks_to_prune_and_store(size, new_size);
+
+		// Now we are going to iterate over elements to insert
+		// and keep track of the current `node_index` and `leaf_index`.
+		let mut leaf_index = leaves;
+		let mut node_index = size;
+
+		for elem in elems {
+			// Indexing API is used to store the full node content (both leaf and inner).
+			elem.using_encoded(|elem| {
+				offchain_index::set(&Pallet::<T, I>::offchain_key(node_index), elem)
+			});
+
+			// On-chain we are going to only store new peaks.
+			if peaks_to_store.next_if_eq(&node_index).is_some() {
+				<Nodes<T, I>>::insert(node_index, elem.hash());
 			}
+
+			// Increase the indices.
+			if let Node::Data(..) = elem {
+				leaf_index += 1;
+			}
+			node_index += 1;
 		}

-		NumberOfLeaves::<T, I>::put(leaves);
+		// Update current number of leaves.
+		NumberOfLeaves::<T, I>::put(leaf_index);
+
+		// And remove all remaining items from `peaks_before` collection.
+		for pos in peaks_to_prune {
+			<Nodes<T, I>>::remove(pos);
+		}

 		Ok(())
 	}
 }
+
+fn peaks_to_prune_and_store(
+	old_size: NodeIndex,
+	new_size: NodeIndex,
+) -> (impl Iterator<Item = NodeIndex>, Peekable<impl Iterator<Item = NodeIndex>>) {
+	// A sorted (ascending) collection of peak indices before and after insertion.
+	// both collections may share a common prefix.
+	let peaks_before = if old_size == 0 { vec![] } else { helper::get_peaks(old_size) };
+	let peaks_after = helper::get_peaks(new_size);
+	sp_std::if_std! {
+		log::trace!("peaks_before: {:?}", peaks_before);
+		log::trace!("peaks_after: {:?}", peaks_after);
+	}
+	let mut peaks_before = peaks_before.into_iter().peekable();
+	let mut peaks_after = peaks_after.into_iter().peekable();
+
+	// Consume a common prefix between `peaks_before` and `peaks_after`,
+	// since that's something we will not be touching anyway.
+	while peaks_before.peek() == peaks_after.peek() {
+		peaks_before.next();
+		peaks_after.next();
+	}
+
+	// what's left in both collections is:
+	// 1. Old peaks to remove from storage
+	// 2. New peaks to persist in storage
+	(peaks_before, peaks_after)
+}
@@ -17,29 +17,31 @@

 //! Merkle Mountain Range utilities.

+use crate::primitives::{LeafIndex, NodeIndex};
+
 /// MMR nodes & size -related utilities.
 pub struct NodesUtils {
-	no_of_leaves: u64,
+	no_of_leaves: LeafIndex,
 }

 impl NodesUtils {
 	/// Create new instance of MMR nodes utilities for given number of leaves.
-	pub fn new(no_of_leaves: u64) -> Self {
+	pub fn new(no_of_leaves: LeafIndex) -> Self {
 		Self { no_of_leaves }
 	}

 	/// Calculate number of peaks in the MMR.
-	pub fn number_of_peaks(&self) -> u64 {
-		self.number_of_leaves().count_ones() as u64
+	pub fn number_of_peaks(&self) -> NodeIndex {
+		self.number_of_leaves().count_ones() as NodeIndex
 	}

 	/// Return the number of leaves in the MMR.
-	pub fn number_of_leaves(&self) -> u64 {
+	pub fn number_of_leaves(&self) -> LeafIndex {
 		self.no_of_leaves
 	}

 	/// Calculate the total size of MMR (number of nodes).
-	pub fn size(&self) -> u64 {
+	pub fn size(&self) -> NodeIndex {
 		2 * self.no_of_leaves - self.number_of_peaks()
 	}

@@ -15,9 +15,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.

-use crate::{mock::*, *};
+use crate::{mmr::utils, mock::*, *};

 use frame_support::traits::OnInitialize;
+use mmr_lib::helper;
 use pallet_mmr_primitives::{Compact, Proof};
 use sp_core::{
 	offchain::{testing::TestOffchainExt, OffchainDbExt, OffchainWorkerExt},
@@ -48,6 +49,12 @@ fn new_block() -> u64 {
 	MMR::on_initialize(number)
 }

+fn peaks_from_leaves_count(leaves_count: NodeIndex) -> Vec<NodeIndex> {
+	let size = utils::NodesUtils::new(leaves_count).size();
+
+	helper::get_peaks(size)
+}
+
 pub(crate) fn hex(s: &str) -> H256 {
 	s.parse().unwrap()
 }
@@ -115,10 +122,29 @@ fn should_append_to_mmr_when_on_initialize_is_called() {
 	ext.execute_with(|| {
 		// when
 		new_block();
+
+		// then
+		assert_eq!(crate::NumberOfLeaves::<Test>::get(), 1);
+		assert_eq!(
+			(
+				crate::Nodes::<Test>::get(0),
+				crate::Nodes::<Test>::get(1),
+				crate::RootHash::<Test>::get(),
+			),
+			(
+				Some(hex("4320435e8c3318562dba60116bdbcc0b82ffcecb9bb39aae3300cfda3ad0b8b0")),
+				None,
+				hex("0x4320435e8c3318562dba60116bdbcc0b82ffcecb9bb39aae3300cfda3ad0b8b0"),
+			)
+		);
+
+		// when
 		new_block();

 		// then
 		assert_eq!(crate::NumberOfLeaves::<Test>::get(), 2);
+		let peaks = peaks_from_leaves_count(2);
+		assert_eq!(peaks, vec![2]);
 		assert_eq!(
 			(
 				crate::Nodes::<Test>::get(0),
@@ -128,8 +154,8 @@ fn should_append_to_mmr_when_on_initialize_is_called() {
 				crate::RootHash::<Test>::get(),
 			),
 			(
-				Some(hex("4320435e8c3318562dba60116bdbcc0b82ffcecb9bb39aae3300cfda3ad0b8b0")),
-				Some(hex("ad4cbc033833612ccd4626d5f023b9dfc50a35e838514dd1f3c86f8506728705")),
+				None,
+				None,
 				Some(hex("672c04a9cd05a644789d769daa552d35d8de7c33129f8a7cbf49e595234c4854")),
 				None,
 				hex("672c04a9cd05a644789d769daa552d35d8de7c33129f8a7cbf49e595234c4854"),
@@ -166,14 +192,21 @@ fn should_construct_larger_mmr_correctly() {

 		// then
 		assert_eq!(crate::NumberOfLeaves::<Test>::get(), 7);
+		let peaks = peaks_from_leaves_count(7);
+		assert_eq!(peaks, vec![6, 9, 10]);
+		for i in (0..=10).filter(|p| !peaks.contains(p)) {
+			assert!(crate::Nodes::<Test>::get(i).is_none());
+		}
 		assert_eq!(
 			(
-				crate::Nodes::<Test>::get(0),
+				crate::Nodes::<Test>::get(6),
+				crate::Nodes::<Test>::get(9),
 				crate::Nodes::<Test>::get(10),
 				crate::RootHash::<Test>::get(),
 			),
 			(
-				Some(hex("4320435e8c3318562dba60116bdbcc0b82ffcecb9bb39aae3300cfda3ad0b8b0")),
+				Some(hex("ae88a0825da50e953e7a359c55fe13c8015e48d03d301b8bdfc9193874da9252")),
+				Some(hex("7e4316ae2ebf7c3b6821cb3a46ca8b7a4f9351a9b40fcf014bb0a4fd8e8f29da")),
 				Some(hex("611c2174c6164952a66d985cfe1ec1a623794393e3acff96b136d198f37a648c")),
 				hex("e45e25259f7930626431347fa4dd9aae7ac83b4966126d425ca70ab343709d2c"),
 			)
@@ -265,11 +298,7 @@ fn should_verify() {
 		crate::Pallet::<Test>::generate_proof(5).unwrap()
 	});

-	// Now to verify the proof, we really shouldn't require offchain storage or extension.
-	// Hence we initialize the storage once again, using different externalities and then
-	// verify.
-	let mut ext2 = new_test_ext();
-	ext2.execute_with(|| {
+	ext.execute_with(|| {
 		init_chain(7);
 		// then
 		assert_eq!(crate::Pallet::<Test>::verify_leaf(leaf, proof5), Ok(()));