synor/crates/synor-dag/src/pruning.rs

//! DAG pruning for memory management.
//!
//! As the DAG grows, we need to prune old blocks to maintain reasonable
//! memory usage. Pruning removes blocks that are deep enough to be
//! considered finalized while maintaining the ability to verify new blocks.
//!
//! Key concepts:
//! - **Pruning point**: The oldest block we keep full data for
//! - **Virtual genesis**: Acts as the new genesis for pruned DAG
//! - **Finality depth**: How deep a block must be to be considered final

use crate::{
    dag::{BlockDag, DagError},
    ghostdag::GhostdagManager,
    BlockId, BlueScore, FINALITY_DEPTH, PRUNING_DEPTH,
};
use hashbrown::HashSet;
use parking_lot::RwLock;
use thiserror::Error;

/// Configuration for pruning behavior.
#[derive(Clone, Debug)]
pub struct PruningConfig {
    /// Depth at which blocks are considered final.
    pub finality_depth: u64,
    /// Depth at which blocks can be pruned.
    pub pruning_depth: u64,
    /// Minimum blocks to keep in DAG.
    pub min_blocks: usize,
    /// Maximum blocks before triggering pruning.
    pub max_blocks: usize,
}

impl Default for PruningConfig {
    fn default() -> Self {
        PruningConfig {
            finality_depth: FINALITY_DEPTH,
            pruning_depth: PRUNING_DEPTH,
            min_blocks: 10_000,
            max_blocks: 1_000_000,
        }
    }
}

impl PruningConfig {
    /// Creates a config suitable for testing.
    pub fn for_testing() -> Self {
        PruningConfig {
            finality_depth: 100,
            pruning_depth: 200,
            min_blocks: 10,
            max_blocks: 1000,
        }
    }
}

/// Manages pruning of the DAG.
pub struct PruningManager {
    /// Configuration.
    config: PruningConfig,
    /// Current pruning point.
    pruning_point: RwLock<Option<BlockId>>,
    /// Blocks that have been pruned (headers kept, body removed).
    pruned_blocks: RwLock<HashSet<BlockId>>,
}

impl PruningManager {
    /// Creates a new pruning manager.
    pub fn new(config: PruningConfig) -> Self {
        PruningManager {
            config,
            pruning_point: RwLock::new(None),
            pruned_blocks: RwLock::new(HashSet::new()),
        }
    }

    /// Creates a pruning manager with default config.
    pub fn with_defaults() -> Self {
        Self::new(PruningConfig::default())
    }

    /// Returns the current pruning point.
    pub fn pruning_point(&self) -> Option<BlockId> {
        *self.pruning_point.read()
    }

    /// Checks if a block has been pruned.
    pub fn is_pruned(&self, block_id: &BlockId) -> bool {
        self.pruned_blocks.read().contains(block_id)
    }

    /// Calculates the new pruning point based on the current tips.
    ///
    /// The pruning point is the block at `pruning_depth` on the selected chain.
    pub fn calculate_pruning_point(
        &self,
        ghostdag: &GhostdagManager,
        tips: &[BlockId],
    ) -> Result<Option<BlockId>, PruningError> {
        if tips.is_empty() {
            return Ok(None);
        }

        // Find the tip with highest blue score
        let mut best_tip = tips[0];
        let mut best_score = ghostdag.get_blue_score(&tips[0]).unwrap_or(0);

        for tip in tips.iter().skip(1) {
            let score = ghostdag.get_blue_score(tip).unwrap_or(0);
            if score > best_score {
                best_score = score;
                best_tip = *tip;
            }
        }

        // Get the selected chain from best tip
        let selected_chain = ghostdag.get_selected_chain(&best_tip)?;

        // Find block at pruning depth
        if selected_chain.len() > self.config.pruning_depth as usize {
            let pruning_index = self.config.pruning_depth as usize;
            Ok(Some(selected_chain[pruning_index]))
        } else {
            Ok(None) // Not deep enough to prune
        }
    }

    /// Calculates which blocks can be pruned.
    ///
    /// Returns blocks that are:
    /// - Below the pruning point on the selected chain
    /// - Not needed for verification of recent blocks
    pub fn calculate_blocks_to_prune(
        &self,
        dag: &BlockDag,
        ghostdag: &GhostdagManager,
        new_pruning_point: &BlockId,
    ) -> Result<Vec<BlockId>, PruningError> {
        let _current_pruning_point = self.pruning_point();
        let mut blocks_to_prune = Vec::new();

        // Get all blocks below the new pruning point
        let ancestors = dag.get_ancestors(new_pruning_point, usize::MAX);

        for ancestor in ancestors {
            // Skip if already pruned
            if self.is_pruned(&ancestor) {
                continue;
            }

            // Skip genesis
            if ancestor == dag.genesis() {
                continue;
            }

            // Check if block is below finality depth
            let ancestor_score = ghostdag.get_blue_score(&ancestor).unwrap_or(0);
            let pp_score = ghostdag.get_blue_score(new_pruning_point).unwrap_or(0);

            if pp_score - ancestor_score > self.config.finality_depth {
                blocks_to_prune.push(ancestor);
            }
        }

        Ok(blocks_to_prune)
    }

    /// Updates the pruning point and marks blocks as pruned.
    pub fn update_pruning_point(
        &self,
        dag: &BlockDag,
        ghostdag: &GhostdagManager,
        tips: &[BlockId],
    ) -> Result<PruningResult, PruningError> {
        let new_pp = self.calculate_pruning_point(ghostdag, tips)?;

        let Some(new_pruning_point) = new_pp else {
            return Ok(PruningResult::no_change());
        };

        let current_pp = self.pruning_point();

        // Check if pruning point actually changed
        if current_pp == Some(new_pruning_point) {
            return Ok(PruningResult::no_change());
        }

        // Calculate blocks to prune
        let blocks_to_prune = self.calculate_blocks_to_prune(dag, ghostdag, &new_pruning_point)?;

        // Mark blocks as pruned
        let mut pruned = self.pruned_blocks.write();
        for block in &blocks_to_prune {
            pruned.insert(*block);
        }

        // Update pruning point
        *self.pruning_point.write() = Some(new_pruning_point);

        Ok(PruningResult {
            old_pruning_point: current_pp,
            new_pruning_point: Some(new_pruning_point),
            blocks_pruned: blocks_to_prune.len(),
            blocks_to_prune,
        })
    }

    /// Checks if pruning should be triggered.
    pub fn should_prune(&self, dag: &BlockDag) -> bool {
        dag.len() > self.config.max_blocks
    }

    /// Returns the number of pruned blocks.
    pub fn pruned_count(&self) -> usize {
        self.pruned_blocks.read().len()
    }

    /// Returns the configuration.
    pub fn config(&self) -> &PruningConfig {
        &self.config
    }

    /// Validates that a block is not in the pruned set.
    pub fn validate_not_pruned(&self, block_id: &BlockId) -> Result<(), PruningError> {
        if self.is_pruned(block_id) {
            Err(PruningError::BlockPruned(*block_id))
        } else {
            Ok(())
        }
    }
}

impl std::fmt::Debug for PruningManager {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PruningManager")
            .field("pruning_point", &self.pruning_point())
            .field("pruned_count", &self.pruned_count())
            .field("config", &self.config)
            .finish()
    }
}

/// Result of a pruning operation.
#[derive(Clone, Debug)]
pub struct PruningResult {
    /// Previous pruning point.
    pub old_pruning_point: Option<BlockId>,
    /// New pruning point.
    pub new_pruning_point: Option<BlockId>,
    /// Number of blocks pruned.
    pub blocks_pruned: usize,
    /// List of pruned block IDs.
    pub blocks_to_prune: Vec<BlockId>,
}

impl PruningResult {
    /// Creates a result indicating no change.
    pub fn no_change() -> Self {
        PruningResult {
            old_pruning_point: None,
            new_pruning_point: None,
            blocks_pruned: 0,
            blocks_to_prune: Vec::new(),
        }
    }

    /// Returns true if pruning occurred.
    pub fn did_prune(&self) -> bool {
        self.blocks_pruned > 0
    }

    /// Returns true if the pruning point changed.
    pub fn pruning_point_changed(&self) -> bool {
        self.old_pruning_point != self.new_pruning_point
    }
}

/// Errors related to pruning.
#[derive(Debug, Error)]
pub enum PruningError {
    #[error("Block has been pruned: {0}")]
    BlockPruned(BlockId),

    #[error("Invalid pruning point: {0}")]
    InvalidPruningPoint(BlockId),

    #[error("GHOSTDAG error: {0}")]
    GhostdagError(#[from] crate::ghostdag::GhostdagError),

    #[error("DAG error: {0}")]
    DagError(#[from] DagError),
}

/// Proof that a block existed before pruning.
#[derive(Clone, Debug)]
pub struct PruningProof {
    /// The pruning point this proof is relative to.
    pub pruning_point: BlockId,
    /// Headers of blocks in the proof path.
    pub headers: Vec<BlockId>,
    /// Blue scores for verification.
    pub blue_scores: Vec<BlueScore>,
}

impl PruningProof {
    /// Creates a new pruning proof.
    pub fn new(pruning_point: BlockId) -> Self {
        PruningProof {
            pruning_point,
            headers: Vec::new(),
            blue_scores: Vec::new(),
        }
    }

    /// Adds a block to the proof.
    pub fn add_block(&mut self, block_id: BlockId, blue_score: BlueScore) {
        self.headers.push(block_id);
        self.blue_scores.push(blue_score);
    }

    /// Returns the number of blocks in the proof.
    pub fn len(&self) -> usize {
        self.headers.len()
    }

    /// Returns true if the proof is empty.
    pub fn is_empty(&self) -> bool {
        self.headers.is_empty()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{dag::BlockDag, reachability::ReachabilityStore};
    use synor_types::Hash256;

    fn make_block_id(n: u8) -> BlockId {
        let mut bytes = [0u8; 32];
        bytes[0] = n;
        Hash256::from_bytes(bytes)
    }

    fn setup_test() -> (
        std::sync::Arc<BlockDag>,
        std::sync::Arc<ReachabilityStore>,
        std::sync::Arc<GhostdagManager>,
        PruningManager,
    ) {
        let genesis = make_block_id(0);
        let dag = std::sync::Arc::new(BlockDag::new(genesis, 0));
        let reachability = std::sync::Arc::new(ReachabilityStore::new(genesis));
        let ghostdag = std::sync::Arc::new(GhostdagManager::with_k(
            dag.clone(),
            reachability.clone(),
            3,
        ));
        let pruning = PruningManager::new(PruningConfig::for_testing());
        (dag, reachability, ghostdag, pruning)
    }

    #[test]
    fn test_initial_state() {
        let (_, _, _, pruning) = setup_test();

        assert!(pruning.pruning_point().is_none());
        assert_eq!(pruning.pruned_count(), 0);
    }

    #[test]
    fn test_is_pruned() {
        let (_, _, _, pruning) = setup_test();

        let block = make_block_id(1);
        assert!(!pruning.is_pruned(&block));

        // Manually add to pruned set for testing
        pruning.pruned_blocks.write().insert(block);
        assert!(pruning.is_pruned(&block));
    }

    #[test]
    fn test_should_prune() {
        let genesis = make_block_id(0);
        let dag = BlockDag::new(genesis, 0);

        let config = PruningConfig {
            max_blocks: 5,
            ..PruningConfig::for_testing()
        };
        let pruning = PruningManager::new(config);

        // Initially shouldn't prune
        assert!(!pruning.should_prune(&dag));
    }

    #[test]
    fn test_pruning_result() {
        let result = PruningResult::no_change();
        assert!(!result.did_prune());
        assert!(!result.pruning_point_changed());

        let result = PruningResult {
            old_pruning_point: Some(make_block_id(0)),
            new_pruning_point: Some(make_block_id(1)),
            blocks_pruned: 5,
            blocks_to_prune: vec![make_block_id(2)],
        };
        assert!(result.did_prune());
        assert!(result.pruning_point_changed());
    }

    #[test]
    fn test_validate_not_pruned() {
        let (_, _, _, pruning) = setup_test();

        let block = make_block_id(1);
        assert!(pruning.validate_not_pruned(&block).is_ok());

        pruning.pruned_blocks.write().insert(block);
        assert!(pruning.validate_not_pruned(&block).is_err());
    }

    #[test]
    fn test_pruning_proof() {
        let mut proof = PruningProof::new(make_block_id(0));
        assert!(proof.is_empty());

        proof.add_block(make_block_id(1), 100);
        proof.add_block(make_block_id(2), 200);

        assert_eq!(proof.len(), 2);
        assert!(!proof.is_empty());
    }
}