pallet-mmr: move offchain logic to client-side gadget (#12753)

* Move MMR utils methods from pallet to primitives

Signed-off-by: Serban Iorga <serban@parity.io>

* Add method to MmrApi

* Move forks expanding logic from babe to primitives

* Implement MMR gadget

* Remove prunning logic from the MMR pallet

* Code review changes: 1st iteration

* Replace MaybeCanonEngine with CanonEngineBuilder

* fix mmr_leaves_count() for kitchen sink demo

* Update client/merkle-mountain-range/src/canon_engine.rs

Co-authored-by: Adrian Catangiu <adrian@parity.io>

* Code review changes: 2nd iteration

* fix INDEXING_PREFIX

* impl review comments

* add documentation and minor rename

Signed-off-by: Serban Iorga <serban@parity.io>
Co-authored-by: Adrian Catangiu <adrian@parity.io>

substrate/primitives/blockchain/src/backend.rs

@@ -21,9 +21,10 @@ use log::warn;
 use parking_lot::RwLock;
 use sp_runtime::{
 	generic::BlockId,
-	traits::{Block as BlockT, Header as HeaderT, NumberFor},
+	traits::{Block as BlockT, Header as HeaderT, NumberFor, Saturating},
 	Justifications,
 };
+use std::collections::btree_set::BTreeSet;
 
 use crate::header_metadata::HeaderMetadata;
 
@@ -84,6 +85,77 @@ pub trait HeaderBackend<Block: BlockT>: Send + Sync {
 	}
 }
 
+/// Handles stale forks.
+pub trait ForkBackend<Block: BlockT>:
+	HeaderMetadata<Block> + HeaderBackend<Block> + Send + Sync
+{
+	/// Best effort to get all the header hashes that are part of the provided forks
+	/// starting only from the fork heads.
+	///
+	/// The function tries to reconstruct the route from the fork head to the canonical chain.
+	/// If any of the hashes on the route can't be found in the db, the function won't be able
+	/// to reconstruct the route anymore. In this case it will give up expanding the current fork,
+	/// move on to the next ones and at the end it will return an error that also contains
+	/// the partially expanded forks.
+	fn expand_forks(
+		&self,
+		fork_heads: &[Block::Hash],
+	) -> std::result::Result<BTreeSet<Block::Hash>, (BTreeSet<Block::Hash>, Error)> {
+		let mut missing_blocks = vec![];
+		let mut expanded_forks = BTreeSet::new();
+		for fork_head in fork_heads {
+			let mut route_head = *fork_head;
+			// Insert stale blocks hashes until canonical chain is reached.
+			// If we reach a block that is already part of the `expanded_forks` we can stop
+			// processing the fork.
+			while expanded_forks.insert(route_head) {
+				match self.header_metadata(route_head) {
+					Ok(meta) => {
+						// If the parent is part of the canonical chain or there doesn't exist a
+						// block hash for the parent number (bug?!), we can abort adding blocks.
+						let parent_number = meta.number.saturating_sub(1u32.into());
+						match self.hash(parent_number) {
+							Ok(Some(parent_hash)) =>
+								if parent_hash == meta.parent {
+									break
+								},
+							Ok(None) | Err(_) => {
+								missing_blocks.push(BlockId::<Block>::Number(parent_number));
+								break
+							},
+						}
+
+						route_head = meta.parent;
+					},
+					Err(_e) => {
+						missing_blocks.push(BlockId::<Block>::Hash(route_head));
+						break
+					},
+				}
+			}
+		}
+
+		if !missing_blocks.is_empty() {
+			return Err((
+				expanded_forks,
+				Error::UnknownBlocks(format!(
+					"Missing stale headers {:?} while expanding forks {:?}.",
+					fork_heads, missing_blocks
+				)),
+			))
+		}
+
+		Ok(expanded_forks)
+	}
+}
+
+impl<Block, T> ForkBackend<Block> for T
+where
+	Block: BlockT,
+	T: HeaderMetadata<Block> + HeaderBackend<Block> + Send + Sync,
+{
+}
+
 /// Blockchain database backend. Does not perform any validation.
 pub trait Backend<Block: BlockT>:
 	HeaderBackend<Block> + HeaderMetadata<Block, Error = Error>

substrate/primitives/blockchain/src/error.rs

@@ -59,6 +59,9 @@ pub enum Error {
 	#[error("UnknownBlock: {0}")]
 	UnknownBlock(String),
 
+	#[error("UnknownBlocks: {0}")]
+	UnknownBlocks(String),
+
 	#[error(transparent)]
 	ApplyExtrinsicFailed(#[from] ApplyExtrinsicFailed),
 

substrate/primitives/merkle-mountain-range/Cargo.toml

@@ -15,6 +15,7 @@ targets = ["x86_64-unknown-linux-gnu"]
 codec = { package = "parity-scale-codec", version = "3.0.0", default-features = false }
 scale-info = { version = "2.1.1", default-features = false, features = ["derive"] }
 log = { version = "0.4.17", default-features = false }
+mmr-lib = { package = "ckb-merkle-mountain-range", version = "0.5.2", default-features = false }
 serde = { version = "1.0.136", features = ["derive"], optional = true }
 sp-api = { version = "4.0.0-dev", default-features = false, path = "../api" }
 sp-core = { version = "7.0.0", default-features = false, path = "../core" }
@@ -31,6 +32,7 @@ default = ["std"]
 std = [
 	"codec/std",
 	"log/std",
+	"mmr-lib/std",
 	"serde",
 	"sp-api/std",
 	"sp-core/std",

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

@@ -20,6 +20,8 @@
 #![cfg_attr(not(feature = "std"), no_std)]
 #![warn(missing_docs)]
 
+pub use mmr_lib;
+
 use scale_info::TypeInfo;
 use sp_debug_derive::RuntimeDebug;
 use sp_runtime::traits;
@@ -27,6 +29,11 @@ use sp_std::fmt;
 #[cfg(not(feature = "std"))]
 use sp_std::prelude::Vec;
 
+pub mod utils;
+
+/// Prefix for elements stored in the Off-chain DB via Indexing API.
+pub const INDEXING_PREFIX: &'static [u8] = b"mmr";
+
 /// A type to describe node position in the MMR (node index).
 pub type NodeIndex = u64;
 
@@ -357,8 +364,8 @@ pub struct Proof<Hash> {
 #[derive(RuntimeDebug, codec::Encode, codec::Decode, PartialEq, Eq)]
 pub enum Error {
 	/// Error during translation of a block number into a leaf index.
-	#[cfg_attr(feature = "std", error("Error translation block number into leaf index"))]
-	BlockNumToLeafIndex,
+	#[cfg_attr(feature = "std", error("Error performing numeric op"))]
+	InvalidNumericOp,
 	/// Error while pushing new node.
 	#[cfg_attr(feature = "std", error("Error pushing new node"))]
 	Push,
@@ -419,6 +426,9 @@ sp_api::decl_runtime_apis! {
 		/// Return the on-chain MMR root hash.
 		fn mmr_root() -> Result<Hash, Error>;
 
+		/// Return the number of MMR blocks in the chain.
+		fn mmr_leaf_count() -> Result<LeafIndex, Error>;
+
 		/// Generate MMR proof for a series of block numbers. If `best_known_block_number = Some(n)`,
 		/// use historical MMR state at given block height `n`. Else, use current MMR state.
 		fn generate_proof(

substrate/primitives/merkle-mountain-range/src/utils.rs

@@ -0,0 +1,215 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2020-2022 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Merkle Mountain Range utilities.
+
+use codec::Encode;
+use mmr_lib::helper;
+
+use sp_runtime::traits::{CheckedAdd, CheckedSub, Header, One};
+#[cfg(not(feature = "std"))]
+use sp_std::prelude::Vec;
+
+use crate::{Error, LeafIndex, NodeIndex};
+
+/// Get the first block with MMR.
+pub fn first_mmr_block_num<H: Header>(
+	best_block_num: H::Number,
+	mmr_leaf_count: LeafIndex,
+) -> Result<H::Number, Error> {
+	let mmr_blocks_count = mmr_leaf_count.try_into().map_err(|_| {
+		Error::InvalidNumericOp
+			.log_debug("The number of leaves couldn't be converted to a block number.")
+	})?;
+	best_block_num
+		.checked_sub(&mmr_blocks_count)
+		.and_then(|last_non_mmr_block| last_non_mmr_block.checked_add(&One::one()))
+		.ok_or_else(|| {
+			Error::InvalidNumericOp
+				.log_debug("The best block should be greater than the number of mmr blocks.")
+		})
+}
+
+/// Convert a block number into a leaf index.
+pub fn block_num_to_leaf_index<H: Header>(
+	block_num: H::Number,
+	first_mmr_block_num: H::Number,
+) -> Result<LeafIndex, Error> {
+	let leaf_idx = block_num.checked_sub(&first_mmr_block_num).ok_or_else(|| {
+		Error::InvalidNumericOp
+			.log_debug("The provided block should be greater than the first mmr block.")
+	})?;
+
+	leaf_idx.try_into().map_err(|_| {
+		Error::InvalidNumericOp.log_debug("Couldn't convert the leaf index to `LeafIndex`.")
+	})
+}
+
+/// MMR nodes & size -related utilities.
+pub struct NodesUtils {
+	no_of_leaves: LeafIndex,
+}
+
+impl NodesUtils {
+	/// Create new instance of MMR nodes utilities for given number of leaves.
+	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) -> NodeIndex {
+		self.number_of_leaves().count_ones() as NodeIndex
+	}
+
+	/// Return the number of leaves in the MMR.
+	pub fn number_of_leaves(&self) -> LeafIndex {
+		self.no_of_leaves
+	}
+
+	/// Calculate the total size of MMR (number of nodes).
+	pub fn size(&self) -> NodeIndex {
+		2 * self.no_of_leaves - self.number_of_peaks()
+	}
+
+	/// 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) -> 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()
+	}
+
+	/// 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.
+	pub fn node_temp_offchain_key<H: Header>(
+		prefix: &[u8],
+		pos: NodeIndex,
+		parent_hash: H::Hash,
+	) -> Vec<u8> {
+		(prefix, pos, parent_hash).encode()
+	}
+
+	/// Build canonical offchain key for node `pos` in MMR.
+	///
+	/// Used for nodes added by now finalized blocks.
+	/// Never read keys using `node_canon_offchain_key` unless you sure that
+	/// there's no `node_offchain_key` key in the storage.
+	pub fn node_canon_offchain_key(prefix: &[u8], pos: NodeIndex) -> sp_std::prelude::Vec<u8> {
+		(prefix, pos).encode()
+	}
+}
+
+#[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_depth_correctly() {
+		assert_eq!(
+			vec![0, 1, 2, 3, 4, 9, 15, 21]
+				.into_iter()
+				.map(|n| NodesUtils::new(n).number_of_leaves())
+				.collect::<Vec<_>>(),
+			vec![0, 1, 2, 3, 4, 9, 15, 21]
+		);
+	}
+
+	#[test]
+	fn should_calculate_number_of_peaks_correctly() {
+		assert_eq!(
+			vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 21]
+				.into_iter()
+				.map(|n| NodesUtils::new(n).number_of_peaks())
+				.collect::<Vec<_>>(),
+			vec![0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4, 3]
+		);
+	}
+
+	#[test]
+	fn should_calculate_the_size_correctly() {
+		let leaves = vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 21];
+		let sizes = vec![0, 1, 3, 4, 7, 8, 10, 11, 15, 16, 18, 19, 22, 23, 25, 26, 39];
+		assert_eq!(
+			leaves
+				.clone()
+				.into_iter()
+				.map(|n| NodesUtils::new(n).size())
+				.collect::<Vec<_>>(),
+			sizes.clone()
+		);
+	}
+}