Skip to content
This repository was archived by the owner on Nov 15, 2023. It is now read-only.

Commit 7a531b0

Browse files
aurexavtomusdrw
andauthored
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 <[email protected]>
1 parent c3de648 commit 7a531b0

File tree

11 files changed

+163
-59
lines changed

11 files changed

+163
-59
lines changed

Cargo.lock

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

bin/node/runtime/src/lib.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1559,7 +1559,7 @@ impl_runtime_apis! {
15591559
Block,
15601560
mmr::Hash,
15611561
> for Runtime {
1562-
fn generate_proof(leaf_index: u64)
1562+
fn generate_proof(leaf_index: pallet_mmr::primitives::LeafIndex)
15631563
-> Result<(mmr::EncodableOpaqueLeaf, mmr::Proof<mmr::Hash>), mmr::Error>
15641564
{
15651565
Mmr::generate_proof(leaf_index)

frame/merkle-mountain-range/Cargo.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ targets = ["x86_64-unknown-linux-gnu"]
1414
[dependencies]
1515
codec = { package = "parity-scale-codec", version = "2.0.0", default-features = false }
1616
scale-info = { version = "1.0", default-features = false, features = ["derive"] }
17-
mmr-lib = { package = "ckb-merkle-mountain-range", default-features = false, version = "0.3.1" }
17+
mmr-lib = { package = "ckb-merkle-mountain-range", default-features = false, version = "0.3.2" }
1818

1919
sp-core = { version = "4.0.0-dev", default-features = false, path = "../../primitives/core" }
2020
sp-io = { version = "4.0.0-dev", default-features = false, path = "../../primitives/io" }

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

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,16 @@ use sp_std::fmt;
2626
#[cfg(not(feature = "std"))]
2727
use sp_std::prelude::Vec;
2828

29+
/// A type to describe node position in the MMR (node index).
30+
pub type NodeIndex = u64;
31+
32+
/// A type to describe leaf position in the MMR.
33+
///
34+
/// Note this is different from [`NodeIndex`], which can be applied to
35+
/// both leafs and inner nodes. Leafs will always have consecutive `LeafIndex`,
36+
/// but might be actually at different positions in the MMR `NodeIndex`.
37+
pub type LeafIndex = u64;
38+
2939
/// A provider of the MMR's leaf data.
3040
pub trait LeafDataProvider {
3141
/// 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);
275285
#[derive(codec::Encode, codec::Decode, RuntimeDebug, Clone, PartialEq, Eq)]
276286
pub struct Proof<Hash> {
277287
/// The index of the leaf the proof is for.
278-
pub leaf_index: u64,
288+
pub leaf_index: LeafIndex,
279289
/// Number of leaves in MMR, when the proof was generated.
280-
pub leaf_count: u64,
290+
pub leaf_count: NodeIndex,
281291
/// Proof elements (hashes of siblings of inner nodes on the path to the leaf).
282292
pub items: Vec<Hash>,
283293
}
@@ -402,7 +412,7 @@ sp_api::decl_runtime_apis! {
402412
/// API to interact with MMR pallet.
403413
pub trait MmrApi<Hash: codec::Codec> {
404414
/// Generate MMR proof for a leaf under given index.
405-
fn generate_proof(leaf_index: u64) -> Result<(EncodableOpaqueLeaf, Proof<Hash>), Error>;
415+
fn generate_proof(leaf_index: LeafIndex) -> Result<(EncodableOpaqueLeaf, Proof<Hash>), Error>;
406416

407417
/// Verify MMR proof against on-chain MMR.
408418
///

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ use sp_blockchain::HeaderBackend;
3232
use sp_core::Bytes;
3333
use sp_runtime::{generic::BlockId, traits::Block as BlockT};
3434

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

3737
/// Retrieved MMR leaf and its proof.
3838
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
@@ -71,7 +71,7 @@ pub trait MmrApi<BlockHash> {
7171
#[rpc(name = "mmr_generateProof")]
7272
fn generate_proof(
7373
&self,
74-
leaf_index: u64,
74+
leaf_index: LeafIndex,
7575
at: Option<BlockHash>,
7676
) -> Result<LeafProof<BlockHash>>;
7777
}
@@ -98,7 +98,7 @@ where
9898
{
9999
fn generate_proof(
100100
&self,
101-
leaf_index: u64,
101+
leaf_index: LeafIndex,
102102
at: Option<<Block as BlockT>::Hash>,
103103
) -> Result<LeafProof<<Block as BlockT>::Hash>> {
104104
let api = self.client.runtime_api();

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ benchmarks_instance_pallet! {
2525
on_initialize {
2626
let x in 1 .. 1_000;
2727

28-
let leaves = x as u64;
28+
let leaves = x as NodeIndex;
2929
}: {
3030
for b in 0..leaves {
3131
Pallet::<T, I>::on_initialize((b as u32).into());

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

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -70,10 +70,10 @@ mod mock;
7070
mod tests;
7171

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

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

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

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

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

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

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

266-
let mmr: ModuleMmr<mmr::storage::RuntimeStorage, T, I> = mmr::Mmr::new(proof.leaf_count);
266+
let mmr: ModuleMmr<mmr::storage::OffchainStorage, T, I> = mmr::Mmr::new(proof.leaf_count);
267267
let is_valid = mmr.verify_leaf_proof(leaf, proof)?;
268268
if is_valid {
269269
Ok(())

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

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ use crate::{
2121
utils::NodesUtils,
2222
Hasher, Node, NodeOf,
2323
},
24-
primitives::{self, Error},
24+
primitives::{self, Error, NodeIndex},
2525
Config, HashingOf,
2626
};
2727
#[cfg(not(feature = "std"))]
@@ -60,7 +60,7 @@ where
6060
Storage<StorageType, T, I, L>: mmr_lib::MMRStore<NodeOf<T, I, L>>,
6161
{
6262
mmr: mmr_lib::MMR<NodeOf<T, I, L>, Hasher<HashingOf<T, I>, L>, Storage<StorageType, T, I, L>>,
63-
leaves: u64,
63+
leaves: NodeIndex,
6464
}
6565

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

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

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

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

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

Lines changed: 81 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -18,19 +18,24 @@
1818
//! A MMR storage implementations.
1919
2020
use codec::Encode;
21+
use frame_support::log;
22+
use mmr_lib::helper;
23+
use sp_io::offchain_index;
24+
use sp_std::iter::Peekable;
2125
#[cfg(not(feature = "std"))]
22-
use sp_std::prelude::Vec;
26+
use sp_std::prelude::*;
2327

2428
use crate::{
25-
mmr::{Node, NodeOf},
26-
primitives, Config, Nodes, NumberOfLeaves, Pallet,
29+
mmr::{utils::NodesUtils, Node, NodeOf},
30+
primitives::{self, NodeIndex},
31+
Config, Nodes, NumberOfLeaves, Pallet,
2732
};
2833

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

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

85-
fn append(&mut self, pos: u64, elems: Vec<NodeOf<T, I, L>>) -> mmr_lib::Result<()> {
86-
let mut leaves = crate::NumberOfLeaves::<T, I>::get();
87-
let mut size = crate::mmr::utils::NodesUtils::new(leaves).size();
90+
fn append(&mut self, pos: NodeIndex, elems: Vec<NodeOf<T, I, L>>) -> mmr_lib::Result<()> {
91+
if elems.is_empty() {
92+
return Ok(())
93+
}
94+
95+
sp_std::if_std! {
96+
log::trace!("elems: {:?}", elems.iter().map(|elem| elem.hash()).collect::<Vec<_>>());
97+
}
98+
99+
let leaves = NumberOfLeaves::<T, I>::get();
100+
let size = NodesUtils::new(leaves).size();
101+
88102
if pos != size {
89103
return Err(mmr_lib::Error::InconsistentStore)
90104
}
91105

106+
let new_size = size + elems.len() as NodeIndex;
107+
108+
// A sorted (ascending) iterator over peak indices to prune and persist.
109+
let (peaks_to_prune, mut peaks_to_store) = peaks_to_prune_and_store(size, new_size);
110+
111+
// Now we are going to iterate over elements to insert
112+
// and keep track of the current `node_index` and `leaf_index`.
113+
let mut leaf_index = leaves;
114+
let mut node_index = size;
115+
92116
for elem in elems {
93-
// on-chain we only store the hash (even if it's a leaf)
94-
<Nodes<T, I>>::insert(size, elem.hash());
95-
// Indexing API is used to store the full leaf content.
96-
let key = Pallet::<T, I>::offchain_key(size);
97-
elem.using_encoded(|elem| sp_io::offchain_index::set(&key, elem));
98-
size += 1;
117+
// Indexing API is used to store the full node content (both leaf and inner).
118+
elem.using_encoded(|elem| {
119+
offchain_index::set(&Pallet::<T, I>::offchain_key(node_index), elem)
120+
});
121+
122+
// On-chain we are going to only store new peaks.
123+
if peaks_to_store.next_if_eq(&node_index).is_some() {
124+
<Nodes<T, I>>::insert(node_index, elem.hash());
125+
}
99126

127+
// Increase the indices.
100128
if let Node::Data(..) = elem {
101-
leaves += 1;
129+
leaf_index += 1;
102130
}
131+
node_index += 1;
103132
}
104133

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

107142
Ok(())
108143
}
109144
}
145+
146+
fn peaks_to_prune_and_store(
147+
old_size: NodeIndex,
148+
new_size: NodeIndex,
149+
) -> (impl Iterator<Item = NodeIndex>, Peekable<impl Iterator<Item = NodeIndex>>) {
150+
// A sorted (ascending) collection of peak indices before and after insertion.
151+
// both collections may share a common prefix.
152+
let peaks_before = if old_size == 0 { vec![] } else { helper::get_peaks(old_size) };
153+
let peaks_after = helper::get_peaks(new_size);
154+
sp_std::if_std! {
155+
log::trace!("peaks_before: {:?}", peaks_before);
156+
log::trace!("peaks_after: {:?}", peaks_after);
157+
}
158+
let mut peaks_before = peaks_before.into_iter().peekable();
159+
let mut peaks_after = peaks_after.into_iter().peekable();
160+
161+
// Consume a common prefix between `peaks_before` and `peaks_after`,
162+
// since that's something we will not be touching anyway.
163+
while peaks_before.peek() == peaks_after.peek() {
164+
peaks_before.next();
165+
peaks_after.next();
166+
}
167+
168+
// what's left in both collections is:
169+
// 1. Old peaks to remove from storage
170+
// 2. New peaks to persist in storage
171+
(peaks_before, peaks_after)
172+
}

0 commit comments

Comments
 (0)