Support MMR Pruning (paritytech#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 <[email protected]>

bin/node/runtime/src/lib.rs

@@ -1559,7 +1559,7 @@ impl_runtime_apis! {
 		Block,
 		mmr::Hash,
 	> for Runtime {
-		fn generate_proof(leaf_index: u64)
+		fn generate_proof(leaf_index: pallet_mmr::primitives::LeafIndex)
 			-> Result<(mmr::EncodableOpaqueLeaf, mmr::Proof<mmr::Hash>), mmr::Error>
 		{
 			Mmr::generate_proof(leaf_index)

frame/merkle-mountain-range/Cargo.toml

@@ -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" }

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

@@ -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.
 		///

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

@@ -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();

frame/merkle-mountain-range/src/benchmarking.rs

@@ -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());

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

@@ -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(())

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

@@ -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();

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

@@ -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)
 		}
 
+		let new_size = size + elems.len() as NodeIndex;
+
+		// 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 {
-			// 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;
+			// 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 {
-				leaves += 1;
+				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)
+}