From 9decf7b4fb9b8a14d65af8a47627a9d558c67e8c Mon Sep 17 00:00:00 2001
From: Shamil Gadelshin <shamilgadelshin@gmail.com>
Date: Fri, 19 Jan 2024 14:45:42 +0700
Subject: [PATCH 01/19] Change forks pruning algorithm.

- Prune all possible forks after block finalizing without height limit.
---
 substrate/client/api/src/in_mem.rs            | 14 -----
 substrate/client/api/src/leaves.rs            | 56 +++++++----------
 substrate/client/db/src/lib.rs                | 36 +++++------
 substrate/client/service/src/client/client.rs |  9 ++-
 .../primitives/blockchain/src/backend.rs      | 62 ++++++++++++++++---
 5 files changed, 96 insertions(+), 81 deletions(-)

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..45dd6c83bb7c 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,21 @@ 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))
+				.and_modify(|hashes| {
+					hashes.push(hash);
+				})
+				.or_insert(vec![hash]);
+		}
+
+		FinalizationOutcome { removed }
+	}
 }
 
 /// list of leaf hashes ordered by number (descending).
@@ -151,39 +166,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);
+			}
 		}
 	}
 
diff --git a/substrate/client/db/src/lib.rs b/substrate/client/db/src/lib.rs
index 0faa90dfc4f9..6522a913befb 100644
--- a/substrate/client/db/src/lib.rs
+++ b/substrate/client/db/src/lib.rs
@@ -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,12 +1800,16 @@ impl<Block: BlockT> Backend<Block> {
 			apply_state_commit(transaction, commit);
 		}
 
-		let new_displaced = self.blockchain.leaves.write().finalize_height(f_num);
+		let new_displaced = self.blockchain.displaced_leaves_after_finalizing(f_hash, f_num)?;
+		let finalization_outcome = FinalizationOutcome::new(new_displaced.into_iter());
+
+		self.blockchain.leaves.write().remove_displaced_leaves(&finalization_outcome);
+
 		self.prune_blocks(
 			transaction,
 			f_num,
 			f_hash,
-			&new_displaced,
+			&finalization_outcome,
 			current_transaction_justifications,
 		)?;
 
@@ -3190,6 +3181,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 +3195,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/service/src/client/client.rs b/substrate/client/service/src/client/client.rs
index 35e8b53a09cf..876f131c1a7b 100644
--- a/substrate/client/service/src/client/client.rs
+++ b/substrate/client/service/src/client/client.rs
@@ -978,8 +978,13 @@ 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)?
+				.into_iter()
+				.map(|(h, _)| h)
+				.collect();
 
 			let header = self
 				.backend
diff --git a/substrate/primitives/blockchain/src/backend.rs b/substrate/primitives/blockchain/src/backend.rs
index 7a09865f858d..b119acbefbe3 100644
--- a/substrate/primitives/blockchain/src/backend.rs
+++ b/substrate/primitives/blockchain/src/backend.rs
@@ -21,7 +21,7 @@ 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, Saturating, Zero},
 	Justifications,
 };
 use std::collections::btree_set::BTreeSet;
@@ -172,14 +172,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 +247,58 @@ 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<Vec<(Block::Hash, NumberFor<Block>)>, Error> {
+		if finalized_block_number == Zero::zero() {
+			return Ok(Vec::new())
+		}
+
+		let mut displaced_leaves = Vec::new();
+
+		let finalized_block_header = self.expect_header(finalized_block_hash)?;
+		// For each leaf determine whether it belongs to a non-canonical branch.
+		for leaf_hash in self.leaves()? {
+			let mut fork_block_header = self.expect_header(leaf_hash)?;
+
+			let leaf_number = *fork_block_header.number();
+
+			let mut needs_pruning = false;
+			// All leaves will eventually have as ancestor either the finalized block or its
+			// parent. All other forks were cleared on the previous block finalization.
+			loop {
+				// The block ends up in the finalized block. All forks are valid at this point.
+				if fork_block_header.hash() == finalized_block_hash {
+					break
+				}
+
+				// The block ends up in the parent block of the finalized block. It's a stale fork.
+				if fork_block_header.hash() == *finalized_block_header.parent_hash() {
+					needs_pruning = true;
+					break
+				}
+
+				if let Some(parent_header) = self.header(*fork_block_header.parent_hash())? {
+					fork_block_header = parent_header;
+				} else {
+					// Sometimes routes can't be calculated. E.g. after warp sync.
+					needs_pruning = true;
+					break
+				}
+			}
+
+			// Fork ended up in the parent of the finalized block.
+			if needs_pruning {
+				displaced_leaves.push((leaf_hash, leaf_number));
+			}
+		}
+
+		Ok(displaced_leaves)
+	}
 }
 
 /// Blockchain info

From 64094f213f7eb2ed401ffbd9417b290f8027cc24 Mon Sep 17 00:00:00 2001
From: shamil-gadelshin <shamilgadelshin@gmail.com>
Date: Thu, 4 Apr 2024 15:00:26 +0700
Subject: [PATCH 02/19] Update substrate/client/api/src/leaves.rs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Bastian Köcher <git@kchr.de>
---
 substrate/client/api/src/leaves.rs | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/substrate/client/api/src/leaves.rs b/substrate/client/api/src/leaves.rs
index 45dd6c83bb7c..0f84abb56640 100644
--- a/substrate/client/api/src/leaves.rs
+++ b/substrate/client/api/src/leaves.rs
@@ -70,10 +70,7 @@ impl<H: Copy, N: Ord> FinalizationOutcome<H, N> {
 		for (hash, number) in new_displaced {
 			removed
 				.entry(Reverse(number))
-				.and_modify(|hashes| {
-					hashes.push(hash);
-				})
-				.or_insert(vec![hash]);
+				.or_default().push(hash);
 		}
 
 		FinalizationOutcome { removed }

From 62ab0aa1740c91312efc3f0706ca702a5127535d Mon Sep 17 00:00:00 2001
From: shamil-gadelshin <shamilgadelshin@gmail.com>
Date: Thu, 4 Apr 2024 15:00:52 +0700
Subject: [PATCH 03/19] Update substrate/primitives/blockchain/src/backend.rs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Bastian Köcher <git@kchr.de>
---
 substrate/primitives/blockchain/src/backend.rs | 15 ++++++---------
 1 file changed, 6 insertions(+), 9 deletions(-)

diff --git a/substrate/primitives/blockchain/src/backend.rs b/substrate/primitives/blockchain/src/backend.rs
index b119acbefbe3..a1e0da3ba133 100644
--- a/substrate/primitives/blockchain/src/backend.rs
+++ b/substrate/primitives/blockchain/src/backend.rs
@@ -267,29 +267,26 @@ pub trait Backend<Block: BlockT>:
 
 			let leaf_number = *fork_block_header.number();
 
-			let mut needs_pruning = false;
 			// All leaves will eventually have as ancestor either the finalized block or its
 			// parent. All other forks were cleared on the previous block finalization.
-			loop {
+			let needs_pruning = loop {
 				// The block ends up in the finalized block. All forks are valid at this point.
 				if fork_block_header.hash() == finalized_block_hash {
-					break
+					break false
 				}
 
 				// The block ends up in the parent block of the finalized block. It's a stale fork.
-				if fork_block_header.hash() == *finalized_block_header.parent_hash() {
-					needs_pruning = true;
-					break
+				if *fork_block_header.number() <= finalized_block_number {
+					break true
 				}
 
 				if let Some(parent_header) = self.header(*fork_block_header.parent_hash())? {
 					fork_block_header = parent_header;
 				} else {
 					// Sometimes routes can't be calculated. E.g. after warp sync.
-					needs_pruning = true;
-					break
+					break true
 				}
-			}
+			};
 
 			// Fork ended up in the parent of the finalized block.
 			if needs_pruning {

From 962b5d3dbc4dd0cf12657d9367bda0cd2cfc75a5 Mon Sep 17 00:00:00 2001
From: Shamil Gadelshin <shamilgadelshin@gmail.com>
Date: Fri, 5 Apr 2024 19:16:52 +0700
Subject: [PATCH 04/19] Update fork calculation algorithm.

---
 substrate/client/api/src/leaves.rs            |  4 +-
 substrate/client/db/src/lib.rs                | 40 ++++------
 substrate/client/service/src/client/client.rs |  4 +-
 .../primitives/blockchain/src/backend.rs      | 77 +++++++++++--------
 .../blockchain/src/header_metadata.rs         |  2 +-
 5 files changed, 62 insertions(+), 65 deletions(-)

diff --git a/substrate/client/api/src/leaves.rs b/substrate/client/api/src/leaves.rs
index 0f84abb56640..0d420577af95 100644
--- a/substrate/client/api/src/leaves.rs
+++ b/substrate/client/api/src/leaves.rs
@@ -68,9 +68,7 @@ impl<H: Copy, N: Ord> FinalizationOutcome<H, N> {
 	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);
+			removed.entry(Reverse(number)).or_default().push(hash);
 		}
 
 		FinalizationOutcome { removed }
diff --git a/substrate/client/db/src/lib.rs b/substrate/client/db/src/lib.rs
index 6522a913befb..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,
@@ -1801,17 +1801,12 @@ impl<Block: BlockT> Backend<Block> {
 		}
 
 		let new_displaced = self.blockchain.displaced_leaves_after_finalizing(f_hash, f_num)?;
-		let finalization_outcome = FinalizationOutcome::new(new_displaced.into_iter());
+		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,
-			f_hash,
-			&finalization_outcome,
-			current_transaction_justifications,
-		)?;
+		self.prune_blocks(transaction, f_num, &new_displaced, current_transaction_justifications)?;
 
 		Ok(())
 	}
@@ -1820,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 {
@@ -1849,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(())
@@ -1861,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(())
diff --git a/substrate/client/service/src/client/client.rs b/substrate/client/service/src/client/client.rs
index 876f131c1a7b..f721b64f6c53 100644
--- a/substrate/client/service/src/client/client.rs
+++ b/substrate/client/service/src/client/client.rs
@@ -982,9 +982,7 @@ where
 				.backend
 				.blockchain()
 				.displaced_leaves_after_finalizing(hash, block_number)?
-				.into_iter()
-				.map(|(h, _)| h)
-				.collect();
+				.hashes();
 
 			let header = self
 				.backend
diff --git a/substrate/primitives/blockchain/src/backend.rs b/substrate/primitives/blockchain/src/backend.rs
index a1e0da3ba133..f72befba553b 100644
--- a/substrate/primitives/blockchain/src/backend.rs
+++ b/substrate/primitives/blockchain/src/backend.rs
@@ -24,11 +24,14 @@ use sp_runtime::{
 	traits::{Block as BlockT, Header as HeaderT, NumberFor, Saturating, 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 {
@@ -253,48 +256,60 @@ pub trait Backend<Block: BlockT>:
 		&self,
 		finalized_block_hash: Block::Hash,
 		finalized_block_number: NumberFor<Block>,
-	) -> std::result::Result<Vec<(Block::Hash, NumberFor<Block>)>, Error> {
+	) -> std::result::Result<DisplacedLeavesAfterFinalization<Block>, Error> {
+		let mut result = DisplacedLeavesAfterFinalization::new();
+
 		if finalized_block_number == Zero::zero() {
-			return Ok(Vec::new())
+			return Ok(result)
 		}
 
-		let mut displaced_leaves = Vec::new();
-
-		let finalized_block_header = self.expect_header(finalized_block_hash)?;
 		// For each leaf determine whether it belongs to a non-canonical branch.
 		for leaf_hash in self.leaves()? {
-			let mut fork_block_header = self.expect_header(leaf_hash)?;
+			let leaf_block_header = self.expect_header(leaf_hash)?;
+			let leaf_number = *leaf_block_header.number();
 
-			let leaf_number = *fork_block_header.number();
-
-			// All leaves will eventually have as ancestor either the finalized block or its
-			// parent. All other forks were cleared on the previous block finalization.
-			let needs_pruning = loop {
-				// The block ends up in the finalized block. All forks are valid at this point.
-				if fork_block_header.hash() == finalized_block_hash {
-					break false
-				}
-
-				// The block ends up in the parent block of the finalized block. It's a stale fork.
-				if *fork_block_header.number() <= finalized_block_number {
-					break true
-				}
-
-				if let Some(parent_header) = self.header(*fork_block_header.parent_hash())? {
-					fork_block_header = parent_header;
-				} else {
+			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.
-					break true
-				}
+					continue;
+				},
+				Err(e) => Err(e)?,
 			};
 
-			// Fork ended up in the parent of the finalized block.
+			// Is it a stale fork?
+			let needs_pruning = leaf_tree_route.common_block().hash != finalized_block_hash;
+
 			if needs_pruning {
-				displaced_leaves.push((leaf_hash, leaf_number));
+				result.displaced_leaves.insert(leaf_hash, leaf_number);
+				result.tree_routes.insert(leaf_hash, leaf_tree_route);
 			}
 		}
 
-		Ok(displaced_leaves)
+		Ok(result)
+	}
+}
+
+/// Calculation result of the displaced leaves after the block finalization.
+#[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> DisplacedLeavesAfterFinalization<Block> {
+	/// Returns a collection of hashes for the displaced leaves.
+	pub fn hashes(&self) -> Vec<Block::Hash> {
+		self.displaced_leaves.keys().map(|h| *h).collect()
+	}
+
+	/// Constructor. We need this explicit initialization to not introduce a requirement
+	/// `Block: Default`
+	pub fn new() -> Self {
+		Self { displaced_leaves: BTreeMap::new(), tree_routes: BTreeMap::new() }
 	}
 }
 
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,

From d325870387b5d5aad1d145b297dc593349140367 Mon Sep 17 00:00:00 2001
From: Shamil Gadelshin <shamilgadelshin@gmail.com>
Date: Fri, 5 Apr 2024 22:05:15 +0700
Subject: [PATCH 05/19] Remove obsolete tests.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

`finalize_height` method doesn’t exist. It was used to determine the forks and that algorithm changed.
---
 substrate/client/api/src/leaves.rs | 57 ------------------------------
 1 file changed, 57 deletions(-)

diff --git a/substrate/client/api/src/leaves.rs b/substrate/client/api/src/leaves.rs
index 0d420577af95..e129de8bf3fa 100644
--- a/substrate/client/api/src/leaves.rs
+++ b/substrate/client/api/src/leaves.rs
@@ -403,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";
@@ -462,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);
-	}
 }

From 7e05f0d4e4180bb197dc5820bbc8aa6ca357b6b1 Mon Sep 17 00:00:00 2001
From: Shamil Gadelshin <shamilgadelshin@gmail.com>
Date: Fri, 5 Apr 2024 23:52:31 +0700
Subject: [PATCH 06/19] Add pr_3962.prdoc

---
 prdoc/pr_3962.prdoc | 11 +++++++++++
 1 file changed, 11 insertions(+)
 create mode 100644 prdoc/pr_3962.prdoc

diff --git a/prdoc/pr_3962.prdoc b/prdoc/pr_3962.prdoc
new file mode 100644
index 000000000000..3d5a878b8987
--- /dev/null
+++ b/prdoc/pr_3962.prdoc
@@ -0,0 +1,11 @@
+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.
+
+crates:
+  - name: sc-client-api
+  - name: sc-client-db
+  - name: sp-blockchain

From 4aeeeb6911fc5069187bd32dbb9e8418a0dc2792 Mon Sep 17 00:00:00 2001
From: Shamil Gadelshin <shamilgadelshin@gmail.com>
Date: Tue, 9 Apr 2024 14:12:52 +0700
Subject: [PATCH 07/19] Update expand_fork() function and fix test.

---
 substrate/client/consensus/babe/src/tests.rs   |  4 ++--
 substrate/primitives/blockchain/src/backend.rs | 12 ++++++++++++
 2 files changed, 14 insertions(+), 2 deletions(-)

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/primitives/blockchain/src/backend.rs b/substrate/primitives/blockchain/src/backend.rs
index f72befba553b..309b3e265c91 100644
--- a/substrate/primitives/blockchain/src/backend.rs
+++ b/substrate/primitives/blockchain/src/backend.rs
@@ -107,6 +107,18 @@ pub trait ForkBackend<Block: BlockT>:
 		let mut missing_blocks = vec![];
 		let mut expanded_forks = BTreeSet::new();
 		for fork_head in fork_heads {
+			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(_) => {
+					// Continue with fallback algorithm
+				},
+			}
+
 			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

From 27fcc5f9727ce8baa74649254024158edb1dc538 Mon Sep 17 00:00:00 2001
From: Shamil Gadelshin <shamilgadelshin@gmail.com>
Date: Tue, 9 Apr 2024 14:39:42 +0700
Subject: [PATCH 08/19] Update `follow_report_multiple_pruned_block` test.

---
 substrate/client/rpc-spec-v2/src/chain_head/tests.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

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 c3f10a201c58..3174b91ad2b2 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
+++ b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
@@ -2436,7 +2436,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);
 
@@ -2481,7 +2481,7 @@ async fn follow_report_multiple_pruned_block() {
 	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);
 }

From 57816e9767e068f73c130344378cf94543e09c81 Mon Sep 17 00:00:00 2001
From: Shamil Gadelshin <shamilgadelshin@gmail.com>
Date: Tue, 9 Apr 2024 20:49:16 +0700
Subject: [PATCH 09/19] Update `sc-service-test` tests.

---
 .../client/service/test/src/client/mod.rs      | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/substrate/client/service/test/src/client/mod.rs b/substrate/client/service/test/src/client/mod.rs
index 51dcc4966e58..dc4b04bcc98e 100644
--- a/substrate/client/service/test/src/client/mod.rs
+++ b/substrate/client/service/test/src/client/mod.rs
@@ -1165,7 +1165,7 @@ fn finalizing_diverged_block_should_trigger_reorg() {
 
 	// G -> A1 -> A2
 	//   \
-	//    -> B1 -> B2
+	//    -> B1 -> B2 -> B3
 
 	let mut finality_notifications = client.finality_notification_stream();
 
@@ -1250,8 +1250,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));
 }
 
@@ -1372,8 +1372,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));
 }
 
@@ -1602,9 +1606,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));
 }

From 6fbc756fa3eaee55cd9f6b06bb6fbd888ab638bf Mon Sep 17 00:00:00 2001
From: shamil-gadelshin <shamilgadelshin@gmail.com>
Date: Wed, 17 Apr 2024 21:23:58 +0700
Subject: [PATCH 10/19] Update substrate/primitives/blockchain/src/backend.rs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Bastian Köcher <git@kchr.de>
---
 substrate/primitives/blockchain/src/backend.rs | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/substrate/primitives/blockchain/src/backend.rs b/substrate/primitives/blockchain/src/backend.rs
index 309b3e265c91..ebbadbfba927 100644
--- a/substrate/primitives/blockchain/src/backend.rs
+++ b/substrate/primitives/blockchain/src/backend.rs
@@ -302,7 +302,7 @@ pub trait Backend<Block: BlockT>:
 	}
 }
 
-/// Calculation result of the displaced leaves after the block finalization.
+/// Result of  [`HeaderBackend::displaced_leaves_after_finalizing`].
 #[derive(Clone, Debug)]
 pub struct DisplacedLeavesAfterFinalization<Block: BlockT> {
 	/// A collection of hashes and block numbers for displaced leaves.
@@ -312,6 +312,12 @@ pub struct DisplacedLeavesAfterFinalization<Block: BlockT> {
 	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) -> Vec<Block::Hash> {

From af6afe4a866a55e692e322ba9d37af8cd88e71c7 Mon Sep 17 00:00:00 2001
From: shamil-gadelshin <shamilgadelshin@gmail.com>
Date: Wed, 17 Apr 2024 21:24:14 +0700
Subject: [PATCH 11/19] Update substrate/primitives/blockchain/src/backend.rs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Bastian Köcher <git@kchr.de>
---
 substrate/primitives/blockchain/src/backend.rs | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/substrate/primitives/blockchain/src/backend.rs b/substrate/primitives/blockchain/src/backend.rs
index ebbadbfba927..174dbf90f666 100644
--- a/substrate/primitives/blockchain/src/backend.rs
+++ b/substrate/primitives/blockchain/src/backend.rs
@@ -323,12 +323,6 @@ impl<Block: BlockT> DisplacedLeavesAfterFinalization<Block> {
 	pub fn hashes(&self) -> Vec<Block::Hash> {
 		self.displaced_leaves.keys().map(|h| *h).collect()
 	}
-
-	/// Constructor. We need this explicit initialization to not introduce a requirement
-	/// `Block: Default`
-	pub fn new() -> Self {
-		Self { displaced_leaves: BTreeMap::new(), tree_routes: BTreeMap::new() }
-	}
 }
 
 /// Blockchain info

From 75ad2200136a0b200a083d9f2b3210ec52e62bfc Mon Sep 17 00:00:00 2001
From: shamil-gadelshin <shamilgadelshin@gmail.com>
Date: Wed, 17 Apr 2024 21:24:28 +0700
Subject: [PATCH 12/19] Update substrate/primitives/blockchain/src/backend.rs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Bastian Köcher <git@kchr.de>
---
 substrate/primitives/blockchain/src/backend.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/substrate/primitives/blockchain/src/backend.rs b/substrate/primitives/blockchain/src/backend.rs
index 174dbf90f666..68ca5f3ff6ee 100644
--- a/substrate/primitives/blockchain/src/backend.rs
+++ b/substrate/primitives/blockchain/src/backend.rs
@@ -269,7 +269,7 @@ pub trait Backend<Block: BlockT>:
 		finalized_block_hash: Block::Hash,
 		finalized_block_number: NumberFor<Block>,
 	) -> std::result::Result<DisplacedLeavesAfterFinalization<Block>, Error> {
-		let mut result = DisplacedLeavesAfterFinalization::new();
+		let mut result = DisplacedLeavesAfterFinalization::default();
 
 		if finalized_block_number == Zero::zero() {
 			return Ok(result)

From 027739fc0b43bdb797a62c004bf7441594afff8e Mon Sep 17 00:00:00 2001
From: Shamil Gadelshin <shamilgadelshin@gmail.com>
Date: Wed, 17 Apr 2024 23:28:50 +0700
Subject: [PATCH 13/19] Apply review suggestions.

- update prdoc file
- update test comments
- minor refactoring
---
 prdoc/pr_3962.prdoc                                  | 1 +
 substrate/client/rpc-spec-v2/src/chain_head/tests.rs | 4 ++--
 substrate/client/service/src/client/client.rs        | 3 ++-
 substrate/primitives/blockchain/src/backend.rs       | 4 ++--
 4 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/prdoc/pr_3962.prdoc b/prdoc/pr_3962.prdoc
index 3d5a878b8987..7ef59d38ce5c 100644
--- a/prdoc/pr_3962.prdoc
+++ b/prdoc/pr_3962.prdoc
@@ -4,6 +4,7 @@ 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
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 3174b91ad2b2..c7525fd07b0d 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
+++ b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
@@ -2429,6 +2429,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![
@@ -2446,7 +2447,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())
@@ -2477,7 +2477,7 @@ 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)],
diff --git a/substrate/client/service/src/client/client.rs b/substrate/client/service/src/client/client.rs
index f721b64f6c53..3c25c233775b 100644
--- a/substrate/client/service/src/client/client.rs
+++ b/substrate/client/service/src/client/client.rs
@@ -982,7 +982,8 @@ where
 				.backend
 				.blockchain()
 				.displaced_leaves_after_finalizing(hash, block_number)?
-				.hashes();
+				.hashes()
+				.collect();
 
 			let header = self
 				.backend
diff --git a/substrate/primitives/blockchain/src/backend.rs b/substrate/primitives/blockchain/src/backend.rs
index 68ca5f3ff6ee..765a08aba5dc 100644
--- a/substrate/primitives/blockchain/src/backend.rs
+++ b/substrate/primitives/blockchain/src/backend.rs
@@ -320,8 +320,8 @@ impl<Block: BlockT> Default for DisplacedLeavesAfterFinalization<Block> {
 
 impl<Block: BlockT> DisplacedLeavesAfterFinalization<Block> {
 	/// Returns a collection of hashes for the displaced leaves.
-	pub fn hashes(&self) -> Vec<Block::Hash> {
-		self.displaced_leaves.keys().map(|h| *h).collect()
+	pub fn hashes(&self) -> impl Iterator<Item = Block::Hash> + '_ {
+		self.displaced_leaves.keys().cloned()
 	}
 }
 

From 0ea2dfbbd5847c18a9685052642aeaec89c56494 Mon Sep 17 00:00:00 2001
From: Shamil Gadelshin <shamilgadelshin@gmail.com>
Date: Wed, 8 May 2024 18:21:29 +0700
Subject: [PATCH 14/19] Update follow_unique_pruned_blocks test.

- update `follow_unique_pruned_blocks` test from `sc-rpc-spec-v2` crate
---
 .../rpc-spec-v2/src/chain_head/tests.rs       | 25 +++++++++----------
 1 file changed, 12 insertions(+), 13 deletions(-)

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 dab5fde89053..d0323c4ff3b6 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
+++ b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
@@ -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,7 +3768,7 @@ async fn follow_unique_pruned_blocks() {
 	});
 	assert_eq!(event, expected);
 
-	// Block 2 must be reported as pruned, even if it was the previous best.
+	// All blocks from stale forks are pruned when we finalize block 6.
 	let event: FollowEvent<String> = get_next_event(&mut sub).await;
 	let expected = FollowEvent::Finalized(Finalized {
 		finalized_block_hashes: vec![
@@ -3784,7 +3776,12 @@ async fn follow_unique_pruned_blocks() {
 			format!("{:?}", block_2_f_hash),
 			format!("{:?}", block_6_hash),
 		],
-		pruned_block_hashes: vec![],
+		pruned_block_hashes: vec![
+			format!("{:?}", block_2_hash),
+			format!("{:?}", block_3_hash),
+			format!("{:?}", block_4_hash),
+			format!("{:?}", block_5_hash),
+		],
 	});
 	assert_eq!(event, expected);
 
@@ -3802,9 +3799,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 +3813,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);
 }

From 98e47bd0f33f7a05bc4966e2c74f87c8885964ad Mon Sep 17 00:00:00 2001
From: Shamil Gadelshin <shamilgadelshin@gmail.com>
Date: Wed, 8 May 2024 18:48:03 +0700
Subject: [PATCH 15/19] Update the expand_forks() comment.

---
 substrate/primitives/blockchain/src/backend.rs | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/substrate/primitives/blockchain/src/backend.rs b/substrate/primitives/blockchain/src/backend.rs
index 765a08aba5dc..11160c81944a 100644
--- a/substrate/primitives/blockchain/src/backend.rs
+++ b/substrate/primitives/blockchain/src/backend.rs
@@ -100,6 +100,16 @@ pub trait ForkBackend<Block: BlockT>:
 	/// 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.
+	///
+	/// The first part (default path) of the function utilizes `tree_route` function and resorts to
+	/// the another algorithm when the default path fails. The difference between two is that
+	/// `tree_route` works with block hashes and the second algorithm works with block
+	/// numbers: in some cases with finalized blocks the result may differ.
+	/// Example:
+	///  G --- A1 --- A2 --- A3 --- A4           ( < fork1 )
+	///                       \-----C4 --- C5    ( < fork2 )
+	/// We finalize A3 and call expand_fork(C5). The first part of the algorithm will return (C5,
+	/// C4) and the second algorithm will return (C5).
 	fn expand_forks(
 		&self,
 		fork_heads: &[Block::Hash],

From f76be389cc39cb7ffbc903b2168be1cc92e33dec Mon Sep 17 00:00:00 2001
From: Shamil Gadelshin <shamilgadelshin@gmail.com>
Date: Fri, 10 May 2024 17:55:59 +0700
Subject: [PATCH 16/19] Update expand_forks().

---
 .../primitives/blockchain/src/backend.rs      | 64 ++-----------------
 1 file changed, 6 insertions(+), 58 deletions(-)

diff --git a/substrate/primitives/blockchain/src/backend.rs b/substrate/primitives/blockchain/src/backend.rs
index 11160c81944a..a12634374503 100644
--- a/substrate/primitives/blockchain/src/backend.rs
+++ b/substrate/primitives/blockchain/src/backend.rs
@@ -21,7 +21,7 @@ use log::warn;
 use parking_lot::RwLock;
 use sp_runtime::{
 	generic::BlockId,
-	traits::{Block as BlockT, Header as HeaderT, NumberFor, Saturating, Zero},
+	traits::{Block as BlockT, Header as HeaderT, NumberFor, Zero},
 	Justifications,
 };
 use std::collections::{btree_map::BTreeMap, btree_set::BTreeSet};
@@ -92,29 +92,17 @@ 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.
-	///
-	/// The first part (default path) of the function utilizes `tree_route` function and resorts to
-	/// the another algorithm when the default path fails. The difference between two is that
-	/// `tree_route` works with block hashes and the second algorithm works with block
-	/// numbers: in some cases with finalized blocks the result may differ.
 	/// Example:
 	///  G --- A1 --- A2 --- A3 --- A4           ( < fork1 )
 	///                       \-----C4 --- C5    ( < fork2 )
-	/// We finalize A3 and call expand_fork(C5). The first part of the algorithm will return (C5,
-	/// C4) and the second algorithm will return (C5).
+	/// 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 {
 			match tree_route(self, *fork_head, self.info().finalized_hash) {
@@ -125,49 +113,9 @@ pub trait ForkBackend<Block: BlockT>:
 					continue
 				},
 				Err(_) => {
-					// Continue with fallback algorithm
+					// There are cases when blocks are missing (e.g. warp-sync).
 				},
 			}
-
-			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
-					},
-				}
-			}
-		}
-
-		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)

From e56eed69a0009b73504dea8f6771f6ce7b39f5f5 Mon Sep 17 00:00:00 2001
From: Shamil Gadelshin <shamilgadelshin@gmail.com>
Date: Fri, 10 May 2024 17:56:42 +0700
Subject: [PATCH 17/19] Update expand_forks() dependencies.

---
 substrate/client/consensus/babe/src/lib.rs          |  5 +++--
 .../merkle-mountain-range/src/offchain_mmr.rs       | 13 ++++++-------
 2 files changed, 9 insertions(+), 9 deletions(-)

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/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);
 		}

From c875bd2ed0a7e8c0fef967b734a27b13aa8d24e5 Mon Sep 17 00:00:00 2001
From: Shamil Gadelshin <shamilgadelshin@gmail.com>
Date: Fri, 10 May 2024 19:59:36 +0700
Subject: [PATCH 18/19] Update flaky test.

---
 .../rpc-spec-v2/src/chain_head/tests.rs       | 26 ++++++++++++++-----
 1 file changed, 19 insertions(+), 7 deletions(-)

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 d0323c4ff3b6..b195e05b6649 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
+++ b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
@@ -3769,20 +3769,32 @@ async fn follow_unique_pruned_blocks() {
 	assert_eq!(event, expected);
 
 	// All blocks from stale forks are pruned when we finalize block 6.
-	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	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![
-			format!("{:?}", block_2_hash),
-			format!("{:?}", block_3_hash),
-			format!("{:?}", block_4_hash),
-			format!("{:?}", block_5_hash),
-		],
+		pruned_block_hashes: expected_pruned_block_hashes,
 	});
+
 	assert_eq!(event, expected);
 
 	// Pruned hash can be unpinned.

From 7e769f687c524e95efb59978cb35ad9a86a612de Mon Sep 17 00:00:00 2001
From: Shamil Gadelshin <shamilgadelshin@gmail.com>
Date: Wed, 15 May 2024 14:08:49 +0700
Subject: [PATCH 19/19] Fix doc-comment.

---
 substrate/primitives/blockchain/src/backend.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/substrate/primitives/blockchain/src/backend.rs b/substrate/primitives/blockchain/src/backend.rs
index a12634374503..06e5b682964a 100644
--- a/substrate/primitives/blockchain/src/backend.rs
+++ b/substrate/primitives/blockchain/src/backend.rs
@@ -260,7 +260,7 @@ pub trait Backend<Block: BlockT>:
 	}
 }
 
-/// Result of  [`HeaderBackend::displaced_leaves_after_finalizing`].
+/// 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.