client: allow reverting finalized blocks (#4535)

* client: allow reverting blocks past finality

* client: fix leaves reversion

* client: extend docs on revert

* client: add comment on leaves revert

substrate/client/src/client.rs

@@ -1142,10 +1142,24 @@ impl<B, E, Block, RA> Client<B, E, Block, RA> where
 		Ok(())
 	}
 
-	/// Attempts to revert the chain by `n` blocks. Returns the number of blocks that were
-	/// successfully reverted.
+	/// Attempts to revert the chain by `n` blocks guaranteeing that no block is
+	/// reverted past the last finalized block. Returns the number of blocks
+	/// that were successfully reverted.
 	pub fn revert(&self, n: NumberFor<Block>) -> sp_blockchain::Result<NumberFor<Block>> {
-		Ok(self.backend.revert(n)?)
+		Ok(self.backend.revert(n, false)?)
+	}
+
+	/// Attempts to revert the chain by `n` blocks disregarding finality. This
+	/// method will revert any finalized blocks as requested and can potentially
+	/// lead the node in an inconsistent state. Other modules in the system that
+	/// persist data and that rely on finality (e.g. consensus parts) will be
+	/// unaffected by the revert. Use this method with caution and making sure
+	/// that no other data needs to be reverted for consistency aside from the
+	/// block data.
+	///
+	/// Returns the number of blocks that were successfully reverted.
+	pub fn unsafe_revert(&self, n: NumberFor<Block>) -> sp_blockchain::Result<NumberFor<Block>> {
+		Ok(self.backend.revert(n, true)?)
 	}
 
 	/// Get usage info about current client.

substrate/client/src/in_mem.rs

@@ -707,7 +707,11 @@ where
 		}
 	}
 
-	fn revert(&self, _n: NumberFor<Block>) -> sp_blockchain::Result<NumberFor<Block>> {
+	fn revert(
+		&self,
+		_n: NumberFor<Block>,
+		_revert_finalized: bool,
+	) -> sp_blockchain::Result<NumberFor<Block>> {
 		Ok(Zero::zero())
 	}
 

substrate/client/src/leaves.rs

@@ -158,12 +158,35 @@ impl<H, N> LeafSet<H, N> where
 		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.insert_leaf(Reverse(number.clone() - N::one()), parent_hash);
-		self.remove_leaf(&Reverse(number), &hash);
+	/// Revert to the given block height by dropping all leaves in the leaf set
+	/// with a block number higher than the target.
+	pub fn revert(&mut self, best_hash: H, best_number: N) {
+		let items = self.storage.iter()
+			.flat_map(|(number, hashes)| hashes.iter().map(move |h| (h.clone(), number.clone())))
+			.collect::<Vec<_>>();
+
+		for (hash, number) in &items {
+			if number.0 > best_number {
+				assert!(
+					self.remove_leaf(number, hash),
+					"item comes from an iterator over storage; qed",
+				);
+
+				self.pending_removed.push(hash.clone());
+			}
+		}
+
+		let best_number = Reverse(best_number);
+		let leaves_contains_best = self.storage
+			.get(&best_number)
+			.map_or(false, |hashes| hashes.contains(&best_hash));
+
+		// we need to make sure that the best block exists in the leaf set as
+		// this is an invariant of regular block import.
+		if !leaves_contains_best {
+			self.insert_leaf(best_number.clone(), best_hash.clone());
+			self.pending_added.push(LeafSetItem { hash: best_hash, number: best_number });
+		}
 	}
 
 	/// returns an iterator over all hashes in the leaf set

substrate/client/src/light/backend.rs

@@ -213,7 +213,11 @@ impl<S, Block, H> ClientBackend<Block, H> for Backend<S, H> where
 		Ok(GenesisOrUnavailableState::Unavailable)
 	}
 
-	fn revert(&self, _n: NumberFor<Block>) -> ClientResult<NumberFor<Block>> {
+	fn revert(
+		&self,
+		_n: NumberFor<Block>,
+		_revert_finalized: bool,
+	) -> ClientResult<NumberFor<Block>> {
 		Err(ClientError::NotAvailableOnLightClient)
 	}