pallet-mmr: handle forks without collisions in offchain storage (#11594)

* pallet-mmr: fix some typos * pallet-mmr: make the MMR resilient to chain forks * pallet-mmr: get hash for block that added node * beefy-mmr: add debug logging * add explanatory comment * account for block offset of pallet activation * add support for finding all nodes added by leaf * minor improvements * add helper to return all nodes added to mmr with a leaf append * simplify leaf_node_index_to_leaf_index summing the (shifted) differences in peak positions adds up to the (shifted) final position, so don't need to fold over positions. * dead fish: this also doesn't work The idea was to keep a rolling window of `(parent_hash, pos)` leaf entries in the offchain db, with the window matching the one that provides `block_num -> block_hash` mappings in `frame_system`. Once a leaf exits the window it would be "canonicalized" by switching its offchain db key from `(parent_hash, pos)` to simple `pos`. This doesn't work however because there's no way to get leaf contents from offchain db while in runtime context.. so no way to get+clear+set leaf to change its key in offchain db. Ideas: 1. move the "canonicalization" logic to offchain worker 2. enhance IndexingApi with "offchain::move(old_key, new_key)" This is weird, but correct, deterministic and safe AFAICT, so it could be exposed to runtime. * simplify rightmost_leaf_node_index_from_pos * minor fix * move leaf canonicalization to offchain worker * move storage related code to storage.rs * on offchain reads use canonic key for old leaves * fix offchain worker write using canon key * fix pallet-mmr tests * add documentation and fix logging * add offchain mmr canonicalization test * test canon + generate + verify * fix pallet-beefy-mmr tests * implement review suggestions * improve test * pallet-mmr: add offchain pruning of forks * pallet-mmr: improve offchain pruning Instead of keeping pruning map as single blob in offchain db, keep individual parent-hash lists with block-num identifier as part of the offchain key. Signed-off-by: acatangiu <adrian@parity.io> * pallet-mmr: improve MMRStore<OffchainStorage>::get() Do the math and retrieve node using correct (canon or non-canon) offchain db key, instead of blindly looking in both canon and non-canon offchain db locations for each node. Still fallback on looking at both if for any reason it's not where expected. Signed-off-by: acatangiu <adrian@parity.io> * pallet-mmr: storage: improve logs * fix tests: correctly persist overlay runtime indexing API works on overlay, whereas offchain context bypasses overlay, so for loops > canon-window, canon would fail. * pallet-mmr: fix numeric typo in test * add comment around LeafData requirements Signed-off-by: acatangiu <adrian@parity.io> Co-authored-by: Robert Hambrock <roberthambrock@gmail.com>
2026-06-12 15:51:12 +00:00 · 2022-07-07 18:02:37 +03:00
parent 64ebc12a3f
commit 27b538222e
6 changed files with 635 additions and 53 deletions
@@ -44,16 +44,12 @@ pub fn beefy_log(log: ConsensusLog<BeefyId>) -> DigestItem {
 	DigestItem::Consensus(BEEFY_ENGINE_ID, log.encode())
 }

-fn offchain_key(pos: usize) -> Vec<u8> {
-	(<Test as pallet_mmr::Config>::INDEXING_PREFIX, pos as u64).encode()
-}
-
-fn read_mmr_leaf(ext: &mut TestExternalities, index: usize) -> MmrLeaf {
+fn read_mmr_leaf(ext: &mut TestExternalities, key: Vec<u8>) -> MmrLeaf {
 	type Node = pallet_mmr::primitives::DataOrHash<Keccak256, MmrLeaf>;
 	ext.persist_offchain_overlay();
 	let offchain_db = ext.offchain_db();
 	offchain_db
-		.get(&offchain_key(index))
+		.get(&key)
 		.map(|d| Node::decode(&mut &*d).unwrap())
 		.map(|n| match n {
 			Node::Data(d) => d,
@@ -105,12 +101,17 @@ fn should_contain_mmr_digest() {

 #[test]
 fn should_contain_valid_leaf_data() {
+	fn node_offchain_key(parent_hash: H256, pos: usize) -> Vec<u8> {
+		(<Test as pallet_mmr::Config>::INDEXING_PREFIX, parent_hash, pos as u64).encode()
+	}
+
 	let mut ext = new_test_ext(vec![1, 2, 3, 4]);
-	ext.execute_with(|| {
+	let parent_hash = ext.execute_with(|| {
 		init_block(1);
+		<frame_system::Pallet<Test>>::parent_hash()
 	});

-	let mmr_leaf = read_mmr_leaf(&mut ext, 0);
+	let mmr_leaf = read_mmr_leaf(&mut ext, node_offchain_key(parent_hash, 0));
 	assert_eq!(
 		mmr_leaf,
 		MmrLeaf {
@@ -128,11 +129,12 @@ fn should_contain_valid_leaf_data() {
 	);

 	// build second block on top
-	ext.execute_with(|| {
+	let parent_hash = ext.execute_with(|| {
 		init_block(2);
+		<frame_system::Pallet<Test>>::parent_hash()
 	});

-	let mmr_leaf = read_mmr_leaf(&mut ext, 1);
+	let mmr_leaf = read_mmr_leaf(&mut ext, node_offchain_key(parent_hash, 1));
 	assert_eq!(
 		mmr_leaf,
 		MmrLeaf {
@@ -58,7 +58,10 @@

 use codec::Encode;
 use frame_support::weights::Weight;
-use sp_runtime::traits::{self, One, Saturating};
+use sp_runtime::{
+	traits::{self, One, Saturating},
+	SaturatedConversion,
+};

 #[cfg(any(feature = "runtime-benchmarks", test))]
 mod benchmarking;
@@ -116,12 +119,12 @@ pub mod pallet {
 		/// Prefix for elements stored in the Off-chain DB via Indexing API.
 		///
 		/// Each node of the MMR is inserted both on-chain and off-chain via Indexing API.
-		/// The former does not store full leaf content, just it's compact version (hash),
+		/// The former does not store full leaf content, just its compact version (hash),
 		/// and some of the inner mmr nodes might be pruned from on-chain storage.
 		/// The latter will contain all the entries in their full form.
 		///
 		/// Each node is stored in the Off-chain DB under key derived from the
-		/// [`Self::INDEXING_PREFIX`] and it's in-tree index (MMR position).
+		/// [`Self::INDEXING_PREFIX`] and its in-tree index (MMR position).
 		const INDEXING_PREFIX: &'static [u8];

 		/// A hasher type for MMR.
@@ -162,6 +165,12 @@ pub mod pallet {
 		///
 		/// Note that the leaf at each block MUST be unique. You may want to include a block hash or
 		/// block number as an easiest way to ensure that.
+		/// Also note that the leaf added by each block is expected to only reference data coming
+		/// from ancestor blocks (leaves are saved offchain using `(parent_hash, pos)` key to be
+		/// fork-resistant, as such conflicts could only happen on 1-block deep forks, which means
+		/// two forks with identical line of ancestors compete to write the same offchain key, but
+		/// that's fine as long as leaves only contain data coming from ancestors - conflicting
+		/// writes are identical).
 		type LeafData: primitives::LeafDataProvider;

 		/// A hook to act on the new MMR root.
@@ -215,8 +224,31 @@ pub mod pallet {
 			<RootHash<T, I>>::put(root);

 			let peaks_after = mmr::utils::NodesUtils::new(leaves).number_of_peaks();
+
 			T::WeightInfo::on_initialize(peaks_before.max(peaks_after))
 		}
+
+		fn offchain_worker(n: T::BlockNumber) {
+			use mmr::storage::{OffchainStorage, Storage};
+			// MMR pallet uses offchain storage to hold full MMR and leaves.
+			// The leaves are saved under fork-unique keys `(parent_hash, pos)`.
+			// MMR Runtime depends on `frame_system::block_hash(block_num)` mappings to find
+			// parent hashes for particular nodes or leaves.
+			// This MMR offchain worker function moves a rolling window of the same size
+			// as `frame_system::block_hash` map, where nodes/leaves added by blocks that are just
+			// about to exit the window are "canonicalized" so that their offchain key no longer
+			// depends on `parent_hash` therefore on access to `frame_system::block_hash`.
+			//
+			// This approach works to eliminate fork-induced leaf collisions in offchain db,
+			// under the assumption that no fork will be deeper than `frame_system::BlockHashCount`
+			// blocks (2400 blocks on Polkadot, Kusama, Rococo, etc):
+			//   entries pertaining to block `N` where `N < current-2400` are moved to a key based
+			//   solely on block number. The only way to have collisions is if two competing forks
+			//   are deeper than 2400 blocks and they both "canonicalize" their view of block `N`.
+			// Once a block is canonicalized, all MMR entries pertaining to sibling blocks from
+			// other forks are pruned from offchain db.
+			Storage::<OffchainStorage, T, I, LeafOf<T, I>>::canonicalize_and_prune(n);
+		}
 	}
 }

@@ -254,9 +286,38 @@ where
 }

 impl<T: Config<I>, I: 'static> Pallet<T, I> {
-	fn offchain_key(pos: NodeIndex) -> sp_std::prelude::Vec<u8> {
+	/// Build offchain key from `parent_hash` of block that originally added node `pos` to MMR.
+	///
+	/// This combination makes the offchain (key,value) entry resilient to chain forks.
+	fn node_offchain_key(
+		parent_hash: <T as frame_system::Config>::Hash,
+		pos: NodeIndex,
+	) -> sp_std::prelude::Vec<u8> {
+		(T::INDEXING_PREFIX, parent_hash, pos).encode()
+	}
+
+	/// Build canonical offchain key for node `pos` in MMR.
+	///
+	/// Used for nodes added by now finalized blocks.
+	fn node_canon_offchain_key(pos: NodeIndex) -> sp_std::prelude::Vec<u8> {
 		(T::INDEXING_PREFIX, pos).encode()
 	}
+
+	/// Provide the parent number for the block that added `leaf_index` to the MMR.
+	fn leaf_index_to_parent_block_num(
+		leaf_index: LeafIndex,
+		leaves_count: LeafIndex,
+	) -> <T as frame_system::Config>::BlockNumber {
+		// leaves are zero-indexed and were added one per block since pallet activation,
+		// while block numbers are one-indexed, so block number that added `leaf_idx` is:
+		// `block_num = block_num_when_pallet_activated + leaf_idx + 1`
+		// `block_num = (current_block_num - leaves_count) + leaf_idx + 1`
+		// `parent_block_num = current_block_num - leaves_count + leaf_idx`.
+		<frame_system::Pallet<T>>::block_number()
+			.saturating_sub(leaves_count.saturated_into())
+			.saturating_add(leaf_index.saturated_into())
+	}
+
 	/// Generate a MMR proof for the given `leaf_indices`.
 	///
 	/// Note this method can only be used from an off-chain context
@@ -264,7 +325,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_batch_proof(
-		leaf_indices: Vec<NodeIndex>,
+		leaf_indices: Vec<LeafIndex>,
 	) -> Result<
 		(Vec<LeafOf<T, I>>, primitives::BatchProof<<T as Config<I>>::Hash>),
 		primitives::Error,
@@ -18,8 +18,11 @@
 //! A MMR storage implementations.

 use codec::Encode;
+use frame_support::traits::Get;
 use mmr_lib::helper;
-use sp_io::offchain_index;
+use sp_core::offchain::StorageKind;
+use sp_io::{offchain, offchain_index};
+use sp_runtime::traits::UniqueSaturatedInto;
 use sp_std::iter::Peekable;
 #[cfg(not(feature = "std"))]
 use sp_std::prelude::*;
@@ -46,6 +49,51 @@ pub struct RuntimeStorage;
 /// DOES NOT support adding new items to the MMR.
 pub struct OffchainStorage;

+/// Suffix of key for the 'pruning_map'.
+///
+/// Nodes and leaves are initially saved under fork-specific keys in offchain db,
+/// eventually they are "canonicalized" and this map is used to prune non-canon entries.
+const OFFCHAIN_PRUNING_MAP_KEY_SUFFIX: &str = "pruning_map";
+
+/// Used to store offchain mappings of `BlockNumber -> Vec[Hash]` to track all forks.
+/// Size of this offchain map is at most `frame_system::BlockHashCount`, its entries are pruned
+/// as part of the mechanism that prunes the forks this map tracks.
+pub(crate) struct PruningMap<T, I>(sp_std::marker::PhantomData<(T, I)>);
+impl<T, I> PruningMap<T, I>
+where
+	T: Config<I>,
+	I: 'static,
+{
+	pub(crate) fn pruning_map_offchain_key(block: T::BlockNumber) -> sp_std::prelude::Vec<u8> {
+		(T::INDEXING_PREFIX, block, OFFCHAIN_PRUNING_MAP_KEY_SUFFIX).encode()
+	}
+
+	/// Append `hash` to the list of parent hashes for `block` in offchain db.
+	pub fn append(block: T::BlockNumber, hash: <T as frame_system::Config>::Hash) {
+		let map_key = Self::pruning_map_offchain_key(block);
+		offchain::local_storage_get(StorageKind::PERSISTENT, &map_key)
+			.and_then(|v| codec::Decode::decode(&mut &*v).ok())
+			.or_else(|| Some(Vec::<<T as frame_system::Config>::Hash>::new()))
+			.map(|mut parents| {
+				parents.push(hash);
+				offchain::local_storage_set(
+					StorageKind::PERSISTENT,
+					&map_key,
+					&Encode::encode(&parents),
+				);
+			});
+	}
+
+	/// Remove list of parent hashes for `block` from offchain db and return it.
+	pub fn remove(block: T::BlockNumber) -> Option<Vec<<T as frame_system::Config>::Hash>> {
+		let map_key = Self::pruning_map_offchain_key(block);
+		offchain::local_storage_get(StorageKind::PERSISTENT, &map_key).and_then(|v| {
+			offchain::local_storage_clear(StorageKind::PERSISTENT, &map_key);
+			codec::Decode::decode(&mut &*v).ok()
+		})
+	}
+}
+
 /// A storage layer for MMR.
 ///
 /// There are two different implementations depending on the use case.
@@ -58,6 +106,109 @@ impl<StorageType, T, I, L> Default for Storage<StorageType, T, I, L> {
 	}
 }

+impl<T, I, L> Storage<OffchainStorage, T, I, L>
+where
+	T: Config<I>,
+	I: 'static,
+	L: primitives::FullLeaf,
+{
+	/// Move nodes and leaves added by block `N` in offchain db from _fork-aware key_ to
+	/// _canonical key_,
+	/// where `N` is `frame_system::BlockHashCount` blocks behind current block number.
+	///
+	/// This "canonicalization" process is required because the _fork-aware key_ value depends
+	/// on `frame_system::block_hash(block_num)` map which only holds the last
+	/// `frame_system::BlockHashCount` blocks.
+	///
+	/// For the canonicalized block, prune all nodes pertaining to other forks from offchain db.
+	///
+	/// Should only be called from offchain context, because it requires both read and write
+	/// access to offchain db.
+	pub(crate) fn canonicalize_and_prune(block: T::BlockNumber) {
+		// Add "block_num -> hash" mapping to offchain db,
+		// with all forks pushing hashes to same entry (same block number).
+		let parent_hash = <frame_system::Pallet<T>>::parent_hash();
+		PruningMap::<T, I>::append(block, parent_hash);
+
+		// Effectively move a rolling window of fork-unique leaves. Once out of the window, leaves
+		// are "canonicalized" in offchain by moving them under `Pallet::node_canon_offchain_key`.
+		let leaves = NumberOfLeaves::<T, I>::get();
+		let window_size =
+			<T as frame_system::Config>::BlockHashCount::get().unique_saturated_into();
+		if leaves >= window_size {
+			// Move the rolling window towards the end of  `block_num->hash` mappings available
+			// in the runtime: we "canonicalize" the leaf at the end,
+			let to_canon_leaf = leaves.saturating_sub(window_size);
+			// and all the nodes added by that leaf.
+			let to_canon_nodes = NodesUtils::right_branch_ending_in_leaf(to_canon_leaf);
+			frame_support::log::debug!(
+				target: "runtime::mmr::offchain", "Nodes to canon for leaf {}: {:?}",
+				to_canon_leaf, to_canon_nodes
+			);
+			// For this block number there may be node entries saved from multiple forks.
+			let to_canon_block_num =
+				Pallet::<T, I>::leaf_index_to_parent_block_num(to_canon_leaf, leaves);
+			// Only entries under this hash (retrieved from state on current canon fork) are to be
+			// persisted. All other entries added by same block number will be cleared.
+			let to_canon_hash = <frame_system::Pallet<T>>::block_hash(to_canon_block_num);
+
+			Self::canonicalize_nodes_for_hash(&to_canon_nodes, to_canon_hash);
+			// Get all the forks to prune, also remove them from the offchain pruning_map.
+			PruningMap::<T, I>::remove(to_canon_block_num)
+				.map(|forks| {
+					Self::prune_nodes_for_forks(&to_canon_nodes, forks);
+				})
+				.unwrap_or_else(|| {
+					frame_support::log::error!(
+						target: "runtime::mmr::offchain",
+						"Offchain: could not prune: no entry in pruning map for block {:?}",
+						to_canon_block_num
+					);
+				})
+		}
+	}
+
+	fn prune_nodes_for_forks(nodes: &[NodeIndex], forks: Vec<<T as frame_system::Config>::Hash>) {
+		for hash in forks {
+			for pos in nodes {
+				let key = Pallet::<T, I>::node_offchain_key(hash, *pos);
+				frame_support::log::debug!(
+					target: "runtime::mmr::offchain",
+					"Clear elem at pos {} with key {:?}",
+					pos, key
+				);
+				offchain::local_storage_clear(StorageKind::PERSISTENT, &key);
+			}
+		}
+	}
+
+	fn canonicalize_nodes_for_hash(
+		to_canon_nodes: &[NodeIndex],
+		to_canon_hash: <T as frame_system::Config>::Hash,
+	) {
+		for pos in to_canon_nodes {
+			let key = Pallet::<T, I>::node_offchain_key(to_canon_hash, *pos);
+			// Retrieve the element from Off-chain DB under fork-aware key.
+			if let Some(elem) = offchain::local_storage_get(StorageKind::PERSISTENT, &key) {
+				let canon_key = Pallet::<T, I>::node_canon_offchain_key(*pos);
+				// Add under new canon key.
+				offchain::local_storage_set(StorageKind::PERSISTENT, &canon_key, &elem);
+				frame_support::log::debug!(
+					target: "runtime::mmr::offchain",
+					"Moved elem at pos {} from key {:?} to canon key {:?}",
+					pos, key, canon_key
+				);
+			} else {
+				frame_support::log::error!(
+					target: "runtime::mmr::offchain",
+					"Could not canonicalize elem at pos {} using key {:?}",
+					pos, key
+				);
+			}
+		}
+	}
+}
+
 impl<T, I, L> mmr_lib::MMRStore<NodeOf<T, I, L>> for Storage<OffchainStorage, T, I, L>
 where
 	T: Config<I>,
@@ -65,9 +216,49 @@ where
 	L: primitives::FullLeaf + codec::Decode,
 {
 	fn get_elem(&self, pos: NodeIndex) -> mmr_lib::Result<Option<NodeOf<T, I, L>>> {
-		let key = Pallet::<T, I>::offchain_key(pos);
+		let leaves = NumberOfLeaves::<T, I>::get();
+		// Find out which leaf added node `pos` in the MMR.
+		let ancestor_leaf_idx = NodesUtils::leaf_index_that_added_node(pos);
+
+		let window_size =
+			<T as frame_system::Config>::BlockHashCount::get().unique_saturated_into();
+		// Leaves older than this window should have been canonicalized.
+		if leaves.saturating_sub(ancestor_leaf_idx) > window_size {
+			let key = Pallet::<T, I>::node_canon_offchain_key(pos);
+			frame_support::log::debug!(
+				target: "runtime::mmr::offchain", "offchain db get {}: leaf idx {:?}, key {:?}",
+				pos, ancestor_leaf_idx, key
+			);
+			// Just for safety, to easily handle runtime upgrades where any of the window params
+			// change and maybe we mess up storage migration,
+			// return _if and only if_ node is found (in normal conditions it's always found),
+			if let Some(elem) =
+				sp_io::offchain::local_storage_get(sp_core::offchain::StorageKind::PERSISTENT, &key)
+			{
+				return Ok(codec::Decode::decode(&mut &*elem).ok())
+			}
+			// BUT if we DID MESS UP, fall through to searching node using fork-specific key.
+		}
+
+		// Leaves still within the window will be found in offchain db under fork-aware keys.
+		let ancestor_parent_block_num =
+			Pallet::<T, I>::leaf_index_to_parent_block_num(ancestor_leaf_idx, leaves);
+		let ancestor_parent_hash = <frame_system::Pallet<T>>::block_hash(ancestor_parent_block_num);
+		let key = Pallet::<T, I>::node_offchain_key(ancestor_parent_hash, pos);
+		frame_support::log::debug!(
+			target: "runtime::mmr::offchain", "offchain db get {}: leaf idx {:?}, hash {:?}, key {:?}",
+			pos, ancestor_leaf_idx, ancestor_parent_hash, key
+		);
 		// Retrieve the element from Off-chain DB.
 		Ok(sp_io::offchain::local_storage_get(sp_core::offchain::StorageKind::PERSISTENT, &key)
+			.or_else(|| {
+				// Again, this is just us being extra paranoid.
+				// We get here only if we mess up a storage migration for a runtime upgrades where
+				// say the window is increased, and for a little while following the upgrade there's
+				// leaves inside new 'window' that had been already canonicalized before upgrade.
+				let key = Pallet::<T, I>::node_canon_offchain_key(pos);
+				sp_io::offchain::local_storage_get(sp_core::offchain::StorageKind::PERSISTENT, &key)
+			})
 			.and_then(|v| codec::Decode::decode(&mut &*v).ok()))
 	}

@@ -91,9 +282,11 @@ where
 			return Ok(())
 		}

-		sp_std::if_std! {
-			frame_support::log::trace!("elems: {:?}", elems.iter().map(|elem| elem.hash()).collect::<Vec<_>>());
-		}
+		frame_support::log::trace!(
+			target: "runtime::mmr",
+			"elems: {:?}",
+			elems.iter().map(|elem| elem.hash()).collect::<Vec<_>>()
+		);

 		let leaves = NumberOfLeaves::<T, I>::get();
 		let size = NodesUtils::new(leaves).size();
@@ -112,11 +305,24 @@ where
 		let mut leaf_index = leaves;
 		let mut node_index = size;

+		// Use parent hash of block adding new nodes (this block) as extra identifier
+		// in offchain DB to avoid DB collisions and overwrites in case of forks.
+		let parent_hash = <frame_system::Pallet<T>>::parent_hash();
 		for elem in elems {
+			// For now we store this leaf offchain keyed by `(parent_hash, node_index)`
+			// to make it fork-resistant.
+			// Offchain worker task will "canonicalize" it `frame_system::BlockHashCount` blocks
+			// later when we are not worried about forks anymore (highly unlikely to have a fork
+			// in the chain that deep).
+			// "Canonicalization" in this case means moving this leaf under a new key based
+			// only on the leaf's `node_index`.
+			let key = Pallet::<T, I>::node_offchain_key(parent_hash, node_index);
+			frame_support::log::debug!(
+				target: "runtime::mmr::offchain", "offchain db set: pos {} parent_hash {:?} key {:?}",
+				node_index, parent_hash, key
+			);
 			// 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)
-			});
+			elem.using_encoded(|elem| offchain_index::set(&key, elem));

 			// On-chain we are going to only store new peaks.
 			if peaks_to_store.next_if_eq(&node_index).is_some() {
@@ -150,10 +356,8 @@ fn peaks_to_prune_and_store(
 	// 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! {
-		frame_support::log::trace!("peaks_before: {:?}", peaks_before);
-		frame_support::log::trace!("peaks_after: {:?}", peaks_after);
-	}
+	frame_support::log::trace!(target: "runtime::mmr", "peaks_before: {:?}", peaks_before);
+	frame_support::log::trace!(target: "runtime::mmr", "peaks_after: {:?}", peaks_after);
 	let mut peaks_before = peaks_before.into_iter().peekable();
 	let mut peaks_after = peaks_after.into_iter().peekable();

@@ -18,6 +18,7 @@
 //! Merkle Mountain Range utilities.

 use crate::primitives::{LeafIndex, NodeIndex};
+use mmr_lib::helper;

 /// MMR nodes & size -related utilities.
 pub struct NodesUtils {
@@ -53,11 +54,78 @@ impl NodesUtils {

 		64 - self.no_of_leaves.next_power_of_two().leading_zeros()
 	}
+
+	/// Calculate `LeafIndex` for the leaf that added `node_index` to the MMR.
+	pub fn leaf_index_that_added_node(node_index: NodeIndex) -> LeafIndex {
+		let rightmost_leaf_pos = Self::rightmost_leaf_node_index_from_pos(node_index);
+		Self::leaf_node_index_to_leaf_index(rightmost_leaf_pos)
+	}
+
+	// Translate a _leaf_ `NodeIndex` to its `LeafIndex`.
+	fn leaf_node_index_to_leaf_index(pos: NodeIndex) -> LeafIndex {
+		if pos == 0 {
+			return 0
+		}
+		let peaks = helper::get_peaks(pos);
+		(pos + peaks.len() as u64) >> 1
+	}
+
+	// Starting from any node position get position of rightmost leaf; this is the leaf
+	// responsible for the addition of node `pos`.
+	fn rightmost_leaf_node_index_from_pos(pos: NodeIndex) -> NodeIndex {
+		pos - (helper::pos_height_in_tree(pos) as u64)
+	}
+
+	/// Starting from any leaf index, get the sequence of positions of the nodes added
+	/// to the mmr when this leaf was added (inclusive of the leaf's position itself).
+	/// That is, all of these nodes are right children of their respective parents.
+	pub fn right_branch_ending_in_leaf(leaf_index: LeafIndex) -> crate::Vec<NodeIndex> {
+		let pos = helper::leaf_index_to_pos(leaf_index);
+		let num_parents = leaf_index.trailing_ones() as u64;
+		return (pos..=pos + num_parents).collect()
+	}
 }

 #[cfg(test)]
 mod tests {
 	use super::*;
+	use mmr_lib::helper::leaf_index_to_pos;
+
+	#[test]
+	fn should_calculate_node_index_from_leaf_index() {
+		for index in 0..100000 {
+			let pos = leaf_index_to_pos(index);
+			assert_eq!(NodesUtils::leaf_node_index_to_leaf_index(pos), index);
+		}
+	}
+
+	#[test]
+	fn should_calculate_right_branch_correctly() {
+		fn left_jump_sequence(leaf_index: LeafIndex) -> Vec<u64> {
+			let pos = leaf_index_to_pos(leaf_index);
+			let mut right_branch_ending_in_leaf = vec![pos];
+			let mut next_pos = pos + 1;
+			while mmr_lib::helper::pos_height_in_tree(next_pos) > 0 {
+				right_branch_ending_in_leaf.push(next_pos);
+				next_pos += 1;
+			}
+			right_branch_ending_in_leaf
+		}
+
+		for leaf_index in 0..100000 {
+			let pos = mmr_lib::helper::leaf_index_to_pos(leaf_index);
+			assert_eq!(NodesUtils::right_branch_ending_in_leaf(pos), left_jump_sequence(pos));
+		}
+	}
+
+	#[test]
+	fn should_calculate_rightmost_leaf_node_index_from_pos() {
+		for pos in 0..100000 {
+			let leaf_pos = NodesUtils::rightmost_leaf_node_index_from_pos(pos);
+			let leaf_index = NodesUtils::leaf_node_index_to_leaf_index(leaf_pos);
+			assert!(NodesUtils::right_branch_ending_in_leaf(leaf_index).contains(&pos));
+		}
+	}

 	#[test]
 	fn should_calculate_number_of_leaves_correctly() {
@@ -15,9 +15,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.

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

-use frame_support::traits::OnInitialize;
+use frame_support::traits::{Get, OnInitialize};
 use mmr_lib::helper;
 use sp_core::{
 	offchain::{testing::TestOffchainExt, OffchainDbExt, OffchainWorkerExt},
@@ -47,7 +51,6 @@ fn new_block() -> u64 {

 fn peaks_from_leaves_count(leaves_count: NodeIndex) -> Vec<NodeIndex> {
 	let size = utils::NodesUtils::new(leaves_count).size();
-
 	helper::get_peaks(size)
 }

@@ -73,7 +76,7 @@ fn decode_node(
 	}
 }

-fn init_chain(blocks: usize) {
+fn add_blocks(blocks: usize) {
 	// given
 	for _ in 0..blocks {
 		new_block();
@@ -115,9 +118,10 @@ fn should_start_empty() {
 fn should_append_to_mmr_when_on_initialize_is_called() {
 	let _ = env_logger::try_init();
 	let mut ext = new_test_ext();
-	ext.execute_with(|| {
+	let (parent_b1, parent_b2) = ext.execute_with(|| {
 		// when
 		new_block();
+		let parent_b1 = <frame_system::Pallet<Test>>::parent_hash();

 		// then
 		assert_eq!(crate::NumberOfLeaves::<Test>::get(), 1);
@@ -136,6 +140,7 @@ fn should_append_to_mmr_when_on_initialize_is_called() {

 		// when
 		new_block();
+		let parent_b2 = <frame_system::Pallet<Test>>::parent_hash();

 		// then
 		assert_eq!(crate::NumberOfLeaves::<Test>::get(), 2);
@@ -157,26 +162,33 @@ fn should_append_to_mmr_when_on_initialize_is_called() {
 				hex("672c04a9cd05a644789d769daa552d35d8de7c33129f8a7cbf49e595234c4854"),
 			)
 		);
-	});

+		(parent_b1, parent_b2)
+	});
 	// make sure the leaves end up in the offchain DB
 	ext.persist_offchain_overlay();
+
 	let offchain_db = ext.offchain_db();
 	assert_eq!(
-		offchain_db.get(&MMR::offchain_key(0)).map(decode_node),
+		offchain_db.get(&MMR::node_offchain_key(parent_b1, 0)).map(decode_node),
 		Some(mmr::Node::Data(((0, H256::repeat_byte(1)), LeafData::new(1),)))
 	);
 	assert_eq!(
-		offchain_db.get(&MMR::offchain_key(1)).map(decode_node),
+		offchain_db.get(&MMR::node_offchain_key(parent_b2, 1)).map(decode_node),
 		Some(mmr::Node::Data(((1, H256::repeat_byte(2)), LeafData::new(2),)))
 	);
 	assert_eq!(
-		offchain_db.get(&MMR::offchain_key(2)).map(decode_node),
+		offchain_db.get(&MMR::node_offchain_key(parent_b2, 2)).map(decode_node),
 		Some(mmr::Node::Hash(hex(
 			"672c04a9cd05a644789d769daa552d35d8de7c33129f8a7cbf49e595234c4854"
 		)))
 	);
-	assert_eq!(offchain_db.get(&MMR::offchain_key(3)), None);
+	assert_eq!(offchain_db.get(&MMR::node_offchain_key(parent_b2, 3)), None);
+
+	assert_eq!(offchain_db.get(&MMR::node_canon_offchain_key(0)), None);
+	assert_eq!(offchain_db.get(&MMR::node_canon_offchain_key(1)), None);
+	assert_eq!(offchain_db.get(&MMR::node_canon_offchain_key(2)), None);
+	assert_eq!(offchain_db.get(&MMR::node_canon_offchain_key(3)), None);
 }

 #[test]
@@ -184,7 +196,7 @@ fn should_construct_larger_mmr_correctly() {
 	let _ = env_logger::try_init();
 	new_test_ext().execute_with(|| {
 		// when
-		init_chain(7);
+		add_blocks(7);

 		// then
 		assert_eq!(crate::NumberOfLeaves::<Test>::get(), 7);
@@ -215,7 +227,7 @@ fn should_generate_proofs_correctly() {
 	let _ = env_logger::try_init();
 	let mut ext = new_test_ext();
 	// given
-	ext.execute_with(|| init_chain(7));
+	ext.execute_with(|| add_blocks(7));
 	ext.persist_offchain_overlay();

 	// Try to generate proofs now. This requires the offchain extensions to be present
@@ -283,7 +295,7 @@ fn should_generate_batch_proof_correctly() {
 	let _ = env_logger::try_init();
 	let mut ext = new_test_ext();
 	// given
-	ext.execute_with(|| init_chain(7));
+	ext.execute_with(|| add_blocks(7));
 	ext.persist_offchain_overlay();

 	// Try to generate proofs now. This requires the offchain extensions to be present
@@ -316,7 +328,7 @@ fn should_verify() {
 	// Start off with chain initialisation and storing indexing data off-chain
 	// (MMR Leafs)
 	let mut ext = new_test_ext();
-	ext.execute_with(|| init_chain(7));
+	ext.execute_with(|| add_blocks(7));
 	ext.persist_offchain_overlay();

 	// Try to generate proof now. This requires the offchain extensions to be present
@@ -328,7 +340,7 @@ fn should_verify() {
 	});

 	ext.execute_with(|| {
-		init_chain(7);
+		add_blocks(7);
 		// then
 		assert_eq!(crate::Pallet::<Test>::verify_leaves(leaves, proof5), Ok(()));
 	});
@@ -341,7 +353,7 @@ fn should_verify_batch_proof() {
 	// Start off with chain initialisation and storing indexing data off-chain
 	// (MMR Leafs)
 	let mut ext = new_test_ext();
-	ext.execute_with(|| init_chain(7));
+	ext.execute_with(|| add_blocks(7));
 	ext.persist_offchain_overlay();

 	// Try to generate proof now. This requires the offchain extensions to be present
@@ -353,7 +365,7 @@ fn should_verify_batch_proof() {
 	});

 	ext.execute_with(|| {
-		init_chain(7);
+		add_blocks(7);
 		// then
 		assert_eq!(crate::Pallet::<Test>::verify_leaves(leaves, proof), Ok(()));
 	});
@@ -366,7 +378,7 @@ fn verification_should_be_stateless() {
 	// Start off with chain initialisation and storing indexing data off-chain
 	// (MMR Leafs)
 	let mut ext = new_test_ext();
-	ext.execute_with(|| init_chain(7));
+	ext.execute_with(|| add_blocks(7));
 	ext.persist_offchain_overlay();

 	// Try to generate proof now. This requires the offchain extensions to be present
@@ -393,7 +405,7 @@ fn should_verify_batch_proof_statelessly() {
 	// Start off with chain initialisation and storing indexing data off-chain
 	// (MMR Leafs)
 	let mut ext = new_test_ext();
-	ext.execute_with(|| init_chain(7));
+	ext.execute_with(|| add_blocks(7));
 	ext.persist_offchain_overlay();

 	// Try to generate proof now. This requires the offchain extensions to be present
@@ -424,7 +436,7 @@ fn should_verify_on_the_next_block_since_there_is_no_pruning_yet() {
 	let _ = env_logger::try_init();
 	let mut ext = new_test_ext();
 	// given
-	ext.execute_with(|| init_chain(7));
+	ext.execute_with(|| add_blocks(7));

 	ext.persist_offchain_overlay();
 	register_offchain_ext(&mut ext);
@@ -438,3 +450,238 @@ fn should_verify_on_the_next_block_since_there_is_no_pruning_yet() {
 		assert_eq!(crate::Pallet::<Test>::verify_leaves(leaves, proof5), Ok(()));
 	});
 }
+
+#[test]
+fn should_verify_pruning_map() {
+	use sp_core::offchain::StorageKind;
+	use sp_io::offchain;
+
+	let _ = env_logger::try_init();
+	let mut ext = new_test_ext();
+	register_offchain_ext(&mut ext);
+
+	ext.execute_with(|| {
+		type TestPruningMap = PruningMap<Test, ()>;
+		fn offchain_decoded(key: Vec<u8>) -> Option<Vec<H256>> {
+			offchain::local_storage_get(StorageKind::PERSISTENT, &key)
+				.and_then(|v| codec::Decode::decode(&mut &*v).ok())
+		}
+
+		// test append
+		{
+			TestPruningMap::append(1, H256::repeat_byte(1));
+
+			TestPruningMap::append(2, H256::repeat_byte(21));
+			TestPruningMap::append(2, H256::repeat_byte(22));
+
+			TestPruningMap::append(3, H256::repeat_byte(31));
+			TestPruningMap::append(3, H256::repeat_byte(32));
+			TestPruningMap::append(3, H256::repeat_byte(33));
+
+			// `0` not present
+			let map_key = TestPruningMap::pruning_map_offchain_key(0);
+			assert_eq!(offchain::local_storage_get(StorageKind::PERSISTENT, &map_key), None);
+
+			// verify `1` entries
+			let map_key = TestPruningMap::pruning_map_offchain_key(1);
+			let expected = vec![H256::repeat_byte(1)];
+			assert_eq!(offchain_decoded(map_key), Some(expected));
+
+			// verify `2` entries
+			let map_key = TestPruningMap::pruning_map_offchain_key(2);
+			let expected = vec![H256::repeat_byte(21), H256::repeat_byte(22)];
+			assert_eq!(offchain_decoded(map_key), Some(expected));
+
+			// verify `3` entries
+			let map_key = TestPruningMap::pruning_map_offchain_key(3);
+			let expected =
+				vec![H256::repeat_byte(31), H256::repeat_byte(32), H256::repeat_byte(33)];
+			assert_eq!(offchain_decoded(map_key), Some(expected));
+
+			// `4` not present
+			let map_key = TestPruningMap::pruning_map_offchain_key(4);
+			assert_eq!(offchain::local_storage_get(StorageKind::PERSISTENT, &map_key), None);
+		}
+
+		// test remove
+		{
+			// `0` doesn't return anything
+			assert_eq!(TestPruningMap::remove(0), None);
+
+			// remove and verify `1` entries
+			let expected = vec![H256::repeat_byte(1)];
+			assert_eq!(TestPruningMap::remove(1), Some(expected));
+
+			// remove and verify `2` entries
+			let expected = vec![H256::repeat_byte(21), H256::repeat_byte(22)];
+			assert_eq!(TestPruningMap::remove(2), Some(expected));
+
+			// remove and verify `3` entries
+			let expected =
+				vec![H256::repeat_byte(31), H256::repeat_byte(32), H256::repeat_byte(33)];
+			assert_eq!(TestPruningMap::remove(3), Some(expected));
+
+			// `4` doesn't return anything
+			assert_eq!(TestPruningMap::remove(4), None);
+
+			// no entries left in offchain map
+			for block in 0..5 {
+				let map_key = TestPruningMap::pruning_map_offchain_key(block);
+				assert_eq!(offchain::local_storage_get(StorageKind::PERSISTENT, &map_key), None);
+			}
+		}
+	})
+}
+
+#[test]
+fn should_canonicalize_offchain() {
+	use frame_support::traits::Hooks;
+
+	let _ = env_logger::try_init();
+	let mut ext = new_test_ext();
+	register_offchain_ext(&mut ext);
+
+	// adding 13 blocks that we'll later check have been canonicalized,
+	// (test assumes `13 < frame_system::BlockHashCount`).
+	let to_canon_count = 13u32;
+
+	// add 13 blocks and verify leaves and nodes for them have been added to
+	// offchain MMR using fork-proof keys.
+	for blocknum in 0..to_canon_count {
+		ext.execute_with(|| {
+			new_block();
+			<Pallet<Test> as Hooks<BlockNumber>>::offchain_worker(blocknum.into());
+		});
+		ext.persist_offchain_overlay();
+	}
+	let offchain_db = ext.offchain_db();
+	ext.execute_with(|| {
+		// verify leaves added by blocks 1..=13
+		for block_num in 1..=to_canon_count {
+			let parent_num: BlockNumber = (block_num - 1).into();
+			let leaf_index = u64::from(block_num - 1);
+			let pos = helper::leaf_index_to_pos(leaf_index.into());
+			// not canon,
+			assert_eq!(offchain_db.get(&MMR::node_canon_offchain_key(pos)), None);
+			let parent_hash = <frame_system::Pallet<Test>>::block_hash(parent_num);
+			// but available in fork-proof storage.
+			assert_eq!(
+				offchain_db.get(&MMR::node_offchain_key(parent_hash, pos)).map(decode_node),
+				Some(mmr::Node::Data((
+					(leaf_index, H256::repeat_byte(u8::try_from(block_num).unwrap())),
+					LeafData::new(block_num.into()),
+				)))
+			);
+		}
+
+		// verify a couple of nodes and peaks:
+		// 		`pos` is node to verify,
+		// 		`leaf_index` is leaf that added node `pos`,
+		// 		`expected` is expected value of node at `pos`.
+		let verify = |pos: NodeIndex, leaf_index: LeafIndex, expected: H256| {
+			let parent_num: BlockNumber = leaf_index.try_into().unwrap();
+			let parent_hash = <frame_system::Pallet<Test>>::block_hash(parent_num);
+			// not canon,
+			assert_eq!(offchain_db.get(&MMR::node_canon_offchain_key(pos)), None);
+			// but available in fork-proof storage.
+			assert_eq!(
+				offchain_db.get(&MMR::node_offchain_key(parent_hash, pos)).map(decode_node),
+				Some(mmr::Node::Hash(expected))
+			);
+		};
+		verify(2, 1, hex("672c04a9cd05a644789d769daa552d35d8de7c33129f8a7cbf49e595234c4854"));
+		verify(13, 7, hex("441bf63abc7cf9b9e82eb57b8111c883d50ae468d9fd7f301e12269fc0fa1e75"));
+		verify(21, 11, hex("f323ac1a7f56de5f40ed8df3e97af74eec0ee9d72883679e49122ffad2ffd03b"));
+	});
+
+	// add another `frame_system::BlockHashCount` blocks and verify all nodes and leaves
+	// added by our original `to_canon_count` blocks have now been canonicalized in offchain db.
+	let block_hash_size: u64 = <Test as frame_system::Config>::BlockHashCount::get();
+	let base = to_canon_count;
+	for blocknum in base..(base + u32::try_from(block_hash_size).unwrap()) {
+		ext.execute_with(|| {
+			new_block();
+			<Pallet<Test> as Hooks<BlockNumber>>::offchain_worker(blocknum.into());
+		});
+		ext.persist_offchain_overlay();
+	}
+	ext.execute_with(|| {
+		// verify leaves added by blocks 1..=13, should be in offchain under canon key.
+		for block_num in 1..=to_canon_count {
+			let leaf_index = u64::from(block_num - 1);
+			let pos = helper::leaf_index_to_pos(leaf_index.into());
+			let parent_num: BlockNumber = (block_num - 1).into();
+			let parent_hash = <frame_system::Pallet<Test>>::block_hash(parent_num);
+			// no longer available in fork-proof storage (was pruned),
+			assert_eq!(offchain_db.get(&MMR::node_offchain_key(parent_hash, pos)), None);
+			// but available using canon key.
+			assert_eq!(
+				offchain_db.get(&MMR::node_canon_offchain_key(pos)).map(decode_node),
+				Some(mmr::Node::Data((
+					(leaf_index, H256::repeat_byte(u8::try_from(block_num).unwrap())),
+					LeafData::new(block_num.into()),
+				)))
+			);
+		}
+
+		// also check some nodes and peaks:
+		// 		`pos` is node to verify,
+		// 		`leaf_index` is leaf that added node `pos`,
+		// 		`expected` is expected value of node at `pos`.
+		let verify = |pos: NodeIndex, leaf_index: LeafIndex, expected: H256| {
+			let parent_num: BlockNumber = leaf_index.try_into().unwrap();
+			let parent_hash = <frame_system::Pallet<Test>>::block_hash(parent_num);
+			// no longer available in fork-proof storage (was pruned),
+			assert_eq!(offchain_db.get(&MMR::node_offchain_key(parent_hash, pos)), None);
+			// but available using canon key.
+			assert_eq!(
+				offchain_db.get(&MMR::node_canon_offchain_key(pos)).map(decode_node),
+				Some(mmr::Node::Hash(expected))
+			);
+		};
+		verify(2, 1, hex("672c04a9cd05a644789d769daa552d35d8de7c33129f8a7cbf49e595234c4854"));
+		verify(13, 7, hex("441bf63abc7cf9b9e82eb57b8111c883d50ae468d9fd7f301e12269fc0fa1e75"));
+		verify(21, 11, hex("f323ac1a7f56de5f40ed8df3e97af74eec0ee9d72883679e49122ffad2ffd03b"));
+	});
+}
+
+#[test]
+fn should_verify_canonicalized() {
+	use frame_support::traits::Hooks;
+	let _ = env_logger::try_init();
+
+	// How deep is our fork-aware storage (in terms of blocks/leaves, nodes will be more).
+	let block_hash_size: u64 = <Test as frame_system::Config>::BlockHashCount::get();
+
+	// Start off with chain initialisation and storing indexing data off-chain.
+	// Create twice as many leaf entries than our fork-aware capacity,
+	// resulting in ~half of MMR storage to use canonical keys and the other half fork-aware keys.
+	// Verify that proofs can be generated (using leaves and nodes from full set) and verified.
+	let mut ext = new_test_ext();
+	register_offchain_ext(&mut ext);
+	for blocknum in 0u32..(2 * block_hash_size).try_into().unwrap() {
+		ext.execute_with(|| {
+			new_block();
+			<Pallet<Test> as Hooks<BlockNumber>>::offchain_worker(blocknum.into());
+		});
+		ext.persist_offchain_overlay();
+	}
+
+	// Generate proofs for some blocks.
+	let (leaves, proofs) =
+		ext.execute_with(|| crate::Pallet::<Test>::generate_batch_proof(vec![0, 4, 5, 7]).unwrap());
+	// Verify all previously generated proofs.
+	ext.execute_with(|| {
+		assert_eq!(crate::Pallet::<Test>::verify_leaves(leaves, proofs), Ok(()));
+	});
+
+	// Generate proofs for some new blocks.
+	let (leaves, proofs) = ext.execute_with(|| {
+		crate::Pallet::<Test>::generate_batch_proof(vec![block_hash_size + 7]).unwrap()
+	});
+	// Add some more blocks then verify all previously generated proofs.
+	ext.execute_with(|| {
+		add_blocks(7);
+		assert_eq!(crate::Pallet::<Test>::verify_leaves(leaves, proofs), Ok(()));
+	});
+}
@@ -81,7 +81,7 @@ pub struct Proof<Hash> {

 /// A full leaf content stored in the offchain-db.
 pub trait FullLeaf: Clone + PartialEq + fmt::Debug {
-	/// Encode the leaf either in it's full or compact form.
+	/// Encode the leaf either in its full or compact form.
 	///
 	/// NOTE the encoding returned here MUST be `Decode`able into `FullLeaf`.
 	fn using_encoded<R, F: FnOnce(&[u8]) -> R>(&self, f: F, compact: bool) -> R;
@@ -167,18 +167,18 @@ impl EncodableOpaqueLeaf {
 	}
 }

-/// An element representing either full data or it's hash.
+/// An element representing either full data or its hash.
 ///
 /// See [Compact] to see how it may be used in practice to reduce the size
 /// of proofs in case multiple [LeafDataProvider]s are composed together.
 /// This is also used internally by the MMR to differentiate leaf nodes (data)
 /// and inner nodes (hashes).
 ///
-/// [DataOrHash::hash] method calculates the hash of this element in it's compact form,
+/// [DataOrHash::hash] method calculates the hash of this element in its compact form,
 /// so should be used instead of hashing the encoded form (which will always be non-compact).
 #[derive(RuntimeDebug, Clone, PartialEq)]
 pub enum DataOrHash<H: traits::Hash, L> {
-	/// Arbitrary data in it's full form.
+	/// Arbitrary data in its full form.
 	Data(L),
 	/// A hash of some data.
 	Hash(H::Output),
@@ -339,7 +339,7 @@ where
 	A: FullLeaf,
 	B: FullLeaf,
 {
-	/// Retrieve a hash of this item in it's compact form.
+	/// Retrieve a hash of this item in its compact form.
 	pub fn hash(&self) -> H::Output {
 		self.using_encoded(<H as traits::Hash>::hash, true)
 	}
@@ -447,7 +447,7 @@ sp_api::decl_runtime_apis! {
 		/// Note this function does not require any on-chain storage - the
 		/// proof is verified against given MMR root hash.
 		///
-		/// The leaf data is expected to be encoded in it's compact form.
+		/// The leaf data is expected to be encoded in its compact form.
 		fn verify_proof_stateless(root: Hash, leaf: EncodableOpaqueLeaf, proof: Proof<Hash>)
 			-> Result<(), Error>;