Prune leaves on finality (#1871)

* support multiple leaves at same height in leaves set

* prune leaves on finalization

* test leaves behavior in client-db

* fix no-std compilation

substrate/core/client/db/src/lib.rs

@@ -51,7 +51,7 @@ use state_machine::backend::Backend as StateBackend;
 use executor::RuntimeInfo;
 use state_machine::{CodeExecutor, DBValue};
 use crate::utils::{Meta, db_err, meta_keys, open_database, read_db, block_id_to_lookup_key, read_meta};
-use client::LeafSet;
+use client::leaves::{LeafSet, FinalizationDisplaced};
 use client::children;
 use state_db::StateDb;
 use crate::storage_cache::{CachingState, SharedCache, new_shared_cache};
@@ -714,11 +714,18 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 		header: &Block::Header,
 		last_finalized: Option<Block::Hash>,
 		justification: Option<Justification>,
+		finalization_displaced: &mut Option<FinalizationDisplaced<Block::Hash, NumberFor<Block>>>,
 	) -> Result<(Block::Hash, <Block::Header as HeaderT>::Number, bool, bool), client::error::Error> {
 		// TODO: ensure best chain contains this block.
 		let number = *header.number();
 		self.ensure_sequential_finalization(header, last_finalized)?;
-		self.note_finalized(transaction, header, *hash)?;
+		self.note_finalized(
+			transaction,
+			header,
+			*hash,
+			finalization_displaced,
+		)?;
+
 		if let Some(justification) = justification {
 			transaction.put(
 				columns::JUSTIFICATION,
@@ -767,10 +774,13 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 		-> Result<(), client::error::Error>
 	{
 		let mut transaction = DBTransaction::new();
+		let mut finalization_displaced_leaves = None;
+
 		operation.apply_aux(&mut transaction);
 
 		let mut meta_updates = Vec::new();
 		let mut last_finalized_hash = self.blockchain.meta.read().finalized_hash;
+
 		if !operation.finalized_blocks.is_empty() {
 			for (block, justification) in operation.finalized_blocks {
 				let block_hash = self.blockchain.expect_block_hash_from_id(&block)?;
@@ -782,6 +792,7 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 					&block_header,
 					Some(last_finalized_hash),
 					justification,
+					&mut finalization_displaced_leaves,
 				)?);
 				last_finalized_hash = block_hash;
 			}
@@ -846,7 +857,12 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 			if finalized {
 				// TODO: ensure best chain contains this block.
 				self.ensure_sequential_finalization(header, Some(last_finalized_hash))?;
-				self.note_finalized(&mut transaction, header, hash)?;
+				self.note_finalized(
+					&mut transaction,
+					header,
+					hash,
+					&mut finalization_displaced_leaves,
+				)?;
 			} else {
 				// canonicalize blocks which are old enough, regardless of finality.
 				self.force_delayed_canonicalize(&mut transaction, hash, *header.number())?
@@ -892,9 +908,16 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 
 		if let Some((number, hash, enacted, retracted, displaced_leaf, is_best)) = imported {
 			if let Err(e) = write_result {
+				let mut leaves = self.blockchain.leaves.write();
+				let mut undo = leaves.undo();
 				if let Some(displaced_leaf) = displaced_leaf {
-					self.blockchain.leaves.write().undo(displaced_leaf);
+					undo.undo_import(displaced_leaf);
 				}
+
+				if let Some(finalization_displaced) = finalization_displaced_leaves {
+					undo.undo_finalization(finalization_displaced);
+				}
+
 				return Err(e)
 			}
 
@@ -924,6 +947,7 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 		transaction: &mut DBTransaction,
 		f_header: &Block::Header,
 		f_hash: Block::Hash,
+		displaced: &mut Option<FinalizationDisplaced<Block::Hash, NumberFor<Block>>>
 	) -> Result<(), client::error::Error> where
 		Block: BlockT<Hash=H256>,
 	{
@@ -947,6 +971,12 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 			self.changes_tries_storage.prune(changes_trie_config, transaction, f_hash, f_num);
 		}
 
+		let new_displaced = self.blockchain.leaves.write().finalize_height(f_num);
+		match displaced {
+			x @ &mut None => *x = Some(new_displaced),
+			&mut Some(ref mut displaced) => displaced.merge(new_displaced),
+		}
+
 		Ok(())
 	}
 }
@@ -1036,22 +1066,27 @@ impl<Block> client::backend::Backend<Block, Blake2Hasher> for Backend<Block> whe
 		let mut transaction = DBTransaction::new();
 		let hash = self.blockchain.expect_block_hash_from_id(&block)?;
 		let header = self.blockchain.expect_header(block)?;
-		let commit = || {
+		let mut displaced = None;
+		let commit = |displaced| {
 			let (hash, number, is_best, is_finalized) = self.finalize_block_with_transaction(
 				&mut transaction,
 				&hash,
 				&header,
 				None,
 				justification,
+				displaced,
 			)?;
 			self.storage.db.write(transaction).map_err(db_err)?;
 			self.blockchain.update_meta(hash, number, is_best, is_finalized);
 			Ok(())
 		};
-		match commit() {
+		match commit(&mut displaced) {
 			Ok(()) => self.storage.state_db.apply_pending(),
 			e @ Err(_) => {
 				self.storage.state_db.revert_pending();
+				if let Some(displaced) = displaced {
+					self.blockchain.leaves.write().undo().undo_finalization(displaced);
+				}
 				return e;
 			}
 		}
@@ -1152,6 +1187,7 @@ mod tests {
 	use super::*;
 	use crate::columns;
 	use client::backend::Backend as BTrait;
+	use client::blockchain::Backend as BLBTrait;
 	use client::backend::BlockImportOperation as Op;
 	use runtime_primitives::testing::{Header, Block as RawBlock, ExtrinsicWrapper};
 	use runtime_primitives::traits::{Hash, BlakeTwo256};
@@ -1789,8 +1825,6 @@ mod tests {
 				BlockId::Hash(block1),
 			).unwrap();
 
-			println!("{:?}", tree_route);
-
 			assert_eq!(tree_route.common_block().hash, block0);
 			assert!(tree_route.retracted().is_empty());
 			assert_eq!(tree_route.enacted().iter().map(|r| r.hash).collect::<Vec<_>>(), vec![block1]);
@@ -1815,6 +1849,30 @@ mod tests {
 		test_client::trait_tests::test_blockchain_query_by_number_gets_canonical(backend);
 	}
 
+	#[test]
+	fn test_leaves_pruned_on_finality() {
+		let backend: Backend<Block> = Backend::new_test(10, 10);
+		let block0 = insert_header(&backend, 0, Default::default(), Default::default(), Default::default());
+
+		let block1_a = insert_header(&backend, 1, block0, Default::default(), Default::default());
+		let block1_b = insert_header(&backend, 1, block0, Default::default(), [1; 32].into());
+		let block1_c = insert_header(&backend, 1, block0, Default::default(), [2; 32].into());
+
+		assert_eq!(backend.blockchain().leaves().unwrap(), vec![block1_a, block1_b, block1_c]);
+
+		let block2_a = insert_header(&backend, 2, block1_a, Default::default(), Default::default());
+		let block2_b = insert_header(&backend, 2, block1_b, Default::default(), Default::default());
+		let block2_c = insert_header(&backend, 2, block1_b, Default::default(), [1; 32].into());
+
+		assert_eq!(backend.blockchain().leaves().unwrap(), vec![block2_a, block2_b, block2_c, block1_c]);
+
+		backend.finalize_block(BlockId::hash(block1_a), None).unwrap();
+		backend.finalize_block(BlockId::hash(block2_a), None).unwrap();
+
+		// leaves at same height stay. Leaves at lower heights pruned.
+		assert_eq!(backend.blockchain().leaves().unwrap(), vec![block2_a, block2_b, block2_c]);
+	}
+
 	#[test]
 	fn test_aux() {
 		let backend: Backend<test_client::runtime::Block> = Backend::new_test(0, 0);

substrate/core/client/db/src/light.rs

@@ -24,7 +24,8 @@ use kvdb::{KeyValueDB, DBTransaction};
 use client::backend::{AuxStore, NewBlockState};
 use client::blockchain::{BlockStatus, Cache as BlockchainCache,
 	HeaderBackend as BlockchainHeaderBackend, Info as BlockchainInfo};
-use client::{cht, LeafSet};
+use client::cht;
+use client::leaves::{LeafSet, FinalizationDisplaced};
 use client::error::{ErrorKind as ClientErrorKind, Result as ClientResult};
 use client::light::blockchain::Storage as LightBlockchainStorage;
 use parity_codec::{Decode, Encode};
@@ -250,6 +251,7 @@ impl<Block: BlockT> LightStorage<Block> {
 		transaction: &mut DBTransaction,
 		header: &Block::Header,
 		hash: Block::Hash,
+		displaced: &mut Option<FinalizationDisplaced<Block::Hash, NumberFor<Block>>>,
 	) -> ClientResult<()> {
 		let meta = self.meta.read();
 		if &meta.finalized_hash != header.parent_hash() {
@@ -311,6 +313,12 @@ impl<Block: BlockT> LightStorage<Block> {
 			}
 		}
 
+		let new_displaced = self.leaves.write().finalize_height(header.number().clone());
+		match displaced {
+			x @ &mut None => *x = Some(new_displaced),
+			&mut Some(ref mut displaced) => displaced.merge(new_displaced),
+		}
+
 		Ok(())
 	}
 
@@ -366,6 +374,7 @@ impl<Block> LightBlockchainStorage<Block> for LightStorage<Block>
 		leaf_state: NewBlockState,
 		aux_ops: Vec<(Vec<u8>, Option<Vec<u8>>)>,
 	) -> ClientResult<()> {
+		let mut finalization_displaced_leaves = None;
 		let mut transaction = DBTransaction::new();
 
 		let hash = header.hash();
@@ -405,7 +414,12 @@ impl<Block> LightBlockchainStorage<Block> for LightStorage<Block>
 		};
 
 		if finalized {
-			self.note_finalized(&mut transaction, &header, hash)?;
+			self.note_finalized(
+				&mut transaction,
+				&header,
+				hash,
+				&mut finalization_displaced_leaves,
+			)?;
 		}
 
 		{
@@ -425,10 +439,18 @@ impl<Block> LightBlockchainStorage<Block> for LightStorage<Block>
 			debug!("Light DB Commit {:?} ({})", hash, number);
 			let write_result = self.db.write(transaction).map_err(db_err);
 			if let Err(e) = write_result {
+				let mut leaves = self.leaves.write();
+				let mut undo = leaves.undo();
+
 				// revert leaves set update if there was one.
 				if let Some(displaced_leaf) = displaced_leaf {
-					leaves.undo(displaced_leaf);
+					undo.undo_import(displaced_leaf);
 				}
+
+				if let Some(finalization_displaced) = finalization_displaced_leaves {
+					undo.undo_finalization(finalization_displaced);
+				}
+
 				return Err(e);
 			}
 
@@ -464,10 +486,11 @@ impl<Block> LightBlockchainStorage<Block> for LightStorage<Block>
 
 	fn finalize_header(&self, id: BlockId<Block>) -> ClientResult<()> {
 		if let Some(header) = self.header(id)? {
+			let mut displaced = None;
 			let mut transaction = DBTransaction::new();
 			let hash = header.hash();
 			let number = *header.number();
-			self.note_finalized(&mut transaction, &header, hash.clone())?;
+			self.note_finalized(&mut transaction, &header, hash.clone(), &mut displaced)?;
 			{
 				let mut cache = self.cache.0.write();
 				let cache_ops = cache.transaction(&mut transaction)
@@ -477,7 +500,12 @@ impl<Block> LightBlockchainStorage<Block> for LightStorage<Block>
 					)?
 					.into_ops();
 
-				self.db.write(transaction).map_err(db_err)?;
+				if let Err(e) = self.db.write(transaction).map_err(db_err) {
+					if let Some(displaced) = displaced {
+						self.leaves.write().undo().undo_finalization(displaced);
+					}
+					return Err(e);
+				}
 				cache.commit(cache_ops);
 			}
 			self.update_meta(hash, header.number().clone(), false, true);
@@ -940,4 +968,28 @@ pub(crate) mod tests {
 		assert_eq!(db.get_aux(&[2]).unwrap(), Some(vec![102]));
 		assert_eq!(db.get_aux(&[3]).unwrap(), Some(vec![103]));
 	}
+
+	#[test]
+	fn test_leaves_pruned_on_finality() {
+		let db = LightStorage::<Block>::new_test();
+		let block0 = insert_final_block(&db, None, || default_header(&Default::default(), 0));
+
+		let block1_a = insert_block(&db, None, || default_header(&block0, 1));
+		let block1_b = insert_block(&db, None, || header_with_extrinsics_root(&block0, 1, [1; 32].into()));
+		let block1_c = insert_block(&db, None, || header_with_extrinsics_root(&block0, 1, [2; 32].into()));
+
+		assert_eq!(db.leaves.read().hashes(), vec![block1_a, block1_b, block1_c]);
+
+		let block2_a = insert_block(&db, None, || default_header(&block1_a, 2));
+		let block2_b = insert_block(&db, None, || header_with_extrinsics_root(&block1_b, 2, [1; 32].into()));
+		let block2_c = insert_block(&db, None, || header_with_extrinsics_root(&block1_b, 2, [2; 32].into()));
+
+		assert_eq!(db.leaves.read().hashes(), vec![block2_a, block2_b, block2_c, block1_c]);
+
+		db.finalize_header(BlockId::hash(block1_a)).unwrap();
+		db.finalize_header(BlockId::hash(block2_a)).unwrap();
+
+		// leaves at same height stay. Leaves at lower heights pruned.
+		assert_eq!(db.leaves.read().hashes(), vec![block2_a, block2_b, block2_c]);
+	}
 }

substrate/core/client/src/leaves.rs

@@ -14,68 +14,63 @@
 // You should have received a copy of the GNU General Public License
 // along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
 
-use std::collections::BTreeSet;
-use std::cmp::{Ord, Ordering};
+//! Helper for managing the set of available leaves in the chain for DB implementations.
+
+use std::collections::BTreeMap;
+use std::cmp::Reverse;
 use kvdb::{KeyValueDB, DBTransaction};
 use runtime_primitives::traits::SimpleArithmetic;
 use parity_codec::{Encode, Decode};
 use crate::error;
 
-/// helper wrapper type to keep a list of block hashes ordered
-/// by `number` descending in a `BTreeSet` which allows faster and simpler
-/// insertion and removal than keeping them in a list.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 struct LeafSetItem<H, N> {
 	hash: H,
-	number: N,
+	number: Reverse<N>,
 }
 
-impl<H, N> Ord for LeafSetItem<H, N> where N: Ord {
-	fn cmp(&self, other: &Self) -> Ordering {
-		// reverse (descending) order
-		other.number.cmp(&self.number)
-	}
-}
-
-impl<H, N> PartialOrd for LeafSetItem<H, N> where N: PartialOrd {
-	fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
-		// reverse (descending) order
-		other.number.partial_cmp(&self.number)
-	}
-}
-
-impl<H, N> PartialEq for LeafSetItem<H, N> where N: PartialEq {
-	fn eq(&self, other: &LeafSetItem<H, N>) -> bool {
-		self.number == other.number
-	}
-}
-
-impl<H, N> Eq for LeafSetItem<H, N> where N: PartialEq {}
-
 /// A displaced leaf after import.
-pub struct DisplacedLeaf<H, N> {
+#[must_use = "Displaced items from the leaf set must be handled."]
+pub struct ImportDisplaced<H, N> {
 	new_hash: H,
 	displaced: LeafSetItem<H, N>,
 }
 
+/// Displaced leaves after finalization.
+#[must_use = "Displaced items from the leaf set must be handled."]
+pub struct FinalizationDisplaced<H, N> {
+	leaves: BTreeMap<Reverse<N>, Vec<H>>,
+}
+
+impl<H, N: Ord> FinalizationDisplaced<H, N> {
+	/// Merge with another. This should only be used for displaced items that
+	/// are produced within one transaction of each other.
+	pub fn merge(&mut self, mut other: Self) {
+		// this will ignore keys that are in duplicate, however
+		// if these are actually produced correctly via the leaf-set within
+		// one transaction, then there will be no overlap in the keys.
+		self.leaves.append(&mut other.leaves);
+	}
+}
+
 /// list of leaf hashes ordered by number (descending).
 /// stored in memory for fast access.
 /// this allows very fast checking and modification of active leaves.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct LeafSet<H, N> {
-	storage: BTreeSet<LeafSetItem<H, N>>,
+	storage: BTreeMap<Reverse<N>, Vec<H>>,
 	pending_added: Vec<LeafSetItem<H, N>>,
 	pending_removed: Vec<H>,
 }
 
 impl<H, N> LeafSet<H, N> where
-	H: Clone + Decode + Encode,
-	N: Clone + SimpleArithmetic + Decode + Encode,
+	H: Clone + PartialEq + Decode + Encode,
+	N: std::fmt::Debug + Clone + SimpleArithmetic + Decode + Encode,
 {
 	/// Construct a new, blank leaf set.
 	pub fn new() -> Self {
 		Self {
-			storage: BTreeSet::new(),
+			storage: BTreeMap::new(),
 			pending_added: Vec::new(),
 			pending_removed: Vec::new(),
 		}
@@ -83,7 +78,8 @@ impl<H, N> LeafSet<H, N> where
 
 	/// Read the leaf list from the DB, using given prefix for keys.
 	pub fn read_from_db(db: &KeyValueDB, column: Option<u32>, prefix: &[u8]) -> error::Result<Self> {
-		let mut storage = BTreeSet::new();
+		let mut storage = BTreeMap::new();
+
 		for (key, value) in db.iter_from_prefix(column, prefix) {
 			if !key.starts_with(prefix) { break }
 			let raw_hash = &mut &key[prefix.len()..];
@@ -95,7 +91,7 @@ impl<H, N> LeafSet<H, N> where
 				Some(number) => number,
 				None => return Err(error::ErrorKind::Backend("Error decoding number".into()).into()),
 			};
-			storage.insert(LeafSetItem { hash, number });
+			storage.entry(Reverse(number)).or_insert_with(Vec::new).push(hash);
 		}
 		Ok(Self {
 			storage,
@@ -105,20 +101,20 @@ impl<H, N> LeafSet<H, N> where
 	}
 
 	/// update the leaf list on import. returns a displaced leaf if there was one.
-	pub fn import(&mut self, hash: H, number: N, parent_hash: H) -> Option<DisplacedLeaf<H, N>> {
+	pub fn import(&mut self, hash: H, number: N, parent_hash: H) -> Option<ImportDisplaced<H, N>> {
 		// avoid underflow for genesis.
 		let displaced = if number != N::zero() {
-			let displaced = LeafSetItem {
-				hash: parent_hash.clone(),
-				number: number.clone() - N::one(),
-			};
-			let was_displaced = self.storage.remove(&displaced);
+			let new_number = Reverse(number.clone() - N::one());
+			let was_displaced = self.remove_leaf(&new_number, &parent_hash);
 
 			if was_displaced {
-				self.pending_removed.push(parent_hash);
-				Some(DisplacedLeaf {
+				self.pending_removed.push(parent_hash.clone());
+				Some(ImportDisplaced {
 					new_hash: hash.clone(),
-					displaced,
+					displaced: LeafSetItem {
+						hash: parent_hash,
+						number: new_number,
+					},
 				})
 			} else {
 				None
@@ -127,36 +123,53 @@ impl<H, N> LeafSet<H, N> where
 			None
 		};
 
-		let item = LeafSetItem { hash, number };
-		self.storage.insert(item.clone());
-		self.pending_added.push(item);
+		self.insert_leaf(Reverse(number.clone()), hash.clone());
+		self.pending_added.push(LeafSetItem { hash, number: Reverse(number) });
 		displaced
 	}
 
-	/// Undo an import operation, with a displaced leaf.
-	pub fn undo(&mut self, displaced: DisplacedLeaf<H, N>) {
-		let new_number = displaced.displaced.number.clone() + N::one();
-		self.storage.remove(&LeafSetItem { hash: displaced.new_hash, number: new_number });
-		self.storage.insert(displaced.displaced);
-		self.pending_added.clear();
-		self.pending_removed.clear();
+	/// Note a block height finalized, displacing all leaves with number less than the finalized block's.
+	///
+	/// Although it would be more technically correct to also prune out leaves at the
+	/// same number as the finalized block, but with different hashes, the current behavior
+	/// is simpler and our assumptions about how finalization works means that those leaves
+	/// will be pruned soon afterwards anyway.
+	pub fn finalize_height(&mut self, number: N) -> FinalizationDisplaced<H, N> {
+		let boundary = if number == N::zero() {
+			return FinalizationDisplaced { leaves: BTreeMap::new() };
+		} else {
+			number - N::one()
+		};
+
+		let below_boundary = self.storage.split_off(&Reverse(boundary));
+		self.pending_removed.extend(below_boundary.values().flat_map(|h| h.iter()).cloned());
+		FinalizationDisplaced {
+			leaves: below_boundary,
+		}
+	}
+
+	/// Undo all pending operations.
+	///
+	/// This returns an `Undo` struct, where any
+	/// `Displaced` objects that have returned by previous method calls
+	/// should be passed to via the appropriate methods. Otherwise,
+	/// the on-disk state may get out of sync with in-memory state.
+	pub fn undo(&mut self) -> Undo<H, N> {
+		Undo { inner: self }
 	}
 
 	/// currently since revert only affects the canonical chain
 	/// we assume that parent has no further children
 	/// and we add it as leaf again
 	pub fn revert(&mut self, hash: H, number: N, parent_hash: H) {
-		self.storage.insert(LeafSetItem {
-			hash: parent_hash,
-			number: number.clone() - N::one(),
-		});
-		self.storage.remove(&LeafSetItem { hash, number });
+		self.insert_leaf(Reverse(number.clone() - N::one()), parent_hash);
+		self.remove_leaf(&Reverse(number), &hash);
 	}
 
 	/// returns an iterator over all hashes in the leaf set
 	/// ordered by their block number descending.
 	pub fn hashes(&self) -> Vec<H> {
-		self.storage.iter().map(|item| item.hash.clone()).collect()
+		self.storage.iter().flat_map(|(_, hashes)| hashes.iter()).cloned().collect()
 	}
 
 	/// Write the leaf list to the database transaction.
@@ -164,7 +177,7 @@ impl<H, N> LeafSet<H, N> where
 		let mut buf = prefix.to_vec();
 		for LeafSetItem { hash, number } in self.pending_added.drain(..) {
 			hash.using_encoded(|s| buf.extend(s));
-			tx.put_vec(column, &buf[..], number.encode());
+			tx.put_vec(column, &buf[..], number.0.encode());
 			buf.truncate(prefix.len()); // reuse allocation.
 		}
 		for hash in self.pending_removed.drain(..) {
@@ -173,6 +186,68 @@ impl<H, N> LeafSet<H, N> where
 			buf.truncate(prefix.len()); // reuse allocation.
 		}
 	}
+
+	#[cfg(test)]
+	fn contains(&self, number: N, hash: H) -> bool {
+		self.storage.get(&Reverse(number)).map_or(false, |hashes| hashes.contains(&hash))
+	}
+
+	fn insert_leaf(&mut self, number: Reverse<N>, hash: H) {
+		self.storage.entry(number).or_insert_with(Vec::new).push(hash);
+	}
+
+	// returns true if this leaf was contained, false otherwise.
+	fn remove_leaf(&mut self, number: &Reverse<N>, hash: &H) -> bool {
+		let mut empty = false;
+		let removed = self.storage.get_mut(number).map_or(false, |leaves| {
+			let mut found = false;
+			leaves.retain(|h| if h == hash {
+				found = true;
+				false
+			} else {
+				true
+			});
+
+			if leaves.is_empty() { empty = true }
+
+			found
+		});
+
+		if removed && empty {
+			self.storage.remove(number);
+		}
+
+		removed
+	}
+}
+
+/// Helper for undoing operations.
+pub struct Undo<'a, H: 'a, N: 'a> {
+	inner: &'a mut LeafSet<H, N>,
+}
+
+impl<'a, H: 'a, N: 'a> Undo<'a, H, N> where
+	H: Clone + PartialEq + Decode + Encode,
+	N: std::fmt::Debug + Clone + SimpleArithmetic + Decode + Encode,
+{
+	/// Undo an imported block by providing the displaced leaf.
+	pub fn undo_import(&mut self, displaced: ImportDisplaced<H, N>) {
+		let new_number = Reverse(displaced.displaced.number.0.clone() + N::one());
+		self.inner.remove_leaf(&new_number, &displaced.new_hash);
+		self.inner.insert_leaf(new_number, displaced.displaced.hash);
+	}
+
+	/// Undo a finalization operation by providing the displaced leaves.
+	pub fn undo_finalization(&mut self, mut displaced: FinalizationDisplaced<H, N>) {
+		self.inner.storage.append(&mut displaced.leaves);
+	}
+}
+
+impl<'a, H: 'a, N: 'a> Drop for Undo<'a, H, N> {
+	fn drop(&mut self) {
+		self.inner.pending_added.clear();
+		self.inner.pending_removed.clear();
+	}
 }
 
 #[cfg(test)]
@@ -188,15 +263,15 @@ mod tests {
 		set.import(2_1, 2, 1_1);
 		set.import(3_1, 3, 2_1);
 
-		assert!(set.storage.contains(&LeafSetItem { hash: 3_1, number: 3 }));
-		assert!(!set.storage.contains(&LeafSetItem { hash: 2_1, number: 2 }));
-		assert!(!set.storage.contains(&LeafSetItem { hash: 1_1, number: 1 }));
-		assert!(!set.storage.contains(&LeafSetItem { hash: 0, number: 0 }));
+		assert!(set.contains(3, 3_1));
+		assert!(!set.contains(2, 2_1));
+		assert!(!set.contains(1, 1_1));
+		assert!(!set.contains(0, 0));
 
 		set.import(2_2, 2, 1_1);
 
-		assert!(set.storage.contains(&LeafSetItem { hash: 3_1, number: 3 }));
-		assert!(set.storage.contains(&LeafSetItem { hash: 2_2, number: 2 }));
+		assert!(set.contains(3, 3_1));
+		assert!(set.contains(2, 2_2));
 	}
 
 	#[test]
@@ -219,4 +294,63 @@ mod tests {
 		let set2 = LeafSet::read_from_db(&db, None, PREFIX).unwrap();
 		assert_eq!(set, set2);
 	}
+
+	#[test]
+	fn two_leaves_same_height_can_be_included() {
+		let mut set = LeafSet::new();
+
+		set.import(1_1u32, 10u32,0u32);
+		set.import(1_2, 10, 0);
+
+		assert!(set.storage.contains_key(&Reverse(10)));
+		assert!(set.contains(10, 1_1));
+		assert!(set.contains(10, 1_2));
+		assert!(!set.contains(10, 1_3));
+	}
+
+	#[test]
+	fn finalization_consistent_with_disk() {
+		const PREFIX: &[u8] = b"prefix";
+		let db = ::kvdb_memorydb::create(0);
+
+		let mut set = LeafSet::new();
+		set.import(10_1u32, 10u32, 0u32);
+		set.import(11_1, 11, 10_2);
+		set.import(11_2, 11, 10_2);
+		set.import(12_1, 12, 11_123);
+
+		assert!(set.contains(10, 10_1));
+
+		let mut tx = DBTransaction::new();
+		set.prepare_transaction(&mut tx, None, PREFIX);
+		db.write(tx).unwrap();
+
+		let _ = set.finalize_height(11);
+		let mut tx = DBTransaction::new();
+		set.prepare_transaction(&mut tx, None, PREFIX);
+		db.write(tx).unwrap();
+
+		assert!(set.contains(11, 11_1));
+		assert!(set.contains(11, 11_2));
+		assert!(set.contains(12, 12_1));
+		assert!(!set.contains(10, 10_1));
+
+		let set2 = LeafSet::read_from_db(&db, None, PREFIX).unwrap();
+		assert_eq!(set, set2);
+	}
+
+	#[test]
+	fn undo_finalization() {
+		let mut set = LeafSet::new();
+		set.import(10_1u32, 10u32, 0u32);
+		set.import(11_1, 11, 10_2);
+		set.import(11_2, 11, 10_2);
+		set.import(12_1, 12, 11_123);
+
+		let displaced = set.finalize_height(11);
+		assert!(!set.contains(10, 10_1));
+
+		set.undo().undo_finalization(displaced);
+		assert!(set.contains(10, 10_1));
+	}
 }

substrate/core/client/src/lib.rs

@@ -38,9 +38,9 @@ pub mod block_builder;
 #[cfg(feature = "std")]
 pub mod light;
 #[cfg(feature = "std")]
-pub mod children;
+pub mod leaves;
 #[cfg(feature = "std")]
-mod leaves;
+pub mod children;
 #[cfg(feature = "std")]
 mod call_executor;
 #[cfg(feature = "std")]