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

//! Core DAG data structure for block relationships.
//!
//! Maintains parent-child relationships between blocks and provides
//! efficient queries for traversal.

use crate::{BlockId, BlueScore, DaaScore, MAX_BLOCK_PARENTS};
use hashbrown::{HashMap, HashSet};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use thiserror::Error;

/// Relations for a single block in the DAG.
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct BlockRelations {
    /// Parent block IDs (blocks this block references).
    pub parents: Vec<BlockId>,
    /// Child block IDs (blocks that reference this block).
    pub children: Vec<BlockId>,
}

impl BlockRelations {
    /// Creates new relations with the given parents.
    pub fn new(parents: Vec<BlockId>) -> Self {
        BlockRelations {
            parents,
            children: Vec::new(),
        }
    }

    /// Returns true if this block has no parents (genesis).
    pub fn is_genesis(&self) -> bool {
        self.parents.is_empty()
    }

    /// Returns the number of parents.
    pub fn parent_count(&self) -> usize {
        self.parents.len()
    }

    /// Returns the number of children.
    pub fn child_count(&self) -> usize {
        self.children.len()
    }
}

/// Metadata stored for each block in the DAG.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct BlockMeta {
    /// Block relations (parents/children).
    pub relations: BlockRelations,
    /// Blue score (cumulative blue ancestor count).
    pub blue_score: BlueScore,
    /// DAA score (difficulty adjustment score).
    pub daa_score: DaaScore,
    /// Timestamp when block was received.
    pub timestamp: u64,
    /// Whether this block is in the blue set.
    pub is_blue: bool,
    /// The selected parent (main chain continuation).
    pub selected_parent: Option<BlockId>,
    /// Merge set blues (blue blocks merged at this block).
    pub merge_set_blues: Vec<BlockId>,
    /// Merge set reds (red blocks merged at this block).
    pub merge_set_reds: Vec<BlockId>,
}

impl BlockMeta {
    /// Creates metadata for a new block.
    pub fn new(parents: Vec<BlockId>, timestamp: u64) -> Self {
        BlockMeta {
            relations: BlockRelations::new(parents),
            blue_score: 0,
            daa_score: 0,
            timestamp,
            is_blue: false,
            selected_parent: None,
            merge_set_blues: Vec::new(),
            merge_set_reds: Vec::new(),
        }
    }

    /// Creates metadata for the genesis block.
    pub fn genesis(timestamp: u64) -> Self {
        BlockMeta {
            relations: BlockRelations::default(),
            blue_score: 0,
            daa_score: 0,
            timestamp,
            is_blue: true, // Genesis is always blue
            selected_parent: None,
            merge_set_blues: Vec::new(),
            merge_set_reds: Vec::new(),
        }
    }
}

/// The main DAG data structure.
pub struct BlockDag {
    /// All blocks in the DAG.
    blocks: RwLock<HashMap<BlockId, BlockMeta>>,
    /// Tips of the DAG (blocks with no children).
    tips: RwLock<HashSet<BlockId>>,
    /// The genesis block ID.
    genesis: BlockId,
    /// Virtual genesis (for pruning).
    virtual_genesis: RwLock<BlockId>,
}

impl BlockDag {
    /// Creates a new DAG with the given genesis block.
    pub fn new(genesis_id: BlockId, genesis_timestamp: u64) -> Self {
        let mut blocks = HashMap::new();
        blocks.insert(genesis_id, BlockMeta::genesis(genesis_timestamp));

        let mut tips = HashSet::new();
        tips.insert(genesis_id);

        BlockDag {
            blocks: RwLock::new(blocks),
            tips: RwLock::new(tips),
            genesis: genesis_id,
            virtual_genesis: RwLock::new(genesis_id),
        }
    }

    /// Returns the genesis block ID.
    pub fn genesis(&self) -> BlockId {
        self.genesis
    }

    /// Returns the current virtual genesis (for pruned DAGs).
    pub fn virtual_genesis(&self) -> BlockId {
        *self.virtual_genesis.read()
    }

    /// Returns the current tips (blocks with no children).
    pub fn tips(&self) -> Vec<BlockId> {
        self.tips.read().iter().copied().collect()
    }

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

    /// Returns true if the DAG is empty (only genesis).
    pub fn is_empty(&self) -> bool {
        self.len() <= 1
    }

    /// Checks if a block exists in the DAG.
    pub fn contains(&self, block_id: &BlockId) -> bool {
        self.blocks.read().contains_key(block_id)
    }

    /// Gets metadata for a block.
    pub fn get_block(&self, block_id: &BlockId) -> Option<BlockMeta> {
        self.blocks.read().get(block_id).cloned()
    }

    /// Gets the parents of a block.
    pub fn get_parents(&self, block_id: &BlockId) -> Option<Vec<BlockId>> {
        self.blocks
            .read()
            .get(block_id)
            .map(|meta| meta.relations.parents.clone())
    }

    /// Gets the children of a block.
    pub fn get_children(&self, block_id: &BlockId) -> Option<Vec<BlockId>> {
        self.blocks
            .read()
            .get(block_id)
            .map(|meta| meta.relations.children.clone())
    }

    /// Gets the blue score of a block.
    pub fn get_blue_score(&self, block_id: &BlockId) -> Option<BlueScore> {
        self.blocks.read().get(block_id).map(|meta| meta.blue_score)
    }

    /// Gets the selected parent of a block.
    pub fn get_selected_parent(&self, block_id: &BlockId) -> Option<BlockId> {
        self.blocks
            .read()
            .get(block_id)
            .and_then(|meta| meta.selected_parent)
    }

    /// Inserts a new block into the DAG.
    ///
    /// Returns an error if:
    /// - Block already exists
    /// - Any parent doesn't exist
    /// - Too many parents
    pub fn insert_block(
        &self,
        block_id: BlockId,
        parents: Vec<BlockId>,
        timestamp: u64,
    ) -> Result<(), DagError> {
        // Validate parent count
        if parents.len() > MAX_BLOCK_PARENTS {
            return Err(DagError::TooManyParents {
                count: parents.len(),
                max: MAX_BLOCK_PARENTS,
            });
        }

        // Validate parents exist (except for genesis check)
        if parents.is_empty() && block_id != self.genesis {
            return Err(DagError::NoParents);
        }

        let mut blocks = self.blocks.write();
        let mut tips = self.tips.write();

        // Check block doesn't already exist
        if blocks.contains_key(&block_id) {
            return Err(DagError::BlockExists(block_id));
        }

        // Verify all parents exist
        for parent in &parents {
            if !blocks.contains_key(parent) {
                return Err(DagError::ParentNotFound(*parent));
            }
        }

        // Create block metadata
        let meta = BlockMeta::new(parents.clone(), timestamp);
        blocks.insert(block_id, meta);

        // Update parent-child relationships
        for parent in &parents {
            if let Some(parent_meta) = blocks.get_mut(parent) {
                parent_meta.relations.children.push(block_id);
            }
            // Parent is no longer a tip
            tips.remove(parent);
        }

        // New block is a tip
        tips.insert(block_id);

        Ok(())
    }

    /// Updates GHOSTDAG data for a block.
    pub fn update_ghostdag_data(
        &self,
        block_id: &BlockId,
        blue_score: BlueScore,
        selected_parent: Option<BlockId>,
        is_blue: bool,
        merge_set_blues: Vec<BlockId>,
        merge_set_reds: Vec<BlockId>,
    ) -> Result<(), DagError> {
        let mut blocks = self.blocks.write();
        let meta = blocks
            .get_mut(block_id)
            .ok_or(DagError::BlockNotFound(*block_id))?;

        meta.blue_score = blue_score;
        meta.selected_parent = selected_parent;
        meta.is_blue = is_blue;
        meta.merge_set_blues = merge_set_blues;
        meta.merge_set_reds = merge_set_reds;

        Ok(())
    }

    /// Returns all ancestors of a block up to a given depth.
    pub fn get_ancestors(&self, block_id: &BlockId, max_depth: usize) -> HashSet<BlockId> {
        let blocks = self.blocks.read();
        let mut ancestors = HashSet::new();
        let mut frontier = vec![*block_id];
        let mut depth = 0;

        while depth < max_depth && !frontier.is_empty() {
            let mut next_frontier = Vec::new();
            for current in frontier {
                if let Some(meta) = blocks.get(&current) {
                    for parent in &meta.relations.parents {
                        if ancestors.insert(*parent) {
                            next_frontier.push(*parent);
                        }
                    }
                }
            }
            frontier = next_frontier;
            depth += 1;
        }

        ancestors
    }

    /// Returns all descendants of a block.
    pub fn get_descendants(&self, block_id: &BlockId) -> HashSet<BlockId> {
        let blocks = self.blocks.read();
        let mut descendants = HashSet::new();
        let mut frontier = vec![*block_id];

        while !frontier.is_empty() {
            let mut next_frontier = Vec::new();
            for current in frontier {
                if let Some(meta) = blocks.get(&current) {
                    for child in &meta.relations.children {
                        if descendants.insert(*child) {
                            next_frontier.push(*child);
                        }
                    }
                }
            }
            frontier = next_frontier;
        }

        descendants
    }

    /// Finds the anticone of a block (blocks neither ancestor nor descendant).
    pub fn get_anticone(&self, block_id: &BlockId, depth_limit: usize) -> HashSet<BlockId> {
        let blocks = self.blocks.read();
        let ancestors = self.get_ancestors(block_id, depth_limit);
        let descendants = self.get_descendants(block_id);

        let mut anticone = HashSet::new();
        for (id, _) in blocks.iter() {
            if *id != *block_id && !ancestors.contains(id) && !descendants.contains(id) {
                anticone.insert(*id);
            }
        }

        anticone
    }

    /// Gets the selected chain from a block back to genesis.
    pub fn get_selected_chain(&self, from: &BlockId) -> Vec<BlockId> {
        let blocks = self.blocks.read();
        let mut chain = vec![*from];
        let mut current = *from;

        while let Some(meta) = blocks.get(&current) {
            if let Some(selected_parent) = meta.selected_parent {
                chain.push(selected_parent);
                current = selected_parent;
            } else {
                break;
            }
        }

        chain
    }

    /// Iterates over all blocks in topological order (parents before children).
    pub fn iter_topological(&self) -> Vec<BlockId> {
        let blocks = self.blocks.read();
        let mut result = Vec::with_capacity(blocks.len());
        let mut in_degree: HashMap<BlockId, usize> = HashMap::new();
        let mut queue = Vec::new();

        // Calculate in-degrees
        for (id, meta) in blocks.iter() {
            let degree = meta.relations.parents.len();
            in_degree.insert(*id, degree);
            if degree == 0 {
                queue.push(*id);
            }
        }

        // Process in topological order
        while let Some(current) = queue.pop() {
            result.push(current);
            if let Some(meta) = blocks.get(&current) {
                for child in &meta.relations.children {
                    if let Some(degree) = in_degree.get_mut(child) {
                        *degree -= 1;
                        if *degree == 0 {
                            queue.push(*child);
                        }
                    }
                }
            }
        }

        result
    }
}

impl std::fmt::Debug for BlockDag {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("BlockDag")
            .field("block_count", &self.len())
            .field("tip_count", &self.tips.read().len())
            .field("genesis", &self.genesis)
            .finish()
    }
}

/// Errors that can occur with the DAG.
#[derive(Debug, Error)]
pub enum DagError {
    #[error("Block already exists: {0}")]
    BlockExists(BlockId),

    #[error("Block not found: {0}")]
    BlockNotFound(BlockId),

    #[error("Parent not found: {0}")]
    ParentNotFound(BlockId),

    #[error("No parents specified (non-genesis block)")]
    NoParents,

    #[error("Too many parents: {count} (max {max})")]
    TooManyParents { count: usize, max: usize },

    #[error("Cycle detected in DAG")]
    CycleDetected,

    #[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)
    }

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

        assert_eq!(dag.genesis(), genesis_id);
        assert_eq!(dag.len(), 1);
        assert!(dag.contains(&genesis_id));
    }

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

        let block1 = make_block_id(1);
        dag.insert_block(block1, vec![genesis_id], 100).unwrap();

        assert_eq!(dag.len(), 2);
        assert!(dag.contains(&block1));

        let parents = dag.get_parents(&block1).unwrap();
        assert_eq!(parents.len(), 1);
        assert_eq!(parents[0], genesis_id);
    }

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

        // Initially, genesis is the only tip
        assert_eq!(dag.tips().len(), 1);
        assert!(dag.tips().contains(&genesis_id));

        // Add a child block
        let block1 = make_block_id(1);
        dag.insert_block(block1, vec![genesis_id], 100).unwrap();

        // Now only block1 is a tip
        assert_eq!(dag.tips().len(), 1);
        assert!(dag.tips().contains(&block1));
        assert!(!dag.tips().contains(&genesis_id));

        // Add another block with genesis as parent (parallel mining)
        let block2 = make_block_id(2);
        dag.insert_block(block2, vec![genesis_id], 100).unwrap();

        // Now we have two tips
        assert_eq!(dag.tips().len(), 2);
        assert!(dag.tips().contains(&block1));
        assert!(dag.tips().contains(&block2));
    }

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

        let block1 = make_block_id(1);
        let block2 = make_block_id(2);

        dag.insert_block(block1, vec![genesis_id], 100).unwrap();
        dag.insert_block(block2, vec![genesis_id], 100).unwrap();

        // Block with multiple parents
        let block3 = make_block_id(3);
        dag.insert_block(block3, vec![block1, block2], 200).unwrap();

        let parents = dag.get_parents(&block3).unwrap();
        assert_eq!(parents.len(), 2);
        assert!(parents.contains(&block1));
        assert!(parents.contains(&block2));
    }

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

        let block1 = make_block_id(1);
        let block2 = make_block_id(2);
        let block3 = make_block_id(3);

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

        let ancestors = dag.get_ancestors(&block3, 10);
        assert!(ancestors.contains(&block2));
        assert!(ancestors.contains(&block1));
        assert!(ancestors.contains(&genesis_id));
    }

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

        let block1 = make_block_id(1);
        dag.insert_block(block1, vec![genesis_id], 100).unwrap();

        // Try to insert same block again
        let result = dag.insert_block(block1, vec![genesis_id], 100);
        assert!(matches!(result, Err(DagError::BlockExists(_))));
    }

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

        let missing_parent = make_block_id(99);
        let block1 = make_block_id(1);

        let result = dag.insert_block(block1, vec![missing_parent], 100);
        assert!(matches!(result, Err(DagError::ParentNotFound(_))));
    }

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

        let block1 = make_block_id(1);
        let block2 = make_block_id(2);
        let block3 = make_block_id(3);

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

        let order = dag.iter_topological();

        // Genesis should come first
        assert_eq!(order[0], genesis_id);

        // Check that parents come before children
        let idx_genesis = order.iter().position(|&x| x == genesis_id).unwrap();
        let idx_block1 = order.iter().position(|&x| x == block1).unwrap();
        let idx_block2 = order.iter().position(|&x| x == block2).unwrap();
        let idx_block3 = order.iter().position(|&x| x == block3).unwrap();

        assert!(idx_genesis < idx_block1);
        assert!(idx_block1 < idx_block2);
        assert!(idx_block2 < idx_block3);
    }
}