diff --git a/substrate/client/api/src/leaves.rs b/substrate/client/api/src/leaves.rs
index 26eda46e6f..cdcb80a110 100644
--- a/substrate/client/api/src/leaves.rs
+++ b/substrate/client/api/src/leaves.rs
@@ -32,32 +32,36 @@ struct LeafSetItem<H, N> {
 	number: Reverse<N>,
 }
 
-/// A displaced leaf after import.
-#[must_use = "Displaced items from the leaf set must be handled."]
-pub struct ImportDisplaced<H, N> {
-	new_hash: H,
-	displaced: LeafSetItem<H, N>,
+/// Inserted and removed leaves after an import action.
+pub struct ImportOutcome<H, N> {
+	inserted: LeafSetItem<H, N>,
+	removed: Option<H>,
 }
 
-/// 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>>,
+/// Inserted and removed leaves after a remove action.
+pub struct RemoveOutcome<H, N> {
+	inserted: Option<H>,
+	removed: LeafSetItem<H, N>,
 }
 
-impl<H, N: Ord> FinalizationDisplaced<H, N> {
+/// Removed leaves after a finalization action.
+pub struct FinalizationOutcome<H, N> {
+	removed: BTreeMap<Reverse<N>, Vec<H>>,
+}
+
+impl<H, N: Ord> FinalizationOutcome<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);
+		self.removed.append(&mut other.removed);
 	}
 
 	/// Iterate over all displaced leaves.
 	pub fn leaves(&self) -> impl Iterator<Item = &H> {
-		self.leaves.values().flatten()
+		self.removed.values().flatten()
 	}
 }
 
@@ -99,27 +103,52 @@ 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<ImportDisplaced<H, N>> {
-		// avoid underflow for genesis.
-		let displaced = if number != N::zero() {
-			let parent_number = Reverse(number.clone() - N::one());
-			let was_displaced = self.remove_leaf(&parent_number, &parent_hash);
+	pub fn import(&mut self, hash: H, number: N, parent_hash: H) -> ImportOutcome<H, N> {
+		let number = Reverse(number);
 
-			if was_displaced {
-				Some(ImportDisplaced {
-					new_hash: hash.clone(),
-					displaced: LeafSetItem { hash: parent_hash, number: parent_number },
-				})
-			} else {
-				None
-			}
+		let removed = if number.0 != N::zero() {
+			let parent_number = Reverse(number.0.clone() - N::one());
+			self.remove_leaf(&parent_number, &parent_hash).then(|| parent_hash)
 		} else {
 			None
 		};
 
-		self.insert_leaf(Reverse(number.clone()), hash.clone());
-		displaced
+		self.insert_leaf(number.clone(), hash.clone());
+
+		ImportOutcome { inserted: LeafSetItem { hash, number }, removed }
+	}
+
+	/// Update the leaf list on removal.
+	///
+	/// Note that the leaves set structure doesn't have the information to decide if the
+	/// leaf we're removing is the last children of the parent. Follows that this method requires
+	/// the caller to check this condition and optionally pass the `parent_hash` if `hash` is
+	/// its last child.
+	///
+	/// Returns `None` if no modifications are applied.
+	pub fn remove(
+		&mut self,
+		hash: H,
+		number: N,
+		parent_hash: Option<H>,
+	) -> Option<RemoveOutcome<H, N>> {
+		let number = Reverse(number);
+
+		if !self.remove_leaf(&number, &hash) {
+			return None
+		}
+
+		let inserted = parent_hash.and_then(|parent_hash| {
+			if number.0 != N::zero() {
+				let parent_number = Reverse(number.0.clone() - N::one());
+				self.insert_leaf(parent_number, parent_hash.clone());
+				Some(parent_hash)
+			} else {
+				None
+			}
+		});
+
+		Some(RemoveOutcome { inserted, removed: LeafSetItem { hash, number } })
 	}
 
 	/// Note a block height finalized, displacing all leaves with number less than the finalized
@@ -129,15 +158,15 @@ where
 	/// 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> {
+	pub fn finalize_height(&mut self, number: N) -> FinalizationOutcome<H, N> {
 		let boundary = if number == N::zero() {
-			return FinalizationDisplaced { leaves: BTreeMap::new() }
+			return FinalizationOutcome { removed: BTreeMap::new() }
 		} else {
 			number - N::one()
 		};
 
 		let below_boundary = self.storage.split_off(&Reverse(boundary));
-		FinalizationDisplaced { leaves: below_boundary }
+		FinalizationOutcome { removed: below_boundary }
 	}
 
 	/// The same as [`Self::finalize_height`], but it only simulates the operation.
@@ -145,16 +174,16 @@ where
 	/// This means that no changes are done.
 	///
 	/// Returns the leaves that would be displaced by finalizing the given block.
-	pub fn displaced_by_finalize_height(&self, number: N) -> FinalizationDisplaced<H, N> {
+	pub fn displaced_by_finalize_height(&self, number: N) -> FinalizationOutcome<H, N> {
 		let boundary = if number == N::zero() {
-			return FinalizationDisplaced { leaves: BTreeMap::new() }
+			return FinalizationOutcome { removed: BTreeMap::new() }
 		} else {
 			number - N::one()
 		};
 
 		let below_boundary = self.storage.range(&Reverse(boundary)..);
-		FinalizationDisplaced {
-			leaves: below_boundary.map(|(k, v)| (k.clone(), v.clone())).collect(),
+		FinalizationOutcome {
+			removed: below_boundary.map(|(k, v)| (k.clone(), v.clone())).collect(),
 		}
 	}
 
@@ -276,16 +305,30 @@ where
 	H: Clone + PartialEq + Decode + Encode,
 	N: std::fmt::Debug + Clone + AtLeast32Bit + 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(displaced.displaced.number, displaced.displaced.hash);
+	/// Undo an imported block by providing the import operation outcome.
+	/// No additional operations should be performed between import and undo.
+	pub fn undo_import(&mut self, outcome: ImportOutcome<H, N>) {
+		if let Some(removed_hash) = outcome.removed {
+			let removed_number = Reverse(outcome.inserted.number.0.clone() - N::one());
+			self.inner.insert_leaf(removed_number, removed_hash);
+		}
+		self.inner.remove_leaf(&outcome.inserted.number, &outcome.inserted.hash);
+	}
+
+	/// Undo a removed block by providing the displaced leaf.
+	/// No additional operations should be performed between remove and undo.
+	pub fn undo_remove(&mut self, outcome: RemoveOutcome<H, N>) {
+		if let Some(inserted_hash) = outcome.inserted {
+			let inserted_number = Reverse(outcome.removed.number.0.clone() - N::one());
+			self.inner.remove_leaf(&inserted_number, &inserted_hash);
+		}
+		self.inner.insert_leaf(outcome.removed.number, outcome.removed.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);
+	/// No additional operations should be performed between finalization and undo.
+	pub fn undo_finalization(&mut self, mut outcome: FinalizationOutcome<H, N>) {
+		self.inner.storage.append(&mut outcome.removed);
 	}
 }
 
@@ -295,7 +338,7 @@ mod tests {
 	use std::sync::Arc;
 
 	#[test]
-	fn it_works() {
+	fn import_works() {
 		let mut set = LeafSet::new();
 		set.import(0u32, 0u32, 0u32);
 
@@ -317,6 +360,90 @@ mod tests {
 		assert!(set.contains(3, 3_1));
 		assert!(set.contains(2, 2_2));
 		assert!(set.contains(2, 2_3));
+
+		// Finally test the undo feature
+
+		let outcome = set.import(2_4, 2, 1_1);
+		assert_eq!(outcome.inserted.hash, 2_4);
+		assert_eq!(outcome.removed, None);
+		assert_eq!(set.count(), 4);
+		assert!(set.contains(2, 2_4));
+
+		set.undo().undo_import(outcome);
+		assert_eq!(set.count(), 3);
+		assert!(set.contains(3, 3_1));
+		assert!(set.contains(2, 2_2));
+		assert!(set.contains(2, 2_3));
+
+		let outcome = set.import(3_2, 3, 2_3);
+		assert_eq!(outcome.inserted.hash, 3_2);
+		assert_eq!(outcome.removed, Some(2_3));
+		assert_eq!(set.count(), 3);
+		assert!(set.contains(3, 3_2));
+
+		set.undo().undo_import(outcome);
+		assert_eq!(set.count(), 3);
+		assert!(set.contains(3, 3_1));
+		assert!(set.contains(2, 2_2));
+		assert!(set.contains(2, 2_3));
+	}
+
+	#[test]
+	fn removal_works() {
+		let mut set = LeafSet::new();
+		set.import(10_1u32, 10u32, 0u32);
+		set.import(11_1, 11, 10_1);
+		set.import(11_2, 11, 10_1);
+		set.import(12_1, 12, 11_1);
+
+		let outcome = set.remove(12_1, 12, Some(11_1)).unwrap();
+		assert_eq!(outcome.removed.hash, 12_1);
+		assert_eq!(outcome.inserted, Some(11_1));
+		assert_eq!(set.count(), 2);
+		assert!(set.contains(11, 11_1));
+		assert!(set.contains(11, 11_2));
+
+		let outcome = set.remove(11_1, 11, None).unwrap();
+		assert_eq!(outcome.removed.hash, 11_1);
+		assert_eq!(outcome.inserted, None);
+		assert_eq!(set.count(), 1);
+		assert!(set.contains(11, 11_2));
+
+		let outcome = set.remove(11_2, 11, Some(10_1)).unwrap();
+		assert_eq!(outcome.removed.hash, 11_2);
+		assert_eq!(outcome.inserted, Some(10_1));
+		assert_eq!(set.count(), 1);
+		assert!(set.contains(10, 10_1));
+
+		set.undo().undo_remove(outcome);
+		assert_eq!(set.count(), 1);
+		assert!(set.contains(11, 11_2));
+	}
+
+	#[test]
+	fn finalization_works() {
+		let mut set = LeafSet::new();
+		set.import(9_1u32, 9u32, 0u32);
+		set.import(10_1, 10, 9_1);
+		set.import(10_2, 10, 9_1);
+		set.import(11_1, 11, 10_1);
+		set.import(11_2, 11, 10_1);
+		set.import(12_1, 12, 11_2);
+
+		let outcome = set.finalize_height(11);
+		assert_eq!(set.count(), 2);
+		assert!(set.contains(11, 11_1));
+		assert!(set.contains(12, 12_1));
+		assert_eq!(
+			outcome.removed,
+			[(Reverse(10), vec![10_2])].into_iter().collect::<BTreeMap<_, _>>(),
+		);
+
+		set.undo().undo_finalization(outcome);
+		assert_eq!(set.count(), 3);
+		assert!(set.contains(11, 11_1));
+		assert!(set.contains(12, 12_1));
+		assert!(set.contains(10, 10_2));
 	}
 
 	#[test]
@@ -383,44 +510,4 @@ mod tests {
 		let set2 = LeafSet::read_from_db(&*db, 0, PREFIX).unwrap();
 		assert_eq!(set, set2);
 	}
-
-	#[test]
-	fn undo_import() {
-		let mut set = LeafSet::new();
-		set.import(10_1u32, 10u32, 0u32);
-		set.import(11_1, 11, 10_1);
-		set.import(11_2, 11, 10_1);
-
-		let displaced = set.import(12_1, 12, 11_1).unwrap();
-		assert_eq!(set.count(), 2);
-		assert!(set.contains(11, 11_2));
-		assert!(set.contains(12, 12_1));
-
-		set.undo().undo_import(displaced);
-		assert_eq!(set.count(), 2);
-		assert!(set.contains(11, 11_1));
-		assert!(set.contains(11, 11_2));
-	}
-
-	#[test]
-	fn undo_finalization() {
-		let mut set = LeafSet::new();
-		set.import(9_1u32, 9u32, 0u32);
-		set.import(10_1, 10, 9_1);
-		set.import(10_2, 10, 9_1);
-		set.import(11_1, 11, 10_1);
-		set.import(11_2, 11, 10_1);
-		set.import(12_1, 12, 11_2);
-
-		let displaced = set.finalize_height(11);
-		assert_eq!(set.count(), 2);
-		assert!(set.contains(11, 11_1));
-		assert!(set.contains(12, 12_1));
-
-		set.undo().undo_finalization(displaced);
-		assert_eq!(set.count(), 3);
-		assert!(set.contains(11, 11_1));
-		assert!(set.contains(12, 12_1));
-		assert!(set.contains(10, 10_2));
-	}
 }
diff --git a/substrate/client/db/src/lib.rs b/substrate/client/db/src/lib.rs
index 6b644ad53d..71f291e66e 100644
--- a/substrate/client/db/src/lib.rs
+++ b/substrate/client/db/src/lib.rs
@@ -59,7 +59,7 @@ use codec::{Decode, Encode};
 use hash_db::Prefix;
 use sc_client_api::{
 	backend::NewBlockState,
-	leaves::{FinalizationDisplaced, LeafSet},
+	leaves::{FinalizationOutcome, LeafSet},
 	utils::is_descendent_of,
 	IoInfo, MemoryInfo, MemorySize, UsageInfo,
 };
@@ -1251,7 +1251,7 @@ impl<Block: BlockT> Backend<Block> {
 		header: &Block::Header,
 		last_finalized: Option<Block::Hash>,
 		justification: Option<Justification>,
-		finalization_displaced: &mut Option<FinalizationDisplaced<Block::Hash, NumberFor<Block>>>,
+		finalization_displaced: &mut Option<FinalizationOutcome<Block::Hash, NumberFor<Block>>>,
 	) -> ClientResult<MetaUpdate<Block>> {
 		// TODO: ensure best chain contains this block.
 		let number = *header.number();
@@ -1657,7 +1657,7 @@ impl<Block: BlockT> Backend<Block> {
 		transaction: &mut Transaction<DbHash>,
 		f_header: &Block::Header,
 		f_hash: Block::Hash,
-		displaced: &mut Option<FinalizationDisplaced<Block::Hash, NumberFor<Block>>>,
+		displaced: &mut Option<FinalizationOutcome<Block::Hash, NumberFor<Block>>>,
 		with_state: bool,
 	) -> ClientResult<()> {
 		let f_num = *f_header.number();
@@ -1696,7 +1696,7 @@ impl<Block: BlockT> Backend<Block> {
 		&self,
 		transaction: &mut Transaction<DbHash>,
 		finalized: NumberFor<Block>,
-		displaced: &FinalizationDisplaced<Block::Hash, NumberFor<Block>>,
+		displaced: &FinalizationOutcome<Block::Hash, NumberFor<Block>>,
 	) -> ClientResult<()> {
 		if let BlocksPruning::Some(blocks_pruning) = self.blocks_pruning {
 			// Always keep the last finalized block
@@ -2226,9 +2226,40 @@ impl<Block: BlockT> sc_client_api::backend::Backend<Block> for Backend<Block> {
 			apply_state_commit(&mut transaction, commit);
 		}
 		transaction.remove(columns::KEY_LOOKUP, hash.as_ref());
-		leaves.revert(*hash, hdr.number);
+
+		let children: Vec<_> = self
+			.blockchain()
+			.children(hdr.parent)?
+			.into_iter()
+			.filter(|child_hash| child_hash != hash)
+			.collect();
+		let parent_leaf = if children.is_empty() {
+			children::remove_children(
+				&mut transaction,
+				columns::META,
+				meta_keys::CHILDREN_PREFIX,
+				hdr.parent,
+			);
+			Some(hdr.parent)
+		} else {
+			children::write_children(
+				&mut transaction,
+				columns::META,
+				meta_keys::CHILDREN_PREFIX,
+				hdr.parent,
+				children,
+			);
+			None
+		};
+
+		let remove_outcome = leaves.remove(*hash, hdr.number, parent_leaf);
 		leaves.prepare_transaction(&mut transaction, columns::META, meta_keys::LEAF_PREFIX);
-		self.storage.db.commit(transaction)?;
+		if let Err(e) = self.storage.db.commit(transaction) {
+			if let Some(outcome) = remove_outcome {
+				leaves.undo().undo_remove(outcome);
+			}
+			return Err(e.into())
+		}
 		self.blockchain().remove_header_metadata(*hash);
 		Ok(())
 	}
@@ -3376,7 +3407,21 @@ pub(crate) mod tests {
 			prev_hash = hash;
 		}
 
-		// insert a fork at block 2, which becomes best block
+		for i in 0..2 {
+			let hash = insert_block(
+				&backend,
+				2,
+				blocks[1],
+				None,
+				sp_core::H256::random(),
+				vec![i.into()],
+				None,
+			)
+			.unwrap();
+			blocks.push(hash);
+		}
+
+		// insert a fork at block 1, which becomes best block
 		let best_hash = insert_block(
 			&backend,
 			1,
@@ -3387,11 +3432,36 @@ pub(crate) mod tests {
 			None,
 		)
 		.unwrap();
+
+		assert_eq!(backend.blockchain().info().best_hash, best_hash);
 		assert!(backend.remove_leaf_block(&best_hash).is_err());
-		assert!(backend.have_state_at(&prev_hash, 1));
-		backend.remove_leaf_block(&prev_hash).unwrap();
-		assert_eq!(None, backend.blockchain().header(BlockId::hash(prev_hash)).unwrap());
-		assert!(!backend.have_state_at(&prev_hash, 1));
+
+		assert_eq!(backend.blockchain().leaves().unwrap(), vec![blocks[2], blocks[3], best_hash]);
+		assert_eq!(backend.blockchain().children(blocks[1]).unwrap(), vec![blocks[2], blocks[3]]);
+
+		assert!(backend.have_state_at(&blocks[3], 2));
+		assert!(backend.blockchain().header(BlockId::hash(blocks[3])).unwrap().is_some());
+		backend.remove_leaf_block(&blocks[3]).unwrap();
+		assert!(!backend.have_state_at(&blocks[3], 2));
+		assert!(backend.blockchain().header(BlockId::hash(blocks[3])).unwrap().is_none());
+		assert_eq!(backend.blockchain().leaves().unwrap(), vec![blocks[2], best_hash]);
+		assert_eq!(backend.blockchain().children(blocks[1]).unwrap(), vec![blocks[2]]);
+
+		assert!(backend.have_state_at(&blocks[2], 2));
+		assert!(backend.blockchain().header(BlockId::hash(blocks[2])).unwrap().is_some());
+		backend.remove_leaf_block(&blocks[2]).unwrap();
+		assert!(!backend.have_state_at(&blocks[2], 2));
+		assert!(backend.blockchain().header(BlockId::hash(blocks[2])).unwrap().is_none());
+		assert_eq!(backend.blockchain().leaves().unwrap(), vec![best_hash, blocks[1]]);
+		assert_eq!(backend.blockchain().children(blocks[1]).unwrap(), vec![]);
+
+		assert!(backend.have_state_at(&blocks[1], 1));
+		assert!(backend.blockchain().header(BlockId::hash(blocks[1])).unwrap().is_some());
+		backend.remove_leaf_block(&blocks[1]).unwrap();
+		assert!(!backend.have_state_at(&blocks[1], 1));
+		assert!(backend.blockchain().header(BlockId::hash(blocks[1])).unwrap().is_none());
+		assert_eq!(backend.blockchain().leaves().unwrap(), vec![best_hash]);
+		assert_eq!(backend.blockchain().children(blocks[0]).unwrap(), vec![best_hash]);
 	}
 
 	#[test]