Support MMR Pruning

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,9 @@ use sp_std::fmt;
 #[cfg(not(feature = "std"))]
 use sp_std::prelude::Vec;
 
+/// Node index type.
+pub type NodeIndex = 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 +278,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: NodeIndex,
 	/// 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 +405,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: NodeIndex) -> Result<(EncodableOpaqueLeaf, Proof<Hash>), Error>;
 
 		/// Verify MMR proof against on-chain MMR.
 		///

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

@@ -71,7 +71,7 @@ pub trait MmrApi<BlockHash> {
 	#[rpc(name = "mmr_generateProof")]
 	fn generate_proof(
 		&self,
-		leaf_index: u64,
+		leaf_index: NodeIndex,
 		at: Option<BlockHash>,
 	) -> Result<LeafProof<BlockHash>>;
 }
@@ -98,7 +98,7 @@ where
 {
 	fn generate_proof(
 		&self,
-		leaf_index: u64,
+		leaf_index: NodeIndex,
 		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)

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,23 @@
 //! A MMR storage implementations.
 
 use codec::Encode;
+use frame_support::log;
+use mmr_lib::helper;
+use sp_io::offchain_index;
 #[cfg(not(feature = "std"))]
 use sp_std::prelude::Vec;
 
 use crate::{
-	mmr::{Node, NodeOf},
-	primitives, Config, Nodes, NumberOfLeaves, Pallet,
+	mmr::{utils::NodesUtils, Node, NodeOf},
+	primitives::{self, DataOrHash, 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 +64,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 +82,136 @@ 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());
+		let elems = elems
+			.into_iter()
+			.enumerate()
+			.map(|(i, elem)| (size + i as NodeIndex, elem))
+			.collect::<Vec<_>>();
+		let leaves = elems.iter().fold(leaves, |acc, (pos, elem)| {
 			// 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;
+			elem.using_encoded(|elem| {
+				offchain_index::set(&Pallet::<T, I>::offchain_key(*pos), elem)
+			});
 
 			if let Node::Data(..) = elem {
-				leaves += 1;
+				acc + 1
+			} else {
+				acc
 			}
-		}
+		});
 
 		NumberOfLeaves::<T, I>::put(leaves);
 
+		let peaks_before = if size == 0 { vec![] } else { helper::get_peaks(size) };
+		let peaks_after = helper::get_peaks(size + elems.len() as NodeIndex);
+
+		sp_std::if_std! {
+			log::trace!("peaks_before: {:?}", peaks_before);
+			log::trace!("peaks_after: {:?}", peaks_after);
+		}
+
+		let store = |peaks_to_store: &[_], elems: Vec<(_, DataOrHash<_, _>)>| {
+			let mut peak_to_store =
+				peaks_to_store.get(0).expect("`peaks_to_store` can not be empty; qed");
+			let mut i = 1;
+
+			for (pos, elem) in elems {
+				if &pos == peak_to_store {
+					i += 1;
+
+					<Nodes<T, I>>::insert(peak_to_store, elem.hash());
+
+					if let Some(next_peak_to_store) = peaks_to_store.get(i) {
+						peak_to_store = next_peak_to_store;
+					} else {
+						// No more peaks to store.
+						break
+					}
+				}
+			}
+		};
+
+		// A new tree to build, no need to prune.
+		if peaks_before.is_empty() {
+			store(&peaks_after, elems);
+
+			return Ok(())
+		}
+
+		let mut pivot = None;
+
+		// `peaks_before` and `peaks_after` have a common prefix.
+		for i in 0.. {
+			if let Some(peak_before) = peaks_before.get(i) {
+				if let Some(peak_after) = peaks_after.get(i) {
+					if peak_before == peak_after {
+						pivot = Some(i);
+					} else {
+						break
+					}
+				} else {
+					break
+				}
+			} else {
+				break
+			}
+		}
+
+		sp_std::if_std! {
+			log::trace!("pivot: {:?}", pivot);
+		}
+
+		// According to the MMR specification
+		// `nodes_to_prune` might be empty
+		// `peaks_to_store` must not be empty
+		let mut nodes_to_prune = &[][..];
+		let mut peaks_to_store = &[][..];
+
+		if let Some(pivot) = pivot {
+			let pivot = pivot + 1;
+
+			if pivot < peaks_before.len() {
+				nodes_to_prune = &peaks_before[pivot..];
+			}
+			if pivot < peaks_after.len() {
+				peaks_to_store = &peaks_after[pivot..];
+			}
+		} else {
+			nodes_to_prune = &peaks_before;
+			peaks_to_store = &peaks_after;
+		};
+
+		sp_std::if_std! {
+			log::trace!("nodes_to_prune: {:?}", nodes_to_prune.to_vec());
+			log::trace!("peaks_to_store: {:?}", peaks_to_store.to_vec());
+		}
+
+		store(peaks_to_store, elems);
+
+		for pos in nodes_to_prune {
+			<Nodes<T, I>>::remove(pos);
+		}
+
 		Ok(())
 	}
 }