Skip to main content

reth_prune_types/
mode.rs

1use crate::{segment::PrunePurpose, PruneSegment, PruneSegmentError};
2use alloy_primitives::BlockNumber;
3
4/// Prune mode.
5#[derive(Debug, PartialEq, Eq, Clone, Copy)]
6#[cfg_attr(any(test, feature = "test-utils"), derive(arbitrary::Arbitrary))]
7#[cfg_attr(any(test, feature = "reth-codec"), derive(reth_codecs::Compact))]
8#[cfg_attr(any(test, feature = "reth-codec"), reth_codecs::add_arbitrary_tests(compact))]
9#[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize, serde::Deserialize))]
10#[cfg_attr(any(test, feature = "serde"), serde(rename_all = "lowercase"))]
11pub enum PruneMode {
12    /// Prune all blocks.
13    Full,
14    /// Prune blocks before the `head-N` block number. In other words, keep last N + 1 blocks.
15    Distance(u64),
16    /// Prune blocks before the specified block number. The specified block number is not pruned.
17    Before(BlockNumber),
18}
19
20#[cfg(any(test, feature = "test-utils"))]
21#[expect(clippy::derivable_impls)]
22impl Default for PruneMode {
23    fn default() -> Self {
24        Self::Full
25    }
26}
27
28impl PruneMode {
29    /// Prune blocks up to the specified block number. The specified block number is also pruned.
30    ///
31    /// This acts as `PruneMode::Before(block_number + 1)`.
32    pub const fn before_inclusive(block_number: BlockNumber) -> Self {
33        Self::Before(block_number + 1)
34    }
35
36    /// Returns block up to which variant pruning needs to be done, inclusive, according to the
37    /// provided tip.
38    pub fn prune_target_block(
39        &self,
40        tip: BlockNumber,
41        segment: PruneSegment,
42        purpose: PrunePurpose,
43    ) -> Result<Option<(BlockNumber, Self)>, PruneSegmentError> {
44        self.prune_target_block_with_min(tip, segment, purpose, None)
45    }
46
47    /// Like [`prune_target_block`](Self::prune_target_block), but accepts an optional
48    /// `min_blocks_override` that replaces the segment's default minimum.
49    pub fn prune_target_block_with_min(
50        &self,
51        tip: BlockNumber,
52        segment: PruneSegment,
53        purpose: PrunePurpose,
54        min_blocks_override: Option<u64>,
55    ) -> Result<Option<(BlockNumber, Self)>, PruneSegmentError> {
56        let min_blocks = min_blocks_override.unwrap_or_else(|| segment.min_blocks());
57        let result = match self {
58            Self::Full if min_blocks == 0 => Some((tip, *self)),
59            // For segments with min_blocks > 0, Full mode behaves like Distance(min_blocks)
60            Self::Full if min_blocks <= tip => Some((tip - min_blocks, *self)),
61            Self::Full => None, // Nothing to prune yet
62            Self::Distance(distance) if *distance > tip => None, // Nothing to prune yet
63            Self::Distance(distance) if *distance >= min_blocks => Some((tip - distance, *self)),
64            Self::Before(n) if *n == tip + 1 && purpose.is_static_file() => Some((tip, *self)),
65            Self::Before(n) if *n == 0 || *n > tip => None,
66            Self::Before(n) => (tip - n >= min_blocks).then(|| (*n - 1, *self)),
67            _ => return Err(PruneSegmentError::Configuration(segment)),
68        };
69        Ok(result)
70    }
71
72    /// Check if target block should be pruned according to the provided prune mode and tip.
73    pub const fn should_prune(&self, block: BlockNumber, tip: BlockNumber) -> bool {
74        match self {
75            Self::Full => true,
76            Self::Distance(distance) => {
77                if *distance > tip {
78                    return false
79                }
80                block < tip - *distance
81            }
82            Self::Before(n) => *n > block,
83        }
84    }
85
86    /// Returns true if the prune mode is [`PruneMode::Full`].
87    pub const fn is_full(&self) -> bool {
88        matches!(self, Self::Full)
89    }
90
91    /// Returns true if the prune mode is [`PruneMode::Distance`].
92    pub const fn is_distance(&self) -> bool {
93        matches!(self, Self::Distance(_))
94    }
95
96    /// Returns the next block number that will EVENTUALLY be pruned after the given checkpoint. It
97    /// should not be used to find if there are blocks to be pruned right now. For that, use
98    /// [`Self::prune_target_block`].
99    ///
100    /// This is independent of the current tip and indicates what block is next in the pruning
101    /// sequence according to this mode's configuration. Returns `None` if no more blocks will
102    /// be pruned (i.e., the mode has reached its target).
103    ///
104    /// # Examples
105    ///
106    /// - `Before(10)` with checkpoint at block 5 returns `Some(6)`
107    /// - `Before(10)` with checkpoint at block 9 returns `None` (done)
108    /// - `Distance(100)` with checkpoint at block 1000 returns `Some(1001)` (always has more)
109    /// - `Full` always returns the next block after checkpoint
110    pub const fn next_pruned_block(&self, checkpoint: Option<BlockNumber>) -> Option<BlockNumber> {
111        let next = match checkpoint {
112            Some(c) => c + 1,
113            None => 0,
114        };
115
116        match self {
117            Self::Before(n) => {
118                if next < *n {
119                    Some(next)
120                } else {
121                    None
122                }
123            }
124            Self::Distance(_) | Self::Full => Some(next),
125        }
126    }
127}
128
129#[cfg(test)]
130mod tests {
131    use crate::{PruneMode, PrunePurpose, PruneSegment, MINIMUM_UNWIND_SAFE_DISTANCE};
132    use assert_matches::assert_matches;
133    use serde::Deserialize;
134
135    #[test]
136    fn test_prune_target_block() {
137        let tip = 20000;
138        let segment = PruneSegment::AccountHistory;
139
140        let tests = vec![
141            // Full mode with min_blocks > 0 behaves like Distance(min_blocks)
142            (PruneMode::Full, Ok(Some(tip - segment.min_blocks()))),
143            // Nothing to prune
144            (PruneMode::Distance(tip + 1), Ok(None)),
145            (
146                PruneMode::Distance(segment.min_blocks() + 1),
147                Ok(Some(tip - (segment.min_blocks() + 1))),
148            ),
149            // Nothing to prune
150            (PruneMode::Before(tip + 1), Ok(None)),
151            (
152                PruneMode::Before(tip - MINIMUM_UNWIND_SAFE_DISTANCE),
153                Ok(Some(tip - MINIMUM_UNWIND_SAFE_DISTANCE - 1)),
154            ),
155            (
156                PruneMode::Before(tip - MINIMUM_UNWIND_SAFE_DISTANCE - 1),
157                Ok(Some(tip - MINIMUM_UNWIND_SAFE_DISTANCE - 2)),
158            ),
159            // Nothing to prune
160            (PruneMode::Before(tip - 1), Ok(None)),
161            // Nothing is before genesis
162            (PruneMode::Before(0), Ok(None)),
163            (PruneMode::Before(1), Ok(Some(0))),
164        ];
165
166        for (index, (mode, expected_result)) in tests.into_iter().enumerate() {
167            assert_eq!(
168                mode.prune_target_block(tip, segment, PrunePurpose::User),
169                expected_result.map(|r| r.map(|b| (b, mode))),
170                "Test {} failed",
171                index + 1,
172            );
173        }
174
175        // Test for a scenario where there are no minimum blocks and Full can be used
176        assert_eq!(
177            PruneMode::Full.prune_target_block(
178                tip,
179                PruneSegment::TransactionLookup,
180                PrunePurpose::User
181            ),
182            Ok(Some((tip, PruneMode::Full))),
183        );
184    }
185
186    #[test]
187    fn test_should_prune() {
188        let tip = 20000;
189        let should_prune = true;
190
191        let tests = vec![
192            (PruneMode::Distance(tip + 1), 1, !should_prune),
193            (
194                PruneMode::Distance(MINIMUM_UNWIND_SAFE_DISTANCE + 1),
195                tip - MINIMUM_UNWIND_SAFE_DISTANCE - 1,
196                !should_prune,
197            ),
198            (
199                PruneMode::Distance(MINIMUM_UNWIND_SAFE_DISTANCE + 1),
200                tip - MINIMUM_UNWIND_SAFE_DISTANCE - 2,
201                should_prune,
202            ),
203            (PruneMode::Before(tip + 1), 1, should_prune),
204            (PruneMode::Before(tip + 1), tip + 1, !should_prune),
205        ];
206
207        for (index, (mode, block, expected_result)) in tests.into_iter().enumerate() {
208            assert_eq!(mode.should_prune(block, tip), expected_result, "Test {} failed", index + 1,);
209        }
210    }
211
212    #[test]
213    fn prune_mode_deserialize() {
214        #[derive(Debug, Deserialize)]
215        struct Config {
216            a: Option<PruneMode>,
217            b: Option<PruneMode>,
218            c: Option<PruneMode>,
219            d: Option<PruneMode>,
220        }
221
222        let toml_str = r#"
223        a = "full"
224        b = { distance = 10 }
225        c = { before = 20 }
226    "#;
227
228        assert_matches!(
229            toml::from_str(toml_str),
230            Ok(Config {
231                a: Some(PruneMode::Full),
232                b: Some(PruneMode::Distance(10)),
233                c: Some(PruneMode::Before(20)),
234                d: None
235            })
236        );
237    }
238}