chain-selection: allow reverting current block (#4103)

Block reversion of the current block is technically possible as can be
seen from

https://github.com/paritytech/polkadot-sdk/blob/39b1f50f1c251def87c1625d68567ed252dc6272/polkadot/runtime/parachains/src/disputes.rs#L1215-L1223

- [x] Fix the test

polkadot/node/core/chain-selection/src/tree.rs

@@ -236,7 +236,7 @@ fn propagate_viability_update(
 	Ok(())
 }
 
-/// Imports a new block and applies any reversions to ancestors.
+/// Imports a new block and applies any reversions to ancestors or the block itself.
 pub(crate) fn import_block(
 	backend: &mut OverlayedBackend<impl Backend>,
 	block_hash: Hash,
@@ -246,25 +246,29 @@ pub(crate) fn import_block(
 	weight: BlockWeight,
 	stagnant_at: Timestamp,
 ) -> Result<(), Error> {
-	add_block(backend, block_hash, block_number, parent_hash, weight, stagnant_at)?;
-	apply_ancestor_reversions(backend, block_hash, block_number, reversion_logs)?;
+	let block_entry =
+		add_block(backend, block_hash, block_number, parent_hash, weight, stagnant_at)?;
+	apply_reversions(backend, block_entry, reversion_logs)?;
 
 	Ok(())
 }
 
 // Load the given ancestor's block entry, in descending order from the `block_hash`.
-// The ancestor_number must be at least one block less than the `block_number`.
+// The ancestor_number must be not higher than the `block_entry`'s.
 //
 // The returned entry will be `None` if the range is invalid or any block in the path had
 // no entry present. If any block entry was missing, it can safely be assumed to
 // be finalized.
 fn load_ancestor(
 	backend: &mut OverlayedBackend<impl Backend>,
-	block_hash: Hash,
-	block_number: BlockNumber,
+	block_entry: &BlockEntry,
 	ancestor_number: BlockNumber,
 ) -> Result<Option<BlockEntry>, Error> {
-	if block_number <= ancestor_number {
+	let block_hash = block_entry.block_hash;
+	let block_number = block_entry.block_number;
+	if block_number == ancestor_number {
+		return Ok(Some(block_entry.clone()))
+	} else if block_number < ancestor_number {
 		return Ok(None)
 	}
 
@@ -300,7 +304,7 @@ fn add_block(
 	parent_hash: Hash,
 	weight: BlockWeight,
 	stagnant_at: Timestamp,
-) -> Result<(), Error> {
+) -> Result<BlockEntry, Error> {
 	let mut leaves = backend.load_leaves()?;
 	let parent_entry = backend.load_block_entry(&parent_hash)?;
 
@@ -308,7 +312,7 @@ fn add_block(
 		parent_entry.as_ref().and_then(|parent| parent.non_viable_ancestor_for_child());
 
 	// 1. Add the block to the DB assuming it's not reverted.
-	backend.write_block_entry(BlockEntry {
+	let block_entry = BlockEntry {
 		block_hash,
 		block_number,
 		parent_hash,
@@ -319,7 +323,8 @@ fn add_block(
 			approval: Approval::Unapproved,
 		},
 		weight,
-	});
+	};
+	backend.write_block_entry(block_entry.clone());
 
 	// 2. Update leaves if inherited viability is fine.
 	if inherited_viability.is_none() {
@@ -344,26 +349,25 @@ fn add_block(
 	stagnant_at_list.push(block_hash);
 	backend.write_stagnant_at(stagnant_at, stagnant_at_list);
 
-	Ok(())
+	Ok(block_entry)
 }
 
 /// Assuming that a block is already imported, accepts the number of the block
 /// as well as a list of reversions triggered by the block in ascending order.
-fn apply_ancestor_reversions(
+fn apply_reversions(
 	backend: &mut OverlayedBackend<impl Backend>,
-	block_hash: Hash,
-	block_number: BlockNumber,
+	block_entry: BlockEntry,
 	reversions: Vec<BlockNumber>,
 ) -> Result<(), Error> {
 	// Note: since revert numbers are  in ascending order, the expensive propagation
 	// of unviability is only heavy on the first log.
 	for revert_number in reversions {
-		let maybe_block_entry = load_ancestor(backend, block_hash, block_number, revert_number)?;
-		if let Some(block_entry) = &maybe_block_entry {
+		let maybe_block_entry = load_ancestor(backend, &block_entry, revert_number)?;
+		if let Some(entry) = &maybe_block_entry {
 			gum::trace!(
 				target: LOG_TARGET,
 				?revert_number,
-				revert_hash = ?block_entry.block_hash,
+				revert_hash = ?entry.block_hash,
 				"Block marked as reverted via scraped on-chain reversions"
 			);
 		}
@@ -372,8 +376,8 @@ fn apply_ancestor_reversions(
 			maybe_block_entry,
 			None,
 			revert_number,
-			Some(block_hash),
-			Some(block_number),
+			Some(block_entry.block_hash),
+			Some(block_entry.block_number),
 		)?;
 	}