From 27bb35c61b02b9ab22245e7a7ae8e553c8cc14fd Mon Sep 17 00:00:00 2001
From: joshieDo <93316087+joshieDo@users.noreply.github.com>
Date: Tue, 23 Dec 2025 14:13:39 +0000
Subject: [PATCH] iterate over active stages only

---
 Cargo.lock                                          |  1 +
 crates/node/builder/src/launch/common.rs            | 10 +++++-----
 crates/stages/types/Cargo.toml                      |  2 ++
 crates/stages/types/src/id.rs                       | 13 +++++++++++++
 .../provider/src/providers/database/provider.rs     |  6 ++++--
 5 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index d6e7a6b94c..5ae3fbc86f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -10734,6 +10734,7 @@ dependencies = [
  "proptest-arbitrary-interop",
  "rand 0.9.2",
  "reth-codecs",
+ "reth-prune-types",
  "reth-trie-common",
  "serde",
 ]
diff --git a/crates/node/builder/src/launch/common.rs b/crates/node/builder/src/launch/common.rs
index 7f0e944c92..e59627521d 100644
--- a/crates/node/builder/src/launch/common.rs
+++ b/crates/node/builder/src/launch/common.rs
@@ -938,13 +938,13 @@ where
     ///
     /// A target block hash if the pipeline is inconsistent, otherwise `None`.
     pub fn check_pipeline_consistency(&self) -> ProviderResult<Option<B256>> {
-        // We skip the era stage if it's not enabled
+        // Only check stages that are active based on configuration
         let era_enabled = self.era_import_source().is_some();
-        let mut all_stages =
-            StageId::ALL.into_iter().filter(|id| era_enabled || id != &StageId::Era);
+        let prune_modes = self.prune_modes();
+        let mut active_stages = StageId::active(era_enabled, &prune_modes);
 
         // Get the expected first stage based on config.
-        let first_stage = all_stages.next().expect("there must be at least one stage");
+        let first_stage = active_stages.next().expect("there must be at least one stage");
 
         // If no target was provided, check if the stages are congruent - check if the
         // checkpoint of the last stage matches the checkpoint of the first.
@@ -955,7 +955,7 @@ where
             .block_number;
 
         // Compare all other stages against the first
-        for stage_id in all_stages {
+        for stage_id in active_stages {
             let stage_checkpoint = self
                 .blockchain_db()
                 .get_stage_checkpoint(stage_id)?
diff --git a/crates/stages/types/Cargo.toml b/crates/stages/types/Cargo.toml
index 6e70fbe26a..2712f86a87 100644
--- a/crates/stages/types/Cargo.toml
+++ b/crates/stages/types/Cargo.toml
@@ -13,6 +13,7 @@ workspace = true
 
 [dependencies]
 reth-codecs = { workspace = true, optional = true }
+reth-prune-types.workspace = true
 reth-trie-common.workspace = true
 alloy-primitives.workspace = true
 
@@ -38,6 +39,7 @@ default = ["std"]
 std = [
     "alloy-primitives/std",
     "bytes?/std",
+    "reth-prune-types/std",
     "reth-trie-common/std",
     "serde?/std",
 ]
diff --git a/crates/stages/types/src/id.rs b/crates/stages/types/src/id.rs
index 8c0a91c873..30897bb760 100644
--- a/crates/stages/types/src/id.rs
+++ b/crates/stages/types/src/id.rs
@@ -1,4 +1,5 @@
 use alloc::vec::Vec;
+use reth_prune_types::PruneModes;
 #[cfg(feature = "std")]
 use std::{collections::HashMap, sync::OnceLock};
 
@@ -112,6 +113,18 @@ impl StageId {
         matches!(self, Self::Finish)
     }
 
+    /// Returns an iterator over stages that are active based on configuration.
+    ///
+    /// Optional stages like [`StageId::Era`] and [`StageId::PruneSenderRecovery`] are only
+    /// included when their corresponding configuration is enabled.
+    pub fn active(era_enabled: bool, prune_modes: &PruneModes) -> impl Iterator<Item = Self> {
+        Self::ALL.into_iter().filter(move |id| match id {
+            Self::Era => era_enabled,
+            Self::PruneSenderRecovery => prune_modes.sender_recovery.is_some(),
+            _ => true,
+        })
+    }
+
     /// Get a pre-encoded raw Vec, for example, to be used as the DB key for
     /// `tables::StageCheckpoints` and `tables::StageCheckpointProgresses`
     pub fn get_pre_encoded(&self) -> Option<&Vec<u8>> {
diff --git a/crates/storage/provider/src/providers/database/provider.rs b/crates/storage/provider/src/providers/database/provider.rs
index 9b4b68a606..cf416ce851 100644
--- a/crates/storage/provider/src/providers/database/provider.rs
+++ b/crates/storage/provider/src/providers/database/provider.rs
@@ -1563,9 +1563,11 @@ impl<TX: DbTxMut, N: NodeTypes> StageCheckpointWriter for DatabaseProvider<TX, N
         block_number: BlockNumber,
         drop_stage_checkpoint: bool,
     ) -> ProviderResult<()> {
-        // iterate over all existing stages in the table and update its progress.
+        // iterate over active stages and update their progress.
         let mut cursor = self.tx.cursor_write::<tables::StageCheckpoints>()?;
-        for stage_id in StageId::ALL {
+        // TODO: era_enabled should come from node config
+        let era_enabled = true;
+        for stage_id in StageId::active(era_enabled, self.prune_modes_ref()) {
             let (_, checkpoint) = cursor.seek_exact(stage_id.to_string())?.unwrap_or_default();
             cursor.upsert(
                 stage_id.to_string(),