DevFlow/crates/df-workflow/src/dag.rs

//! DAG（有向无环图）定义与拓扑排序

use std::collections::{HashMap, HashSet, VecDeque};

use df_core::error::Result;
use df_core::types::NodeId;

use crate::node::Node;

/// 图的边：从 source 到 target
#[derive(Debug, Clone)]
pub struct Edge {
    pub source: NodeId,
    pub target: NodeId,
    /// 边的条件表达式（可选）
    pub condition: Option<String>,
}

/// DAG 结构
pub struct Dag {
    /// 节点集合
    pub nodes: HashMap<NodeId, Box<dyn Node>>,
    /// 边集合
    pub edges: Vec<Edge>,
}

impl Dag {
    /// 创建空 DAG
    pub fn new() -> Self {
        Self {
            nodes: HashMap::new(),
            edges: Vec::new(),
        }
    }

    /// 添加节点
    pub fn add_node(&mut self, id: NodeId, node: Box<dyn Node>) {
        self.nodes.insert(id, node);
    }

    /// 添加边
    pub fn add_edge(&mut self, source: NodeId, target: NodeId) {
        self.edges.push(Edge {
            source,
            target,
            condition: None,
        });
    }

    /// 添加带条件的边
    pub fn add_edge_with_condition(
        &mut self,
        source: NodeId,
        target: NodeId,
        condition: String,
    ) {
        self.edges.push(Edge {
            source,
            target,
            condition: Some(condition),
        });
    }

    /// 获取指定节点的所有上游节点 ID
    pub fn predecessors(&self, node_id: &NodeId) -> Vec<NodeId> {
        self.edges
            .iter()
            .filter(|e| &e.target == node_id)
            .map(|e| e.source.clone())
            .collect()
    }

    /// 获取指定节点的所有下游节点 ID
    pub fn successors(&self, node_id: &NodeId) -> Vec<NodeId> {
        self.edges
            .iter()
            .filter(|e| &e.source == node_id)
            .map(|e| e.target.clone())
            .collect()
    }

    /// 拓扑排序 — 返回按执行顺序排列的节点 ID 层级
    ///
    /// 返回 Vec<Vec<NodeId>>，每层内的节点可以并行执行
    pub fn topological_layers(&self) -> Result<Vec<Vec<NodeId>>> {
        let node_ids: HashSet<NodeId> = self.nodes.keys().cloned().collect();
        let mut in_degree: HashMap<NodeId, usize> = HashMap::new();

        // 初始化入度
        for id in &node_ids {
            in_degree.insert(id.clone(), 0);
        }
        for edge in &self.edges {
            if node_ids.contains(&edge.source) && node_ids.contains(&edge.target) {
                *in_degree.entry(edge.target.clone()).or_insert(0) += 1;
            }
        }

        // BFS 分层
        let mut layers = Vec::new();
        let mut queue: VecDeque<NodeId> = in_degree
            .iter()
            .filter(|(_, &deg)| deg == 0)
            .map(|(id, _)| id.clone())
            .collect();

        let mut processed = 0usize;
        while !queue.is_empty() {
            let mut layer: Vec<NodeId> = Vec::new();
            let layer_size = queue.len();
            for _ in 0..layer_size {
                if let Some(id) = queue.pop_front() {
                    layer.push(id.clone());
                    for succ in self.successors(&id) {
                        let deg = in_degree.get_mut(&succ).unwrap();
                        *deg -= 1;
                        if *deg == 0 {
                            queue.push_back(succ);
                        }
                    }
                    processed += 1;
                }
            }
            layers.push(layer);
        }

        if processed != node_ids.len() {
            return Err(df_core::error::Error::Workflow(
                "DAG 中存在环，无法进行拓扑排序".to_string(),
            ));
        }

        Ok(layers)
    }
}

impl Default for Dag {
    fn default() -> Self {
        Self::new()
    }
}