1use crate::{segment::PrunePurpose, PruneSegment, PruneSegmentError};
2use alloy_primitives::BlockNumber;
3
4#[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 Full,
14 Distance(u64),
16 Before(BlockNumber),
18}
19
20#[cfg(any(test, feature = "test-utils"))]
21#[allow(clippy::derivable_impls)]
22impl Default for PruneMode {
23 fn default() -> Self {
24 Self::Full
25 }
26}
27
28impl PruneMode {
29 pub const fn before_inclusive(block_number: BlockNumber) -> Self {
33 Self::Before(block_number + 1)
34 }
35
36 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 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 Self::Full if min_blocks <= tip => Some((tip - min_blocks, *self)),
61 Self::Full => None, Self::Distance(distance) if *distance > tip => None, 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 > tip => None, Self::Before(n) => (tip - n >= min_blocks).then(|| ((*n).saturating_sub(1), *self)),
67 _ => return Err(PruneSegmentError::Configuration(segment)),
68 };
69 Ok(result)
70 }
71
72 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 pub const fn is_full(&self) -> bool {
88 matches!(self, Self::Full)
89 }
90
91 pub const fn is_distance(&self) -> bool {
93 matches!(self, Self::Distance(_))
94 }
95
96 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 (PruneMode::Full, Ok(Some(tip - segment.min_blocks()))),
143 (PruneMode::Distance(tip + 1), Ok(None)),
145 (
146 PruneMode::Distance(segment.min_blocks() + 1),
147 Ok(Some(tip - (segment.min_blocks() + 1))),
148 ),
149 (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 (PruneMode::Before(tip - 1), Ok(None)),
161 ];
162
163 for (index, (mode, expected_result)) in tests.into_iter().enumerate() {
164 assert_eq!(
165 mode.prune_target_block(tip, segment, PrunePurpose::User),
166 expected_result.map(|r| r.map(|b| (b, mode))),
167 "Test {} failed",
168 index + 1,
169 );
170 }
171
172 assert_eq!(
174 PruneMode::Full.prune_target_block(
175 tip,
176 PruneSegment::TransactionLookup,
177 PrunePurpose::User
178 ),
179 Ok(Some((tip, PruneMode::Full))),
180 );
181 }
182
183 #[test]
184 fn test_should_prune() {
185 let tip = 20000;
186 let should_prune = true;
187
188 let tests = vec![
189 (PruneMode::Distance(tip + 1), 1, !should_prune),
190 (
191 PruneMode::Distance(MINIMUM_UNWIND_SAFE_DISTANCE + 1),
192 tip - MINIMUM_UNWIND_SAFE_DISTANCE - 1,
193 !should_prune,
194 ),
195 (
196 PruneMode::Distance(MINIMUM_UNWIND_SAFE_DISTANCE + 1),
197 tip - MINIMUM_UNWIND_SAFE_DISTANCE - 2,
198 should_prune,
199 ),
200 (PruneMode::Before(tip + 1), 1, should_prune),
201 (PruneMode::Before(tip + 1), tip + 1, !should_prune),
202 ];
203
204 for (index, (mode, block, expected_result)) in tests.into_iter().enumerate() {
205 assert_eq!(mode.should_prune(block, tip), expected_result, "Test {} failed", index + 1,);
206 }
207 }
208
209 #[test]
210 fn prune_mode_deserialize() {
211 #[derive(Debug, Deserialize)]
212 struct Config {
213 a: Option<PruneMode>,
214 b: Option<PruneMode>,
215 c: Option<PruneMode>,
216 d: Option<PruneMode>,
217 }
218
219 let toml_str = r#"
220 a = "full"
221 b = { distance = 10 }
222 c = { before = 20 }
223 "#;
224
225 assert_matches!(
226 toml::from_str(toml_str),
227 Ok(Config {
228 a: Some(PruneMode::Full),
229 b: Some(PruneMode::Distance(10)),
230 c: Some(PruneMode::Before(20)),
231 d: None
232 })
233 );
234 }
235}