use crate::{segment::PrunePurpose, PruneSegment, PruneSegmentError};
use alloy_primitives::BlockNumber;
use reth_codecs::{add_arbitrary_tests, Compact};
use serde::{Deserialize, Serialize};
#[derive(Debug, PartialEq, Eq, Clone, Copy, Serialize, Deserialize, Compact)]
#[serde(rename_all = "lowercase")]
#[cfg_attr(any(test, feature = "test-utils"), derive(arbitrary::Arbitrary))]
#[add_arbitrary_tests(compact)]
pub enum PruneMode {
Full,
Distance(u64),
Before(BlockNumber),
}
#[cfg(any(test, feature = "test-utils"))]
impl Default for PruneMode {
fn default() -> Self {
Self::Full
}
}
impl PruneMode {
pub const fn before_inclusive(block_number: BlockNumber) -> Self {
Self::Before(block_number + 1)
}
pub fn prune_target_block(
&self,
tip: BlockNumber,
segment: PruneSegment,
purpose: PrunePurpose,
) -> Result<Option<(BlockNumber, Self)>, PruneSegmentError> {
let result = match self {
Self::Full if segment.min_blocks(purpose) == 0 => Some((tip, *self)),
Self::Distance(distance) if *distance > tip => None, Self::Distance(distance) if *distance >= segment.min_blocks(purpose) => {
Some((tip - distance, *self))
}
Self::Before(n) if *n == tip + 1 && purpose.is_static_file() => Some((tip, *self)),
Self::Before(n) if *n > tip => None, Self::Before(n) if tip - n >= segment.min_blocks(purpose) => {
Some(((*n).saturating_sub(1), *self))
}
_ => return Err(PruneSegmentError::Configuration(segment)),
};
Ok(result)
}
pub const fn should_prune(&self, block: BlockNumber, tip: BlockNumber) -> bool {
match self {
Self::Full => true,
Self::Distance(distance) => {
if *distance > tip {
return false
}
block < tip - *distance
}
Self::Before(n) => *n > block,
}
}
pub const fn is_full(&self) -> bool {
matches!(self, Self::Full)
}
pub const fn is_distance(&self) -> bool {
matches!(self, Self::Distance(_))
}
}
#[cfg(test)]
mod tests {
use crate::{
PruneMode, PrunePurpose, PruneSegment, PruneSegmentError, MINIMUM_PRUNING_DISTANCE,
};
use assert_matches::assert_matches;
use serde::Deserialize;
#[test]
fn test_prune_target_block() {
let tip = 20000;
let segment = PruneSegment::Receipts;
let tests = vec![
(PruneMode::Full, Err(PruneSegmentError::Configuration(segment))),
(PruneMode::Distance(tip + 1), Ok(None)),
(
PruneMode::Distance(segment.min_blocks(PrunePurpose::User) + 1),
Ok(Some(tip - (segment.min_blocks(PrunePurpose::User) + 1))),
),
(PruneMode::Before(tip + 1), Ok(None)),
(
PruneMode::Before(tip - MINIMUM_PRUNING_DISTANCE),
Ok(Some(tip - MINIMUM_PRUNING_DISTANCE - 1)),
),
(
PruneMode::Before(tip - MINIMUM_PRUNING_DISTANCE - 1),
Ok(Some(tip - MINIMUM_PRUNING_DISTANCE - 2)),
),
(PruneMode::Before(tip - 1), Err(PruneSegmentError::Configuration(segment))),
];
for (index, (mode, expected_result)) in tests.into_iter().enumerate() {
assert_eq!(
mode.prune_target_block(tip, segment, PrunePurpose::User),
expected_result.map(|r| r.map(|b| (b, mode))),
"Test {} failed",
index + 1,
);
}
assert_eq!(
PruneMode::Full.prune_target_block(tip, PruneSegment::Transactions, PrunePurpose::User),
Ok(Some((tip, PruneMode::Full))),
);
}
#[test]
fn test_should_prune() {
let tip = 20000;
let should_prune = true;
let tests = vec![
(PruneMode::Distance(tip + 1), 1, !should_prune),
(
PruneMode::Distance(MINIMUM_PRUNING_DISTANCE + 1),
tip - MINIMUM_PRUNING_DISTANCE - 1,
!should_prune,
),
(
PruneMode::Distance(MINIMUM_PRUNING_DISTANCE + 1),
tip - MINIMUM_PRUNING_DISTANCE - 2,
should_prune,
),
(PruneMode::Before(tip + 1), 1, should_prune),
(PruneMode::Before(tip + 1), tip + 1, !should_prune),
];
for (index, (mode, block, expected_result)) in tests.into_iter().enumerate() {
assert_eq!(mode.should_prune(block, tip), expected_result, "Test {} failed", index + 1,);
}
}
#[test]
fn prune_mode_deserialize() {
#[derive(Debug, Deserialize)]
struct Config {
a: Option<PruneMode>,
b: Option<PruneMode>,
c: Option<PruneMode>,
d: Option<PruneMode>,
}
let toml_str = r#"
a = "full"
b = { distance = 10 }
c = { before = 20 }
"#;
assert_matches!(
toml::from_str(toml_str),
Ok(Config {
a: Some(PruneMode::Full),
b: Some(PruneMode::Distance(10)),
c: Some(PruneMode::Before(20)),
d: None
})
);
}
}