diff --git a/prdoc/pr_3962.prdoc b/prdoc/pr_3962.prdoc
new file mode 100644
index 000000000000..7ef59d38ce5c
--- /dev/null
+++ b/prdoc/pr_3962.prdoc
@@ -0,0 +1,12 @@
+title: Change fork calculation algorithm.
+
+doc:
+  - audience: Node Dev
+    description: |
+      This PR changes the fork calculation and pruning algorithm to enable future block header pruning.
+      During the finalization of the block we prune known stale forks, so forks are pruned faster.
+
+crates:
+  - name: sc-client-api
+  - name: sc-client-db
+  - name: sp-blockchain
diff --git a/substrate/client/api/src/in_mem.rs b/substrate/client/api/src/in_mem.rs
index b933ed1f17e0..ba89aede9147 100644
--- a/substrate/client/api/src/in_mem.rs
+++ b/substrate/client/api/src/in_mem.rs
@@ -419,20 +419,6 @@ impl<Block: BlockT> blockchain::Backend<Block> for Blockchain<Block> {
 		Ok(self.storage.read().leaves.hashes())
 	}
 
-	fn displaced_leaves_after_finalizing(
-		&self,
-		block_number: NumberFor<Block>,
-	) -> sp_blockchain::Result<Vec<Block::Hash>> {
-		Ok(self
-			.storage
-			.read()
-			.leaves
-			.displaced_by_finalize_height(block_number)
-			.leaves()
-			.cloned()
-			.collect::<Vec<_>>())
-	}
-
 	fn children(&self, _parent_hash: Block::Hash) -> sp_blockchain::Result<Vec<Block::Hash>> {
 		unimplemented!()
 	}
diff --git a/substrate/client/api/src/leaves.rs b/substrate/client/api/src/leaves.rs
index a8a988771e2f..e129de8bf3fa 100644
--- a/substrate/client/api/src/leaves.rs
+++ b/substrate/client/api/src/leaves.rs
@@ -49,7 +49,7 @@ pub struct FinalizationOutcome<H, N> {
 	removed: BTreeMap<Reverse<N>, Vec<H>>,
 }
 
-impl<H, N: Ord> FinalizationOutcome<H, N> {
+impl<H: Copy, 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) {
@@ -63,6 +63,16 @@ impl<H, N: Ord> FinalizationOutcome<H, N> {
 	pub fn leaves(&self) -> impl Iterator<Item = &H> {
 		self.removed.values().flatten()
 	}
+
+	/// Constructor
+	pub fn new(new_displaced: impl Iterator<Item = (H, N)>) -> Self {
+		let mut removed = BTreeMap::<Reverse<N>, Vec<H>>::new();
+		for (hash, number) in new_displaced {
+			removed.entry(Reverse(number)).or_default().push(hash);
+		}
+
+		FinalizationOutcome { removed }
+	}
 }
 
 /// list of leaf hashes ordered by number (descending).
@@ -151,39 +161,12 @@ where
 		Some(RemoveOutcome { inserted, removed: LeafSetItem { hash, number } })
 	}
 
-	/// 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) -> FinalizationOutcome<H, N> {
-		let boundary = if number == N::zero() {
-			return FinalizationOutcome { removed: BTreeMap::new() }
-		} else {
-			number - N::one()
-		};
-
-		let below_boundary = self.storage.split_off(&Reverse(boundary));
-		FinalizationOutcome { removed: below_boundary }
-	}
-
-	/// The same as [`Self::finalize_height`], but it only simulates the operation.
-	///
-	/// 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) -> FinalizationOutcome<H, N> {
-		let boundary = if number == N::zero() {
-			return FinalizationOutcome { removed: BTreeMap::new() }
-		} else {
-			number - N::one()
-		};
-
-		let below_boundary = self.storage.range(&Reverse(boundary)..);
-		FinalizationOutcome {
-			removed: below_boundary.map(|(k, v)| (k.clone(), v.clone())).collect(),
+	/// Remove all leaves displaced by the last block finalization.
+	pub fn remove_displaced_leaves(&mut self, displaced_leaves: &FinalizationOutcome<H, N>) {
+		for (number, hashes) in &displaced_leaves.removed {
+			for hash in hashes.iter() {
+				self.remove_leaf(number, hash);
+			}
 		}
 	}
 
@@ -420,32 +403,6 @@ mod tests {
 		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]
 	fn flush_to_disk() {
 		const PREFIX: &[u8] = b"abcdefg";
@@ -479,35 +436,4 @@ mod tests {
 		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 = Arc::new(sp_database::MemDb::default());
-
-		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 = Transaction::new();
-		set.prepare_transaction(&mut tx, 0, PREFIX);
-		db.commit(tx).unwrap();
-
-		let _ = set.finalize_height(11);
-		let mut tx = Transaction::new();
-		set.prepare_transaction(&mut tx, 0, PREFIX);
-		db.commit(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, 0, PREFIX).unwrap();
-		assert_eq!(set, set2);
-	}
 }
diff --git a/substrate/client/consensus/babe/src/lib.rs b/substrate/client/consensus/babe/src/lib.rs
index d10bdd8c7e4c..a1ca6fd23b59 100644
--- a/substrate/client/consensus/babe/src/lib.rs
+++ b/substrate/client/consensus/babe/src/lib.rs
@@ -562,9 +562,10 @@ fn aux_storage_cleanup<C: HeaderMetadata<Block> + HeaderBackend<Block>, Block: B
 	// Cleans data for stale forks.
 	let stale_forks = match client.expand_forks(&notification.stale_heads) {
 		Ok(stale_forks) => stale_forks,
-		Err((stale_forks, e)) => {
+		Err(e) => {
 			warn!(target: LOG_TARGET, "{:?}", e);
-			stale_forks
+
+			Default::default()
 		},
 	};
 	hashes.extend(stale_forks.iter());
diff --git a/substrate/client/consensus/babe/src/tests.rs b/substrate/client/consensus/babe/src/tests.rs
index 38c9e1ff6ac2..716067ae4000 100644
--- a/substrate/client/consensus/babe/src/tests.rs
+++ b/substrate/client/consensus/babe/src/tests.rs
@@ -1094,8 +1094,8 @@ async fn obsolete_blocks_aux_data_cleanup() {
 	assert!(aux_data_check(&fork1_hashes[2..3], false));
 	// Present: A4
 	assert!(aux_data_check(&fork1_hashes[3..], true));
-	// Present C4, C5
-	assert!(aux_data_check(&fork3_hashes, true));
+	// Wiped C4, C5
+	assert!(aux_data_check(&fork3_hashes, false));
 }
 
 #[tokio::test]
diff --git a/substrate/client/db/src/lib.rs b/substrate/client/db/src/lib.rs
index 0faa90dfc4f9..36f9aea817c9 100644
--- a/substrate/client/db/src/lib.rs
+++ b/substrate/client/db/src/lib.rs
@@ -68,8 +68,8 @@ use sc_client_api::{
 use sc_state_db::{IsPruned, LastCanonicalized, StateDb};
 use sp_arithmetic::traits::Saturating;
 use sp_blockchain::{
-	Backend as _, CachedHeaderMetadata, Error as ClientError, HeaderBackend, HeaderMetadata,
-	HeaderMetadataCache, Result as ClientResult,
+	Backend as _, CachedHeaderMetadata, DisplacedLeavesAfterFinalization, Error as ClientError,
+	HeaderBackend, HeaderMetadata, HeaderMetadataCache, Result as ClientResult,
 };
 use sp_core::{
 	offchain::OffchainOverlayedChange,
@@ -747,19 +747,6 @@ impl<Block: BlockT> sc_client_api::blockchain::Backend<Block> for BlockchainDb<B
 		Ok(self.leaves.read().hashes())
 	}
 
-	fn displaced_leaves_after_finalizing(
-		&self,
-		block_number: NumberFor<Block>,
-	) -> ClientResult<Vec<Block::Hash>> {
-		Ok(self
-			.leaves
-			.read()
-			.displaced_by_finalize_height(block_number)
-			.leaves()
-			.cloned()
-			.collect::<Vec<_>>())
-	}
-
 	fn children(&self, parent_hash: Block::Hash) -> ClientResult<Vec<Block::Hash>> {
 		children::read_children(&*self.db, columns::META, meta_keys::CHILDREN_PREFIX, parent_hash)
 	}
@@ -1813,14 +1800,13 @@ impl<Block: BlockT> Backend<Block> {
 			apply_state_commit(transaction, commit);
 		}
 
-		let new_displaced = self.blockchain.leaves.write().finalize_height(f_num);
-		self.prune_blocks(
-			transaction,
-			f_num,
-			f_hash,
-			&new_displaced,
-			current_transaction_justifications,
-		)?;
+		let new_displaced = self.blockchain.displaced_leaves_after_finalizing(f_hash, f_num)?;
+		let finalization_outcome =
+			FinalizationOutcome::new(new_displaced.displaced_leaves.clone().into_iter());
+
+		self.blockchain.leaves.write().remove_displaced_leaves(&finalization_outcome);
+
+		self.prune_blocks(transaction, f_num, &new_displaced, current_transaction_justifications)?;
 
 		Ok(())
 	}
@@ -1829,8 +1815,7 @@ impl<Block: BlockT> Backend<Block> {
 		&self,
 		transaction: &mut Transaction<DbHash>,
 		finalized_number: NumberFor<Block>,
-		finalized_hash: Block::Hash,
-		displaced: &FinalizationOutcome<Block::Hash, NumberFor<Block>>,
+		displaced: &DisplacedLeavesAfterFinalization<Block>,
 		current_transaction_justifications: &mut HashMap<Block::Hash, Justification>,
 	) -> ClientResult<()> {
 		match self.blocks_pruning {
@@ -1858,10 +1843,10 @@ impl<Block: BlockT> Backend<Block> {
 
 					self.prune_block(transaction, BlockId::<Block>::number(number))?;
 				}
-				self.prune_displaced_branches(transaction, finalized_hash, displaced)?;
+				self.prune_displaced_branches(transaction, displaced)?;
 			},
 			BlocksPruning::KeepFinalized => {
-				self.prune_displaced_branches(transaction, finalized_hash, displaced)?;
+				self.prune_displaced_branches(transaction, displaced)?;
 			},
 		}
 		Ok(())
@@ -1870,21 +1855,13 @@ impl<Block: BlockT> Backend<Block> {
 	fn prune_displaced_branches(
 		&self,
 		transaction: &mut Transaction<DbHash>,
-		finalized: Block::Hash,
-		displaced: &FinalizationOutcome<Block::Hash, NumberFor<Block>>,
+		displaced: &DisplacedLeavesAfterFinalization<Block>,
 	) -> ClientResult<()> {
 		// Discard all blocks from displaced branches
-		for h in displaced.leaves() {
-			match sp_blockchain::tree_route(&self.blockchain, *h, finalized) {
-				Ok(tree_route) =>
-					for r in tree_route.retracted() {
-						self.blockchain.insert_persisted_body_if_pinned(r.hash)?;
-						self.prune_block(transaction, BlockId::<Block>::hash(r.hash))?;
-					},
-				Err(sp_blockchain::Error::UnknownBlock(_)) => {
-					// Sometimes routes can't be calculated. E.g. after warp sync.
-				},
-				Err(e) => Err(e)?,
+		for (_, tree_route) in displaced.tree_routes.iter() {
+			for r in tree_route.retracted() {
+				self.blockchain.insert_persisted_body_if_pinned(r.hash)?;
+				self.prune_block(transaction, BlockId::<Block>::hash(r.hash))?;
 			}
 		}
 		Ok(())
@@ -3190,6 +3167,9 @@ pub(crate) mod tests {
 
 	#[test]
 	fn test_leaves_pruned_on_finality() {
+		//   / 1b - 2b - 3b
+		// 0 - 1a - 2a
+		//   \ 1c
 		let backend: Backend<Block> = Backend::new_test(10, 10);
 		let block0 = insert_header(&backend, 0, Default::default(), None, Default::default());
 
@@ -3201,18 +3181,16 @@ pub(crate) mod tests {
 
 		let block2_a = insert_header(&backend, 2, block1_a, None, Default::default());
 		let block2_b = insert_header(&backend, 2, block1_b, None, Default::default());
-		let block2_c = insert_header(&backend, 2, block1_b, None, [1; 32].into());
 
-		assert_eq!(
-			backend.blockchain().leaves().unwrap(),
-			vec![block2_a, block2_b, block2_c, block1_c]
-		);
+		let block3_b = insert_header(&backend, 3, block2_b, None, [3; 32].into());
+
+		assert_eq!(backend.blockchain().leaves().unwrap(), vec![block3_b, block2_a, block1_c]);
 
 		backend.finalize_block(block1_a, None).unwrap();
 		backend.finalize_block(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]);
+		// All leaves are pruned that are known to not belong to canonical branch
+		assert_eq!(backend.blockchain().leaves().unwrap(), vec![block2_a]);
 	}
 
 	#[test]
diff --git a/substrate/client/merkle-mountain-range/src/offchain_mmr.rs b/substrate/client/merkle-mountain-range/src/offchain_mmr.rs
index 3c3f0beb6c6a..94593f9c2c7b 100644
--- a/substrate/client/merkle-mountain-range/src/offchain_mmr.rs
+++ b/substrate/client/merkle-mountain-range/src/offchain_mmr.rs
@@ -33,7 +33,7 @@ use sp_runtime::{
 	traits::{Block, Header, NumberFor, One},
 	Saturating,
 };
-use std::{collections::VecDeque, sync::Arc};
+use std::{collections::VecDeque, default::Default, sync::Arc};
 
 /// `OffchainMMR` exposes MMR offchain canonicalization and pruning logic.
 pub struct OffchainMmr<B: Block, BE: Backend<B>, C> {
@@ -273,12 +273,11 @@ where
 		self.write_gadget_state_or_log();
 
 		// Remove offchain MMR nodes for stale forks.
-		let stale_forks = self.client.expand_forks(&notification.stale_heads).unwrap_or_else(
-			|(stale_forks, e)| {
-				warn!(target: LOG_TARGET, "{:?}", e);
-				stale_forks
-			},
-		);
+		let stale_forks = self.client.expand_forks(&notification.stale_heads).unwrap_or_else(|e| {
+			warn!(target: LOG_TARGET, "{:?}", e);
+
+			Default::default()
+		});
 		for hash in stale_forks.iter() {
 			self.prune_branch(hash);
 		}
diff --git a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
index 363d11235dda..b195e05b6649 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
+++ b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
@@ -2487,6 +2487,7 @@ async fn follow_report_multiple_pruned_block() {
 	client.finalize_block(block_3_hash, None).unwrap();
 
 	// Finalizing block 3 directly will also result in block 1 and 2 being finalized.
+	// It will also mark block 2 and block 3 from the fork as pruned.
 	let event: FollowEvent<String> = get_next_event(&mut sub).await;
 	let expected = FollowEvent::Finalized(Finalized {
 		finalized_block_hashes: vec![
@@ -2494,7 +2495,7 @@ async fn follow_report_multiple_pruned_block() {
 			format!("{:?}", block_2_hash),
 			format!("{:?}", block_3_hash),
 		],
-		pruned_block_hashes: vec![],
+		pruned_block_hashes: vec![format!("{:?}", block_2_f_hash), format!("{:?}", block_3_f_hash)],
 	});
 	assert_eq!(event, expected);
 
@@ -2504,7 +2505,6 @@ async fn follow_report_multiple_pruned_block() {
 	//                                                  ^^^ finalized
 	//           -> block 1 -> block 2_f -> block 3_f
 	//
-	// Mark block 4 as finalized to force block 2_f and 3_f to get pruned.
 
 	let block_4 = BlockBuilderBuilder::new(&*client)
 		.on_parent_block(block_3.hash())
@@ -2535,11 +2535,11 @@ async fn follow_report_multiple_pruned_block() {
 	});
 	assert_eq!(event, expected);
 
-	// Block 4 and 5 be reported as pruned, not just the stale head (block 5).
+	// Blocks from the fork were pruned earlier.
 	let event: FollowEvent<String> = get_next_event(&mut sub).await;
 	let expected = FollowEvent::Finalized(Finalized {
 		finalized_block_hashes: vec![format!("{:?}", block_4_hash)],
-		pruned_block_hashes: vec![format!("{:?}", block_2_f_hash), format!("{:?}", block_3_f_hash)],
+		pruned_block_hashes: vec![],
 	});
 	assert_eq!(event, expected);
 }
@@ -3714,16 +3714,8 @@ async fn follow_unique_pruned_blocks() {
 	// The chainHead will see block 5 as the best block. However, the
 	// client will finalize the block 6, which is on another fork.
 	//
-	// When the block 6 is finalized, block 2 block 3 block 4 and block 5 are placed on an invalid
-	// fork. However, pruning of blocks happens on level N - 1.
-	// Therefore, no pruned blocks are reported yet.
+	// When the block 6 is finalized all blocks from the stale forks (2, 3, 4, 5) are pruned.
 	//
-	// When the block 7 is finalized, block 3 is detected as stale. At this step, block 2 and 3
-	// are reported as pruned.
-	//
-	// When the block 8 is finalized, block 5 block 4 and block 2 are detected as stale. However,
-	// only blocks 5 and 4 are reported as pruned. This is because the block 2 was previously
-	// reported.
 
 	// Initial setup steps:
 	let block_1_hash =
@@ -3776,16 +3768,33 @@ async fn follow_unique_pruned_blocks() {
 	});
 	assert_eq!(event, expected);
 
-	// Block 2 must be reported as pruned, even if it was the previous best.
-	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	// All blocks from stale forks are pruned when we finalize block 6.
+	let mut event: FollowEvent<String> = get_next_event(&mut sub).await;
+
+	// Sort pruned block hashes to counter flaky test caused by event generation (get_pruned_hashes)
+	if let FollowEvent::Finalized(Finalized { pruned_block_hashes, .. }) = &mut event {
+		pruned_block_hashes.sort();
+	}
+	let expected_pruned_block_hashes = {
+		let mut hashes = vec![
+			format!("{:?}", block_2_hash),
+			format!("{:?}", block_3_hash),
+			format!("{:?}", block_4_hash),
+			format!("{:?}", block_5_hash),
+		];
+		hashes.sort();
+		hashes
+	};
+
 	let expected = FollowEvent::Finalized(Finalized {
 		finalized_block_hashes: vec![
 			format!("{:?}", block_1_hash),
 			format!("{:?}", block_2_f_hash),
 			format!("{:?}", block_6_hash),
 		],
-		pruned_block_hashes: vec![],
+		pruned_block_hashes: expected_pruned_block_hashes,
 	});
+
 	assert_eq!(event, expected);
 
 	// Pruned hash can be unpinned.
@@ -3802,9 +3811,10 @@ async fn follow_unique_pruned_blocks() {
 	client.finalize_block(block_7_hash, None).unwrap();
 
 	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	// All necessary blocks were pruned on block 6 finalization.
 	let expected = FollowEvent::Finalized(Finalized {
 		finalized_block_hashes: vec![format!("{:?}", block_7_hash)],
-		pruned_block_hashes: vec![format!("{:?}", block_2_hash), format!("{:?}", block_3_hash)],
+		pruned_block_hashes: vec![],
 	});
 	assert_eq!(event, expected);
 
@@ -3815,10 +3825,11 @@ async fn follow_unique_pruned_blocks() {
 	// Finalize the block 8.
 	client.finalize_block(block_8_hash, None).unwrap();
 
+	// All necessary blocks were pruned on block 6 finalization.
 	let event: FollowEvent<String> = get_next_event(&mut sub).await;
 	let expected = FollowEvent::Finalized(Finalized {
 		finalized_block_hashes: vec![format!("{:?}", block_8_hash)],
-		pruned_block_hashes: vec![format!("{:?}", block_4_hash), format!("{:?}", block_5_hash)],
+		pruned_block_hashes: vec![],
 	});
 	assert_eq!(event, expected);
 }
diff --git a/substrate/client/service/src/client/client.rs b/substrate/client/service/src/client/client.rs
index 35e8b53a09cf..3c25c233775b 100644
--- a/substrate/client/service/src/client/client.rs
+++ b/substrate/client/service/src/client/client.rs
@@ -978,8 +978,12 @@ where
 
 			// The stale heads are the leaves that will be displaced after the
 			// block is finalized.
-			let stale_heads =
-				self.backend.blockchain().displaced_leaves_after_finalizing(block_number)?;
+			let stale_heads = self
+				.backend
+				.blockchain()
+				.displaced_leaves_after_finalizing(hash, block_number)?
+				.hashes()
+				.collect();
 
 			let header = self
 				.backend
diff --git a/substrate/client/service/test/src/client/mod.rs b/substrate/client/service/test/src/client/mod.rs
index 4fcb7c160cb4..6542830c998b 100644
--- a/substrate/client/service/test/src/client/mod.rs
+++ b/substrate/client/service/test/src/client/mod.rs
@@ -1164,7 +1164,7 @@ fn finalizing_diverged_block_should_trigger_reorg() {
 
 	// G -> A1 -> A2
 	//   \
-	//    -> B1 -> B2
+	//    -> B1 -> B2 -> B3
 
 	let mut finality_notifications = client.finality_notification_stream();
 
@@ -1249,8 +1249,8 @@ fn finalizing_diverged_block_should_trigger_reorg() {
 
 	ClientExt::finalize_block(&client, b3.hash(), None).unwrap();
 
-	finality_notification_check(&mut finality_notifications, &[b1.hash()], &[]);
-	finality_notification_check(&mut finality_notifications, &[b2.hash(), b3.hash()], &[a2.hash()]);
+	finality_notification_check(&mut finality_notifications, &[b1.hash()], &[a2.hash()]);
+	finality_notification_check(&mut finality_notifications, &[b2.hash(), b3.hash()], &[]);
 	assert!(matches!(finality_notifications.try_recv().unwrap_err(), TryRecvError::Empty));
 }
 
@@ -1371,8 +1371,12 @@ fn finality_notifications_content() {
 	// Import and finalize D4
 	block_on(client.import_as_final(BlockOrigin::Own, d4.clone())).unwrap();
 
-	finality_notification_check(&mut finality_notifications, &[a1.hash(), a2.hash()], &[c1.hash()]);
-	finality_notification_check(&mut finality_notifications, &[d3.hash(), d4.hash()], &[b2.hash()]);
+	finality_notification_check(
+		&mut finality_notifications,
+		&[a1.hash(), a2.hash()],
+		&[c1.hash(), b2.hash()],
+	);
+	finality_notification_check(&mut finality_notifications, &[d3.hash(), d4.hash()], &[a3.hash()]);
 	assert!(matches!(finality_notifications.try_recv().unwrap_err(), TryRecvError::Empty));
 }
 
@@ -1601,9 +1605,9 @@ fn doesnt_import_blocks_that_revert_finality() {
 	block_on(client.import(BlockOrigin::Own, a3.clone())).unwrap();
 	ClientExt::finalize_block(&client, a3.hash(), None).unwrap();
 
-	finality_notification_check(&mut finality_notifications, &[a1.hash(), a2.hash()], &[]);
+	finality_notification_check(&mut finality_notifications, &[a1.hash(), a2.hash()], &[b2.hash()]);
 
-	finality_notification_check(&mut finality_notifications, &[a3.hash()], &[b2.hash()]);
+	finality_notification_check(&mut finality_notifications, &[a3.hash()], &[]);
 
 	assert!(matches!(finality_notifications.try_recv().unwrap_err(), TryRecvError::Empty));
 }
diff --git a/substrate/primitives/blockchain/src/backend.rs b/substrate/primitives/blockchain/src/backend.rs
index 7a09865f858d..06e5b682964a 100644
--- a/substrate/primitives/blockchain/src/backend.rs
+++ b/substrate/primitives/blockchain/src/backend.rs
@@ -21,14 +21,17 @@ use log::warn;
 use parking_lot::RwLock;
 use sp_runtime::{
 	generic::BlockId,
-	traits::{Block as BlockT, Header as HeaderT, NumberFor, Saturating},
+	traits::{Block as BlockT, Header as HeaderT, NumberFor, Zero},
 	Justifications,
 };
-use std::collections::btree_set::BTreeSet;
+use std::collections::{btree_map::BTreeMap, btree_set::BTreeSet};
 
 use crate::header_metadata::HeaderMetadata;
 
-use crate::error::{Error, Result};
+use crate::{
+	error::{Error, Result},
+	tree_route, TreeRoute,
+};
 
 /// Blockchain database header backend. Does not perform any validation.
 pub trait HeaderBackend<Block: BlockT>: Send + Sync {
@@ -89,62 +92,32 @@ pub trait HeaderBackend<Block: BlockT>: Send + Sync {
 pub trait ForkBackend<Block: BlockT>:
 	HeaderMetadata<Block> + HeaderBackend<Block> + Send + Sync
 {
-	/// Best effort to get all the header hashes that are part of the provided forks
-	/// starting only from the fork heads.
+	/// Returns block hashes for provided fork heads. It skips the fork if when blocks are missing
+	/// (e.g. warp-sync) and internal `tree_route` function fails.
 	///
-	/// The function tries to reconstruct the route from the fork head to the canonical chain.
-	/// If any of the hashes on the route can't be found in the db, the function won't be able
-	/// to reconstruct the route anymore. In this case it will give up expanding the current fork,
-	/// move on to the next ones and at the end it will return an error that also contains
-	/// the partially expanded forks.
+	/// Example:
+	///  G --- A1 --- A2 --- A3 --- A4           ( < fork1 )
+	///                       \-----C4 --- C5    ( < fork2 )
+	/// We finalize A3 and call expand_fork(C5). Result = (C5,C4).
 	fn expand_forks(
 		&self,
 		fork_heads: &[Block::Hash],
-	) -> std::result::Result<BTreeSet<Block::Hash>, (BTreeSet<Block::Hash>, Error)> {
-		let mut missing_blocks = vec![];
+	) -> std::result::Result<BTreeSet<Block::Hash>, Error> {
 		let mut expanded_forks = BTreeSet::new();
 		for fork_head in fork_heads {
-			let mut route_head = *fork_head;
-			// Insert stale blocks hashes until canonical chain is reached.
-			// If we reach a block that is already part of the `expanded_forks` we can stop
-			// processing the fork.
-			while expanded_forks.insert(route_head) {
-				match self.header_metadata(route_head) {
-					Ok(meta) => {
-						// If the parent is part of the canonical chain or there doesn't exist a
-						// block hash for the parent number (bug?!), we can abort adding blocks.
-						let parent_number = meta.number.saturating_sub(1u32.into());
-						match self.hash(parent_number) {
-							Ok(Some(parent_hash)) =>
-								if parent_hash == meta.parent {
-									break
-								},
-							Ok(None) | Err(_) => {
-								missing_blocks.push(BlockId::<Block>::Number(parent_number));
-								break
-							},
-						}
-
-						route_head = meta.parent;
-					},
-					Err(_e) => {
-						missing_blocks.push(BlockId::<Block>::Hash(route_head));
-						break
-					},
-				}
+			match tree_route(self, *fork_head, self.info().finalized_hash) {
+				Ok(tree_route) => {
+					for block in tree_route.retracted() {
+						expanded_forks.insert(block.hash);
+					}
+					continue
+				},
+				Err(_) => {
+					// There are cases when blocks are missing (e.g. warp-sync).
+				},
 			}
 		}
 
-		if !missing_blocks.is_empty() {
-			return Err((
-				expanded_forks,
-				Error::UnknownBlocks(format!(
-					"Missing stale headers {:?} while expanding forks {:?}.",
-					fork_heads, missing_blocks
-				)),
-			))
-		}
-
 		Ok(expanded_forks)
 	}
 }
@@ -172,14 +145,6 @@ pub trait Backend<Block: BlockT>:
 	/// Results must be ordered best (longest, highest) chain first.
 	fn leaves(&self) -> Result<Vec<Block::Hash>>;
 
-	/// Returns displaced leaves after the given block would be finalized.
-	///
-	/// The returned leaves do not contain the leaves from the same height as `block_number`.
-	fn displaced_leaves_after_finalizing(
-		&self,
-		block_number: NumberFor<Block>,
-	) -> Result<Vec<Block::Hash>>;
-
 	/// Return hashes of all blocks that are children of the block with `parent_hash`.
 	fn children(&self, parent_hash: Block::Hash) -> Result<Vec<Block::Hash>>;
 
@@ -255,6 +220,67 @@ pub trait Backend<Block: BlockT>:
 	}
 
 	fn block_indexed_body(&self, hash: Block::Hash) -> Result<Option<Vec<Vec<u8>>>>;
+
+	/// Returns all leaves that will be displaced after the block finalization.
+	fn displaced_leaves_after_finalizing(
+		&self,
+		finalized_block_hash: Block::Hash,
+		finalized_block_number: NumberFor<Block>,
+	) -> std::result::Result<DisplacedLeavesAfterFinalization<Block>, Error> {
+		let mut result = DisplacedLeavesAfterFinalization::default();
+
+		if finalized_block_number == Zero::zero() {
+			return Ok(result)
+		}
+
+		// For each leaf determine whether it belongs to a non-canonical branch.
+		for leaf_hash in self.leaves()? {
+			let leaf_block_header = self.expect_header(leaf_hash)?;
+			let leaf_number = *leaf_block_header.number();
+
+			let leaf_tree_route = match tree_route(self, leaf_hash, finalized_block_hash) {
+				Ok(tree_route) => tree_route,
+				Err(Error::UnknownBlock(_)) => {
+					// Sometimes routes can't be calculated. E.g. after warp sync.
+					continue;
+				},
+				Err(e) => Err(e)?,
+			};
+
+			// Is it a stale fork?
+			let needs_pruning = leaf_tree_route.common_block().hash != finalized_block_hash;
+
+			if needs_pruning {
+				result.displaced_leaves.insert(leaf_hash, leaf_number);
+				result.tree_routes.insert(leaf_hash, leaf_tree_route);
+			}
+		}
+
+		Ok(result)
+	}
+}
+
+/// Result of  [`Backend::displaced_leaves_after_finalizing`].
+#[derive(Clone, Debug)]
+pub struct DisplacedLeavesAfterFinalization<Block: BlockT> {
+	/// A collection of hashes and block numbers for displaced leaves.
+	pub displaced_leaves: BTreeMap<Block::Hash, NumberFor<Block>>,
+
+	/// A collection of tree routes from the leaves to finalized block.
+	pub tree_routes: BTreeMap<Block::Hash, TreeRoute<Block>>,
+}
+
+impl<Block: BlockT> Default for DisplacedLeavesAfterFinalization<Block> {
+	fn default() -> Self {
+		Self { displaced_leaves: Default::default(), tree_routes: Default::default() }
+	}
+}
+
+impl<Block: BlockT> DisplacedLeavesAfterFinalization<Block> {
+	/// Returns a collection of hashes for the displaced leaves.
+	pub fn hashes(&self) -> impl Iterator<Item = Block::Hash> + '_ {
+		self.displaced_leaves.keys().cloned()
+	}
 }
 
 /// Blockchain info
diff --git a/substrate/primitives/blockchain/src/header_metadata.rs b/substrate/primitives/blockchain/src/header_metadata.rs
index ccd640c0567a..27caaae71add 100644
--- a/substrate/primitives/blockchain/src/header_metadata.rs
+++ b/substrate/primitives/blockchain/src/header_metadata.rs
@@ -97,7 +97,7 @@ pub fn lowest_common_ancestor<Block: BlockT, T: HeaderMetadata<Block> + ?Sized>(
 }
 
 /// Compute a tree-route between two blocks. See tree-route docs for more details.
-pub fn tree_route<Block: BlockT, T: HeaderMetadata<Block>>(
+pub fn tree_route<Block: BlockT, T: HeaderMetadata<Block> + ?Sized>(
 	backend: &T,
 	from: Block::Hash,
 	to: Block::Hash,