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 层级
    ///
    /// 返回 Vec<Vec<NodeId>>，每层内的节点可以并行执行。
    ///
    /// 复杂度：O(V+E)。一次性遍历边构建邻接索引（出边表）与入度表，
    /// BFS 分层只走索引查询（每条边仅被访问一次），不再为每个节点全表扫描边集合。
    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();
        // 邻接出边表：source → 直接后继列表（仅含两端均在 nodes 内的有效边）
        let mut adjacency_out: HashMap<NodeId, Vec<NodeId>> = HashMap::new();

        // 初始化入度（确保每个节点都有表项，便于后续 O(1) 修改）
        for id in &node_ids {
            in_degree.insert(id.clone(), 0);
        }
        // 单次 O(E) 遍历：同时构建入度表与出边表
        for edge in &self.edges {
            // 仅收录两端均为已注册节点的边，跳过野节点
            if node_ids.contains(&edge.source) && node_ids.contains(&edge.target) {
                *in_degree.get_mut(&edge.target).unwrap() += 1;
                adjacency_out
                    .entry(edge.source.clone())
                    .or_default()
                    .push(edge.target.clone());
            }
        }

        // 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());
                    // 出边已索引，O(出度) 遍历而非 O(E) 全表扫描
                    if let Some(succs) = adjacency_out.get(&id) {
                        for succ in succs {
                            let deg = in_degree.get_mut(succ).unwrap();
                            *deg -= 1;
                            if *deg == 0 {
                                queue.push_back(succ.clone());
                            }
                        }
                    }
                    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()
    }
}