feat: add state pruning mechanism to sync pipeline

This commit introduces a new pruning mechanism for managing historical
state in the sync pipeline. The system allows users to control state
retention with flexible pruning modes.

## Changes

### Stage Trait Enhancement
- Added `Stage::prune` method to enable stages to prune historical data
- Introduced `PruneInput` and `PruneOutput` types for pruning operations
- Added `PruningMode` enum supporting three modes:
  - `Full`: Keep all historical state (no pruning)
  - `HistoricalBlocks(n)`: Keep last N blocks of historical state
  - `LatestOnly`: Keep only the latest state

### Pipeline Integration
- Added `PruningConfig` to configure pruning behavior
- Integrated pruning into the pipeline loop with configurable intervals
- Pruning runs automatically after processing a specified number of blocks
- Added `should_prune()` and `prune()` methods to Pipeline

### Stage Implementations
- **StateTrie**: Added no-op prune implementation (TODO for provider API)
- **Blocks**: Added no-op prune implementation (TODO for provider API)
- **Classes**: Added no-op prune with rationale (classes are immutable)

## Design Notes

The pruning mechanism is designed with a single, unified API that handles
all three pruning scenarios. The `PruneInput::prune_before()` method
calculates the cutoff block based on the pruning mode, making it easy
for stages to determine what to prune.

For now, stage implementations are placeholders (no-ops) as actual
pruning requires provider API extensions. The infrastructure is in
place for future implementation.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

Author: kariy

crates/sync/pipeline/src/lib.rs

@@ -80,7 +80,7 @@ use futures::future::BoxFuture;
 use katana_primitives::block::BlockNumber;
 use katana_provider_api::stage::StageCheckpointProvider;
 use katana_provider_api::ProviderError;
-use katana_stage::{Stage, StageExecutionInput, StageExecutionOutput};
+use katana_stage::{PruneInput, PruneOutput, PruningMode, Stage, StageExecutionInput, StageExecutionOutput};
 use tokio::sync::watch;
 use tokio::task::yield_now;
 use tracing::{debug, error, info, info_span, Instrument};
@@ -99,6 +99,9 @@ pub enum Error {
     #[error("stage {id} execution failed: {error}")]
     StageExecution { id: &'static str, error: katana_stage::Error },
 
+    #[error("stage {id} pruning failed: {error}")]
+    StagePruning { id: &'static str, error: katana_stage::Error },
+
     #[error(transparent)]
     Provider(#[from] ProviderError),
 }
@@ -199,6 +202,35 @@ impl PipelineHandle {
     }
 }
 
+/// Configuration for pruning behavior in the pipeline.
+#[derive(Debug, Clone)]
+pub struct PruningConfig {
+    /// The pruning mode to use.
+    pub mode: PruningMode,
+    /// How many blocks to process between pruning runs.
+    /// Pruning will be triggered after every `interval` blocks are synced.
+    /// If `None`, pruning is disabled.
+    pub interval: Option<u64>,
+}
+
+impl Default for PruningConfig {
+    fn default() -> Self {
+        Self { mode: PruningMode::Full, interval: None }
+    }
+}
+
+impl PruningConfig {
+    /// Creates a new pruning configuration with the specified mode and interval.
+    pub fn new(mode: PruningMode, interval: Option<u64>) -> Self {
+        Self { mode, interval }
+    }
+
+    /// Returns whether pruning is enabled.
+    pub fn is_enabled(&self) -> bool {
+        !matches!(self.mode, PruningMode::Full) && self.interval.is_some()
+    }
+}
+
 /// Syncing pipeline.
 ///
 /// The pipeline drives the execution of stages, running each stage to completion in the order they
@@ -221,6 +253,8 @@ pub struct Pipeline<P> {
     command_tx: watch::Sender<Option<PipelineCommand>>,
     block_tx: watch::Sender<Option<BlockNumber>>,
     tip: Option<BlockNumber>,
+    pruning_config: PruningConfig,
+    last_pruned_block: Option<BlockNumber>,
 }
 
 impl<P> Pipeline<P> {
@@ -246,10 +280,24 @@ impl<P> Pipeline<P> {
             provider,
             chunk_size,
             tip: None,
+            pruning_config: PruningConfig::default(),
+            last_pruned_block: None,
         };
         (pipeline, handle)
     }
 
+    /// Sets the pruning configuration for the pipeline.
+    ///
+    /// This controls how and when historical state is pruned during synchronization.
+    pub fn set_pruning_config(&mut self, config: PruningConfig) {
+        self.pruning_config = config;
+    }
+
+    /// Returns the current pruning configuration.
+    pub fn pruning_config(&self) -> &PruningConfig {
+        &self.pruning_config
+    }
+
     /// Adds a new stage to the end of the pipeline.
     ///
     /// Stages are executed in the order they are added.
@@ -418,6 +466,13 @@ impl<P: StageCheckpointProvider> Pipeline<P> {
                 // Notify subscribers about the newly processed block
                 let _ = self.block_tx.send(Some(last_block_processed));
 
+                // Check if we should run pruning
+                if self.should_prune(last_block_processed) {
+                    info!(target: "pipeline", block = %last_block_processed, "Starting pruning.");
+                    self.prune(last_block_processed).await?;
+                    self.last_pruned_block = Some(last_block_processed);
+                }
+
                 if last_block_processed >= tip {
                     info!(target: "pipeline", %tip, "Finished syncing until tip.");
                     self.tip = None;
@@ -440,6 +495,52 @@ impl<P: StageCheckpointProvider> Pipeline<P> {
             yield_now().await;
         }
     }
+
+    /// Determines if pruning should be performed based on the current block and configuration.
+    fn should_prune(&self, current_block: BlockNumber) -> bool {
+        if !self.pruning_config.is_enabled() {
+            return false;
+        }
+
+        let interval = match self.pruning_config.interval {
+            Some(i) => i,
+            None => return false,
+        };
+
+        match self.last_pruned_block {
+            Some(last_pruned) => current_block.saturating_sub(last_pruned) >= interval,
+            None => current_block >= interval,
+        }
+    }
+
+    /// Runs pruning on all stages.
+    async fn prune(&mut self, tip: BlockNumber) -> PipelineResult<()> {
+        if self.stages.is_empty() {
+            return Ok(());
+        }
+
+        let prune_input = PruneInput::new(tip, self.pruning_config.mode);
+
+        for stage in self.stages.iter_mut() {
+            let id = stage.id();
+            let span = info_span!(target: "pipeline", "stage.prune", stage = %id, %tip);
+            let enter = span.entered();
+
+            info!(target: "pipeline", mode = ?self.pruning_config.mode, "Pruning stage.");
+
+            let span_inner = enter.exit();
+            let PruneOutput { pruned_count } = stage
+                .prune(&prune_input)
+                .instrument(span_inner.clone())
+                .await
+                .map_err(|error| Error::StagePruning { id, error })?;
+
+            let _enter = span_inner.enter();
+            info!(target: "pipeline", %pruned_count, "Stage pruning completed.");
+        }
+
+        Ok(())
+    }
 }
 
 impl<P> IntoFuture for Pipeline<P>
@@ -467,6 +568,8 @@ where
             .field("command", &self.command_rx)
             .field("provider", &self.provider)
             .field("chunk_size", &self.chunk_size)
+            .field("pruning_config", &self.pruning_config)
+            .field("last_pruned_block", &self.last_pruned_block)
             .field("stages", &self.stages.iter().map(|s| s.id()).collect::<Vec<_>>())
             .finish()
     }

crates/sync/stage/src/blocks/mod.rs

@@ -17,7 +17,7 @@ use num_traits::ToPrimitive;
 use starknet::core::types::ResourcePrice;
 use tracing::{error, info_span, Instrument};
 
-use crate::{Stage, StageExecutionInput, StageExecutionOutput, StageResult};
+use crate::{PruneInput, PruneOutput, PruneResult, Stage, StageExecutionInput, StageExecutionOutput, StageResult};
 
 mod downloader;
 
@@ -133,6 +133,17 @@ where
             Ok(StageExecutionOutput { last_block_processed: input.to() })
         })
     }
+
+    fn prune<'a>(&'a mut self, _input: &'a PruneInput) -> BoxFuture<'a, PruneResult> {
+        Box::pin(async move {
+            // TODO: Implement block pruning once the BlockWriter provider API supports it.
+            // For now, this is a no-op. Block pruning would involve:
+            // 1. Determining which blocks to prune based on input.prune_before()
+            // 2. Removing block headers, bodies, receipts, and state updates for those blocks
+            // 3. Ensuring the chain tip and recent blocks remain intact
+            Ok(PruneOutput::default())
+        })
+    }
 }
 
 #[derive(Debug, thiserror::Error)]

crates/sync/stage/src/classes.rs

@@ -14,7 +14,7 @@ use katana_rpc_types::class::ConversionError;
 use rayon::prelude::*;
 use tracing::{debug, error, info_span, Instrument};
 
-use super::{Stage, StageExecutionInput, StageExecutionOutput, StageResult};
+use super::{PruneInput, PruneOutput, PruneResult, Stage, StageExecutionInput, StageExecutionOutput, StageResult};
 use crate::downloader::{BatchDownloader, Downloader, DownloaderResult};
 
 /// A stage for downloading and storing contract classes.
@@ -158,6 +158,15 @@ where
             Ok(StageExecutionOutput { last_block_processed: input.to() })
         })
     }
+
+    fn prune<'a>(&'a mut self, _input: &'a PruneInput) -> BoxFuture<'a, PruneResult> {
+        Box::pin(async move {
+            // Classes are immutable once declared and don't need pruning.
+            // A class declared at block N can still be used to deploy contracts at block N+1000.
+            // Therefore, we cannot safely prune classes based on block age alone.
+            Ok(PruneOutput::default())
+        })
+    }
 }
 
 #[derive(Debug, thiserror::Error)]

crates/sync/stage/src/lib.rs

@@ -18,6 +18,9 @@ pub use trie::StateTrie;
 /// The result type of a stage execution. See [Stage::execute].
 pub type StageResult = Result<StageExecutionOutput, Error>;
 
+/// The result type of a stage pruning. See [Stage::prune].
+pub type PruneResult = Result<PruneOutput, Error>;
+
 /// Input parameters for stage execution.
 ///
 /// # Invariant
@@ -62,6 +65,76 @@ pub struct StageExecutionOutput {
     pub last_block_processed: BlockNumber,
 }
 
+/// Pruning mode configuration that determines how much historical state to retain.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum PruningMode {
+    /// Keep all historical state since genesis (no pruning).
+    Full,
+    /// Keep only the last N blocks of historical state.
+    HistoricalBlocks(u64),
+    /// Keep only the latest state, pruning all historical data.
+    LatestOnly,
+}
+
+impl Default for PruningMode {
+    fn default() -> Self {
+        Self::Full
+    }
+}
+
+/// Input parameters for stage pruning.
+#[derive(Debug, Clone)]
+pub struct PruneInput {
+    /// The current tip of the chain (highest synced block).
+    tip: BlockNumber,
+    /// The pruning mode to apply.
+    mode: PruningMode,
+}
+
+impl PruneInput {
+    /// Creates a new [`PruneInput`] with the given tip and pruning mode.
+    pub fn new(tip: BlockNumber, mode: PruningMode) -> Self {
+        Self { tip, mode }
+    }
+
+    /// Returns the current chain tip.
+    #[inline]
+    pub fn tip(&self) -> BlockNumber {
+        self.tip
+    }
+
+    /// Returns the pruning mode.
+    #[inline]
+    pub fn mode(&self) -> PruningMode {
+        self.mode
+    }
+
+    /// Calculates the block number before which all state should be pruned.
+    ///
+    /// Returns `None` if no pruning should occur (e.g., in `Full` mode).
+    /// Returns `Some(block_number)` indicating that all state before this block can be pruned.
+    pub fn prune_before(&self) -> Option<BlockNumber> {
+        match self.mode {
+            PruningMode::Full => None,
+            PruningMode::LatestOnly => Some(self.tip.saturating_sub(1)),
+            PruningMode::HistoricalBlocks(n) => {
+                if self.tip > n {
+                    Some(self.tip - n)
+                } else {
+                    None
+                }
+            }
+        }
+    }
+}
+
+/// Output from a stage pruning operation.
+#[derive(Debug, Default)]
+pub struct PruneOutput {
+    /// The number of items (blocks, state entries, etc.) that were pruned.
+    pub pruned_count: u64,
+}
+
 #[derive(Debug, thiserror::Error)]
 pub enum Error {
     #[error(transparent)]
@@ -96,9 +169,10 @@ pub enum Error {
 ///
 /// # Implementation Note
 ///
-/// The [`execute`](Stage::execute) method returns a [`BoxFuture`] instead of `impl Future` to
-/// maintain dyn-compatibility. This allows the pipeline to store different stage implementations
-/// in a `Vec<Box<dyn Stage>>`, enabling dynamic composition of sync stages at runtime.
+/// The [`execute`](Stage::execute) and [`prune`](Stage::prune) methods return a [`BoxFuture`]
+/// instead of `impl Future` to maintain dyn-compatibility. This allows the pipeline to store
+/// different stage implementations in a `Vec<Box<dyn Stage>>`, enabling dynamic composition of
+/// sync stages at runtime.
 ///
 /// While this introduces a small heap allocation for the future, it's negligible compared to
 /// the actual async work performed by stages (network I/O, database operations, etc.).
@@ -126,16 +200,72 @@ pub trait Stage: Send + Sync {
     /// Implementors are expected to perform any necessary processings on all blocks in the range
     /// `[input.from, input.to]`.
     fn execute<'a>(&'a mut self, input: &'a StageExecutionInput) -> BoxFuture<'a, StageResult>;
+
+    /// Prunes historical data for this stage according to the pruning configuration.
+    ///
+    /// This method is called by the pipeline to remove old historical state that is no longer
+    /// needed according to the pruning mode. The pruning operation is non-blocking and runs
+    /// asynchronously.
+    ///
+    /// # Arguments
+    ///
+    /// * `input` - The pruning input containing the current chain tip and pruning mode
+    ///
+    /// # Returns
+    ///
+    /// A future that resolves to a [`PruneResult`] containing [`PruneOutput`] with the
+    /// number of items that were pruned.
+    ///
+    /// # Implementation Notes
+    ///
+    /// - Stages that don't store historical state (e.g., Classes) can provide a no-op
+    ///   implementation that returns `Ok(PruneOutput::default())`.
+    /// - Stages that store state (e.g., Blocks, StateTrie) should implement pruning logic
+    ///   appropriate to their data model.
+    /// - The pruning operation must be non-blocking, just like [`execute`](Stage::execute).
+    /// - Implementors should use [`PruneInput::prune_before`] to determine which blocks to prune.
+    fn prune<'a>(&'a mut self, input: &'a PruneInput) -> BoxFuture<'a, PruneResult>;
 }
 
 #[cfg(test)]
 mod tests {
-    use crate::StageExecutionInput;
+    use crate::{PruneInput, PruningMode, StageExecutionInput};
 
     #[tokio::test]
     #[should_panic(expected = "Invalid block range")]
     async fn invalid_range_panics() {
         // When from > to, the range is invalid and should panic at construction time
         let _ = StageExecutionInput::new(100, 99);
     }
+
+    #[test]
+    fn prune_before_full_mode() {
+        let input = PruneInput::new(1000, PruningMode::Full);
+        assert_eq!(input.prune_before(), None);
+    }
+
+    #[test]
+    fn prune_before_latest_only() {
+        let input = PruneInput::new(1000, PruningMode::LatestOnly);
+        assert_eq!(input.prune_before(), Some(999));
+
+        // Edge case: tip at block 0
+        let input = PruneInput::new(0, PruningMode::LatestOnly);
+        assert_eq!(input.prune_before(), Some(0));
+    }
+
+    #[test]
+    fn prune_before_historical_blocks() {
+        // Keep last 100 blocks, tip at 1000
+        let input = PruneInput::new(1000, PruningMode::HistoricalBlocks(100));
+        assert_eq!(input.prune_before(), Some(900));
+
+        // Keep last 100 blocks, tip at 50 (not enough blocks yet)
+        let input = PruneInput::new(50, PruningMode::HistoricalBlocks(100));
+        assert_eq!(input.prune_before(), None);
+
+        // Keep last 100 blocks, tip at exactly 100
+        let input = PruneInput::new(100, PruningMode::HistoricalBlocks(100));
+        assert_eq!(input.prune_before(), Some(0));
+    }
 }