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

//! GHOSTDAG blue set selection algorithm.
//!
//! GHOSTDAG selects a "blue set" of blocks that form the main chain, while
//! still including transactions from parallel (red) blocks. The algorithm
//! ensures that the blue set forms a k-cluster: any chain of blocks in the
//! blue set has at most k blocks in its anticone.
//!
//! Key concepts:
//! - **k parameter**: Controls the allowed parallelism (antichain width)
//! - **Blue set**: Blocks considered "honest" and part of main chain
//! - **Red set**: Parallel blocks that don't fit k-cluster property
//! - **Selected parent**: The blue parent with highest blue score
//! - **Blue score**: Cumulative count of blue blocks in ancestry

use crate::{
    dag::{BlockDag, DagError},
    reachability::ReachabilityStore,
    BlockId, BlueScore, GHOSTDAG_K,
};
use hashbrown::{HashMap, HashSet};
use parking_lot::RwLock;
use thiserror::Error;

/// GHOSTDAG data computed for each block.
#[derive(Clone, Debug)]
pub struct GhostdagData {
    /// Blue score (cumulative blue ancestor count).
    pub blue_score: BlueScore,
    /// The selected parent (highest blue score parent).
    pub selected_parent: BlockId,
    /// Blocks in the merge set that are blue.
    pub merge_set_blues: Vec<BlockId>,
    /// Blocks in the merge set that are red.
    pub merge_set_reds: Vec<BlockId>,
    /// Blues anticone sizes (for k-cluster verification).
    pub blues_anticone_sizes: HashMap<BlockId, usize>,
}

impl GhostdagData {
    /// Creates data for the genesis block.
    pub fn genesis(genesis_id: BlockId) -> Self {
        GhostdagData {
            blue_score: 0,
            selected_parent: genesis_id, // Genesis is its own selected parent
            merge_set_blues: Vec::new(),
            merge_set_reds: Vec::new(),
            blues_anticone_sizes: HashMap::new(),
        }
    }

    /// Returns the total merge set size.
    pub fn merge_set_size(&self) -> usize {
        self.merge_set_blues.len() + self.merge_set_reds.len()
    }
}

/// Manager for GHOSTDAG calculations.
pub struct GhostdagManager {
    /// The DAG structure.
    dag: std::sync::Arc<BlockDag>,
    /// Reachability queries.
    reachability: std::sync::Arc<ReachabilityStore>,
    /// Cached GHOSTDAG data.
    data_cache: RwLock<HashMap<BlockId, GhostdagData>>,
    /// The k parameter.
    k: u8,
}

impl GhostdagManager {
    /// Creates a new GHOSTDAG manager.
    pub fn new(
        dag: std::sync::Arc<BlockDag>,
        reachability: std::sync::Arc<ReachabilityStore>,
    ) -> Self {
        Self::with_k(dag, reachability, GHOSTDAG_K)
    }

    /// Creates a new GHOSTDAG manager with custom k.
    pub fn with_k(
        dag: std::sync::Arc<BlockDag>,
        reachability: std::sync::Arc<ReachabilityStore>,
        k: u8,
    ) -> Self {
        let mut data_cache = HashMap::new();

        // Initialize with genesis data
        let genesis = dag.genesis();
        data_cache.insert(genesis, GhostdagData::genesis(genesis));

        GhostdagManager {
            dag,
            reachability,
            data_cache: RwLock::new(data_cache),
            k,
        }
    }

    /// Processes a new block and computes its GHOSTDAG data.
    pub fn add_block(
        &self,
        block_id: BlockId,
        parents: &[BlockId],
    ) -> Result<GhostdagData, GhostdagError> {
        if parents.is_empty() {
            return Err(GhostdagError::NoParents);
        }

        // 1. Find the selected parent (parent with highest blue score)
        let selected_parent = self.find_selected_parent(parents)?;

        // 2. Get the selected parent's GHOSTDAG data
        let selected_parent_data = self.get_data(&selected_parent)?;

        // 3. Calculate merge set (blocks not in selected parent's past)
        let merge_set = self.calculate_merge_set(block_id, parents, &selected_parent)?;

        // 4. Partition merge set into blues and reds using k-cluster
        let (merge_set_blues, merge_set_reds, blues_anticone_sizes) =
            self.partition_merge_set(&selected_parent_data, &merge_set)?;

        // 5. Calculate blue score
        let blue_score = selected_parent_data.blue_score + merge_set_blues.len() as u64 + 1;

        let data = GhostdagData {
            blue_score,
            selected_parent,
            merge_set_blues,
            merge_set_reds,
            blues_anticone_sizes,
        };

        // Cache the result
        self.data_cache.write().insert(block_id, data.clone());

        // Update the DAG with GHOSTDAG data
        self.dag.update_ghostdag_data(
            &block_id,
            data.blue_score,
            Some(data.selected_parent),
            true, // New blocks start as blue (may be re-evaluated)
            data.merge_set_blues.clone(),
            data.merge_set_reds.clone(),
        )?;

        Ok(data)
    }

    /// Finds the parent with the highest blue score.
    fn find_selected_parent(&self, parents: &[BlockId]) -> Result<BlockId, GhostdagError> {
        let mut best_parent = parents[0];
        let mut best_score = self.get_blue_score(&parents[0])?;

        for parent in parents.iter().skip(1) {
            let score = self.get_blue_score(parent)?;
            if score > best_score {
                best_score = score;
                best_parent = *parent;
            }
        }

        Ok(best_parent)
    }

    /// Calculates the merge set for a new block.
    ///
    /// The merge set contains all blocks that are:
    /// - In the past of the new block's parents
    /// - Not in the past of the selected parent
    fn calculate_merge_set(
        &self,
        _block_id: BlockId,
        parents: &[BlockId],
        selected_parent: &BlockId,
    ) -> Result<Vec<BlockId>, GhostdagError> {
        let mut merge_set = Vec::new();

        // For each non-selected parent, add blocks not in selected parent's past
        for parent in parents {
            if parent == selected_parent {
                continue;
            }

            // Check if this parent is in the selected parent's past
            let is_in_past = self
                .reachability
                .is_ancestor(selected_parent, parent)
                .map_err(|e| GhostdagError::ReachabilityError(e.to_string()))?;

            if !is_in_past {
                merge_set.push(*parent);

                // Also add ancestors of this parent that aren't in selected parent's past
                // (up to merge depth limit)
                let parent_ancestors = self.get_ancestors_not_in_past(parent, selected_parent)?;
                merge_set.extend(parent_ancestors);
            }
        }

        // Remove duplicates
        merge_set.sort();
        merge_set.dedup();

        Ok(merge_set)
    }

    /// Gets ancestors of a block that are not in another block's past.
    fn get_ancestors_not_in_past(
        &self,
        block: &BlockId,
        reference: &BlockId,
    ) -> Result<Vec<BlockId>, GhostdagError> {
        let mut result = Vec::new();
        let mut visited = HashSet::new();
        let mut frontier = vec![*block];

        // Limit traversal depth to avoid expensive computations
        const MAX_DEPTH: usize = 100;
        let mut depth = 0;

        while !frontier.is_empty() && depth < MAX_DEPTH {
            let mut next_frontier = Vec::new();

            for current in frontier {
                if visited.contains(&current) {
                    continue;
                }
                visited.insert(current);

                // Check if current is in reference's past
                let is_in_past = self
                    .reachability
                    .is_ancestor(reference, &current)
                    .unwrap_or(true); // Assume in past if error

                if !is_in_past && current != *block {
                    result.push(current);

                    // Add parents to frontier
                    if let Some(parents) = self.dag.get_parents(&current) {
                        next_frontier.extend(parents.iter().copied());
                    }
                }
            }

            frontier = next_frontier;
            depth += 1;
        }

        Ok(result)
    }

    /// Partitions the merge set into blues and reds using k-cluster.
    ///
    /// A block is blue if adding it to the blue set maintains the k-cluster property:
    /// - For any blue block B in the merge set, |anticone(B) ∩ blues| ≤ k
    ///
    /// OPTIMIZATION: Uses incremental anticone tracking to avoid O(n²) complexity.
    /// Instead of recalculating anticone sizes, we track them incrementally as we
    /// add new blues and update existing sizes in O(n) per candidate.
    fn partition_merge_set(
        &self,
        selected_parent_data: &GhostdagData,
        merge_set: &[BlockId],
    ) -> Result<(Vec<BlockId>, Vec<BlockId>, HashMap<BlockId, usize>), GhostdagError> {
        let mut blues: Vec<BlockId> = Vec::with_capacity(merge_set.len());
        let mut reds: Vec<BlockId> = Vec::new();
        let mut blues_anticone_sizes: HashMap<BlockId, usize> = HashMap::with_capacity(
            merge_set.len() + selected_parent_data.blues_anticone_sizes.len(),
        );

        // Initialize with selected parent's blues anticone sizes
        blues_anticone_sizes.extend(selected_parent_data.blues_anticone_sizes.iter());

        // Sort merge set by blue score (process higher-scoring blocks first)
        // Pre-fetch scores to avoid repeated lookups during sort
        let mut sorted_merge_set: Vec<_> = merge_set
            .iter()
            .map(|id| (*id, self.get_blue_score(id).unwrap_or(0)))
            .collect();
        sorted_merge_set.sort_by(|a, b| b.1.cmp(&a.1)); // Descending by score

        for (candidate, _score) in sorted_merge_set {
            // OPTIMIZATION: Calculate anticone relationships in a single pass
            // and track which existing blues are in the candidate's anticone
            let (anticone_size, anticone_blues) =
                self.calculate_blues_anticone_optimized(&candidate, &blues)?;

            // Check if adding this block maintains k-cluster property
            if anticone_size <= self.k as usize {
                // OPTIMIZATION: Check k-cluster incrementally
                // Only check blues that are in the candidate's anticone
                let can_add =
                    self.check_k_cluster_incremental(&anticone_blues, &blues_anticone_sizes)?;

                if can_add {
                    // Update anticone sizes for existing blues that are in candidate's anticone
                    for blue in &anticone_blues {
                        if let Some(size) = blues_anticone_sizes.get_mut(blue) {
                            *size += 1;
                        }
                    }
                    blues_anticone_sizes.insert(candidate, anticone_size);
                    blues.push(candidate);
                } else {
                    reds.push(candidate);
                }
            } else {
                reds.push(candidate);
            }
        }

        Ok((blues, reds, blues_anticone_sizes))
    }

    /// Calculates anticone size and returns which blues are in the anticone.
    ///
    /// OPTIMIZATION: Returns both the count AND the list of anticone blues,
    /// avoiding redundant reachability queries later.
    fn calculate_blues_anticone_optimized(
        &self,
        candidate: &BlockId,
        current_blues: &[BlockId],
    ) -> Result<(usize, Vec<BlockId>), GhostdagError> {
        let mut anticone_blues = Vec::with_capacity(self.k as usize + 1);

        for blue in current_blues {
            let is_anticone = self
                .reachability
                .is_anticone(candidate, blue)
                .map_err(|e| GhostdagError::ReachabilityError(e.to_string()))?;

            if is_anticone {
                anticone_blues.push(*blue);
                // Early exit if we already exceed k (no point continuing)
                if anticone_blues.len() > self.k as usize {
                    return Ok((anticone_blues.len(), anticone_blues));
                }
            }
        }

        Ok((anticone_blues.len(), anticone_blues))
    }

    /// Checks if adding a new blue would violate k-cluster for existing blues.
    ///
    /// OPTIMIZATION: Only checks blues that are in the new block's anticone,
    /// and uses the pre-computed anticone sizes instead of recalculating.
    fn check_k_cluster_incremental(
        &self,
        anticone_blues: &[BlockId],
        blues_anticone_sizes: &HashMap<BlockId, usize>,
    ) -> Result<bool, GhostdagError> {
        // For each blue in the candidate's anticone, check if adding 1 to its
        // anticone size would exceed k
        for blue in anticone_blues {
            let current_size = blues_anticone_sizes.get(blue).copied().unwrap_or(0);
            if current_size >= self.k as usize {
                return Ok(false);
            }
        }
        Ok(true)
    }

    /// Legacy method for compatibility - calculates blues anticone size.
    #[allow(dead_code)]
    fn calculate_blues_anticone_size(
        &self,
        candidate: &BlockId,
        current_blues: &[BlockId],
    ) -> Result<usize, GhostdagError> {
        let (size, _) = self.calculate_blues_anticone_optimized(candidate, current_blues)?;
        Ok(size)
    }

    /// Legacy method for compatibility - checks k-cluster constraint.
    #[allow(dead_code)]
    fn check_k_cluster_with_new_blue(
        &self,
        new_blue: &BlockId,
        current_blues: &[BlockId],
    ) -> Result<bool, GhostdagError> {
        let (_, anticone_blues) =
            self.calculate_blues_anticone_optimized(new_blue, current_blues)?;
        // Build a temporary anticone sizes map for the check
        let mut temp_sizes = HashMap::new();
        for blue in current_blues {
            let (size, _) = self.calculate_blues_anticone_optimized(blue, current_blues)?;
            temp_sizes.insert(*blue, size);
        }
        self.check_k_cluster_incremental(&anticone_blues, &temp_sizes)
    }

    /// Gets the GHOSTDAG data for a block.
    pub fn get_data(&self, block_id: &BlockId) -> Result<GhostdagData, GhostdagError> {
        self.data_cache
            .read()
            .get(block_id)
            .cloned()
            .ok_or(GhostdagError::DataNotFound(*block_id))
    }

    /// Gets the blue score for a block.
    pub fn get_blue_score(&self, block_id: &BlockId) -> Result<BlueScore, GhostdagError> {
        self.data_cache
            .read()
            .get(block_id)
            .map(|d| d.blue_score)
            .ok_or(GhostdagError::DataNotFound(*block_id))
    }

    /// Gets the selected parent of a block.
    pub fn get_selected_parent(&self, block_id: &BlockId) -> Result<BlockId, GhostdagError> {
        self.data_cache
            .read()
            .get(block_id)
            .map(|d| d.selected_parent)
            .ok_or(GhostdagError::DataNotFound(*block_id))
    }

    /// Returns the chain of selected parents from a block to genesis.
    pub fn get_selected_chain(&self, from: &BlockId) -> Result<Vec<BlockId>, GhostdagError> {
        let mut chain = vec![*from];
        let mut current = *from;
        let genesis = self.dag.genesis();

        while current != genesis {
            let parent = self.get_selected_parent(&current)?;
            chain.push(parent);
            current = parent;
        }

        Ok(chain)
    }

    /// Checks if a block is in the blue set.
    pub fn is_blue(&self, block_id: &BlockId) -> bool {
        self.dag
            .get_block(block_id)
            .map(|meta| meta.is_blue)
            .unwrap_or(false)
    }

    /// Returns the k parameter.
    pub fn k(&self) -> u8 {
        self.k
    }
}

impl std::fmt::Debug for GhostdagManager {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("GhostdagManager")
            .field("k", &self.k)
            .field("cached_blocks", &self.data_cache.read().len())
            .finish()
    }
}

/// Errors that can occur during GHOSTDAG calculations.
#[derive(Debug, Error)]
pub enum GhostdagError {
    #[error("No parents specified")]
    NoParents,

    #[error("GHOSTDAG data not found for block: {0}")]
    DataNotFound(BlockId),

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

    #[error("Reachability error: {0}")]
    ReachabilityError(String),

    #[error("Invalid block structure: {0}")]
    InvalidStructure(String),
}

#[cfg(test)]
mod tests {
    use super::*;
    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_dag() -> (
        std::sync::Arc<BlockDag>,
        std::sync::Arc<ReachabilityStore>,
        GhostdagManager,
    ) {
        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 = GhostdagManager::with_k(dag.clone(), reachability.clone(), 3);
        (dag, reachability, ghostdag)
    }

    #[test]
    fn test_genesis_data() {
        let genesis = make_block_id(0);
        let (dag, reachability, ghostdag) = setup_test_dag();

        let data = ghostdag.get_data(&genesis).unwrap();
        assert_eq!(data.blue_score, 0);
        assert_eq!(data.selected_parent, genesis);
    }

    #[test]
    fn test_simple_chain() {
        let genesis = make_block_id(0);
        let (dag, reachability, ghostdag) = setup_test_dag();

        // Add a simple chain: genesis -> block1 -> block2
        let block1 = make_block_id(1);
        dag.insert_block(block1, vec![genesis], 100).unwrap();
        reachability.add_block(block1, genesis, &[genesis]).unwrap();
        let data1 = ghostdag.add_block(block1, &[genesis]).unwrap();

        assert_eq!(data1.blue_score, 1);
        assert_eq!(data1.selected_parent, genesis);

        let block2 = make_block_id(2);
        dag.insert_block(block2, vec![block1], 200).unwrap();
        reachability.add_block(block2, block1, &[block1]).unwrap();
        let data2 = ghostdag.add_block(block2, &[block1]).unwrap();

        assert_eq!(data2.blue_score, 2);
        assert_eq!(data2.selected_parent, block1);
    }

    #[test]
    fn test_selected_parent_highest_score() {
        let genesis = make_block_id(0);
        let (dag, reachability, ghostdag) = setup_test_dag();

        // Create two branches from genesis
        let block1 = make_block_id(1);
        let block2 = make_block_id(2);

        dag.insert_block(block1, vec![genesis], 100).unwrap();
        dag.insert_block(block2, vec![genesis], 100).unwrap();
        reachability.add_block(block1, genesis, &[genesis]).unwrap();
        reachability.add_block(block2, genesis, &[genesis]).unwrap();

        ghostdag.add_block(block1, &[genesis]).unwrap();
        ghostdag.add_block(block2, &[genesis]).unwrap();

        // Extend block1's chain to make it have higher blue score
        let block3 = make_block_id(3);
        dag.insert_block(block3, vec![block1], 200).unwrap();
        reachability.add_block(block3, block1, &[block1]).unwrap();
        ghostdag.add_block(block3, &[block1]).unwrap();

        // Now create a block with both branches as parents
        let block4 = make_block_id(4);
        dag.insert_block(block4, vec![block3, block2], 300).unwrap();
        reachability
            .add_block(block4, block3, &[block3, block2])
            .unwrap();
        let data4 = ghostdag.add_block(block4, &[block3, block2]).unwrap();

        // Block3 should be selected parent (higher chain)
        assert_eq!(data4.selected_parent, block3);
    }

    #[test]
    fn test_blue_score_increases() {
        let genesis = make_block_id(0);
        let (dag, reachability, ghostdag) = setup_test_dag();

        let mut current = genesis;
        for i in 1..=5 {
            let block = make_block_id(i);
            dag.insert_block(block, vec![current], i as u64 * 100)
                .unwrap();
            reachability.add_block(block, current, &[current]).unwrap();
            let data = ghostdag.add_block(block, &[current]).unwrap();

            assert_eq!(data.blue_score, i as u64);
            current = block;
        }
    }

    #[test]
    fn test_selected_chain() {
        let genesis = make_block_id(0);
        let (dag, reachability, ghostdag) = setup_test_dag();

        // Build a chain
        let block1 = make_block_id(1);
        let block2 = make_block_id(2);
        let block3 = make_block_id(3);

        dag.insert_block(block1, vec![genesis], 100).unwrap();
        dag.insert_block(block2, vec![block1], 200).unwrap();
        dag.insert_block(block3, vec![block2], 300).unwrap();

        reachability.add_block(block1, genesis, &[genesis]).unwrap();
        reachability.add_block(block2, block1, &[block1]).unwrap();
        reachability.add_block(block3, block2, &[block2]).unwrap();

        ghostdag.add_block(block1, &[genesis]).unwrap();
        ghostdag.add_block(block2, &[block1]).unwrap();
        ghostdag.add_block(block3, &[block2]).unwrap();

        let chain = ghostdag.get_selected_chain(&block3).unwrap();
        assert_eq!(chain, vec![block3, block2, block1, genesis]);
    }
}