consortium/dag/
graph.rs

1//! DAG data structure with validation and topological sorting.
2
3use std::collections::{HashMap, HashSet, VecDeque};
4
5use crate::dag::cache::CacheStrategy;
6use crate::dag::error::{DagError, Result};
7use crate::dag::types::{ConcurrencyLimit, DagTask, TaskId};
8
9/// Internal DAG representation using adjacency lists.
10pub struct DagGraph {
11    /// All tasks keyed by ID.
12    pub(crate) tasks: HashMap<TaskId, Box<dyn DagTask>>,
13    /// Forward edges: task -> set of tasks that depend on it (dependents).
14    pub(crate) dependents: HashMap<TaskId, HashSet<TaskId>>,
15    /// Reverse edges: task -> set of tasks it depends on (dependencies).
16    pub(crate) dependencies: HashMap<TaskId, HashSet<TaskId>>,
17    /// Concurrency groups with their limits.
18    pub(crate) concurrency_groups: HashMap<String, ConcurrencyLimit>,
19    /// Which concurrency group each task belongs to.
20    pub(crate) task_groups: HashMap<TaskId, String>,
21    /// Per-task cache strategy.
22    pub(crate) task_cache: HashMap<TaskId, Box<dyn CacheStrategy>>,
23}
24
25impl DagGraph {
26    /// Create an empty graph.
27    pub fn new() -> Self {
28        Self {
29            tasks: HashMap::new(),
30            dependents: HashMap::new(),
31            dependencies: HashMap::new(),
32            concurrency_groups: HashMap::new(),
33            task_groups: HashMap::new(),
34            task_cache: HashMap::new(),
35        }
36    }
37
38    /// Add a task to the graph.
39    pub fn add_task(&mut self, id: TaskId, task: Box<dyn DagTask>) -> Result<()> {
40        if self.tasks.contains_key(&id) {
41            return Err(DagError::DuplicateTask(id));
42        }
43        self.tasks.insert(id.clone(), task);
44        self.dependents.entry(id.clone()).or_default();
45        self.dependencies.entry(id).or_default();
46        Ok(())
47    }
48
49    /// Add a dependency edge: `dependent` cannot start until `dependency` succeeds.
50    pub fn add_dep(&mut self, dependent: &TaskId, dependency: &TaskId) -> Result<()> {
51        if !self.tasks.contains_key(dependent) {
52            return Err(DagError::UnknownTask(dependent.clone()));
53        }
54        if !self.tasks.contains_key(dependency) {
55            return Err(DagError::UnknownTask(dependency.clone()));
56        }
57        self.dependents
58            .entry(dependency.clone())
59            .or_default()
60            .insert(dependent.clone());
61        self.dependencies
62            .entry(dependent.clone())
63            .or_default()
64            .insert(dependency.clone());
65        Ok(())
66    }
67
68    /// Set a concurrency group with a limit.
69    pub fn set_concurrency_group(&mut self, name: &str, max_concurrent: usize) {
70        self.concurrency_groups.insert(
71            name.to_string(),
72            ConcurrencyLimit {
73                max_concurrent: Some(max_concurrent),
74            },
75        );
76    }
77
78    /// Assign a task to a concurrency group.
79    pub fn assign_group(&mut self, task_id: &TaskId, group: &str) {
80        self.task_groups.insert(task_id.clone(), group.to_string());
81    }
82
83    /// Validate the graph: check for cycles using Kahn's algorithm.
84    pub fn validate(&self) -> Result<()> {
85        if self.tasks.is_empty() {
86            return Err(DagError::EmptyDag);
87        }
88
89        // Kahn's algorithm for cycle detection
90        let mut in_degree: HashMap<&TaskId, usize> = HashMap::new();
91        for id in self.tasks.keys() {
92            let deps = self.dependencies.get(id).map(|d| d.len()).unwrap_or(0);
93            in_degree.insert(id, deps);
94        }
95
96        let mut queue: VecDeque<&TaskId> = in_degree
97            .iter()
98            .filter(|(_, &deg)| deg == 0)
99            .map(|(&id, _)| id)
100            .collect();
101
102        let mut visited = 0;
103
104        while let Some(id) = queue.pop_front() {
105            visited += 1;
106            if let Some(deps) = self.dependents.get(id) {
107                for dep in deps {
108                    if let Some(degree) = in_degree.get_mut(dep) {
109                        *degree -= 1;
110                        if *degree == 0 {
111                            queue.push_back(dep);
112                        }
113                    }
114                }
115            }
116        }
117
118        if visited != self.tasks.len() {
119            // Find a task involved in the cycle
120            let cycle_task = in_degree
121                .iter()
122                .find(|(_, &deg)| deg > 0)
123                .map(|(&id, _)| id.clone())
124                .unwrap_or_else(|| TaskId("unknown".into()));
125            return Err(DagError::CycleDetected(cycle_task));
126        }
127
128        Ok(())
129    }
130
131    /// Return task IDs with zero dependencies (roots/starting points).
132    pub fn roots(&self) -> Vec<TaskId> {
133        self.tasks
134            .keys()
135            .filter(|id| {
136                self.dependencies
137                    .get(*id)
138                    .map(|d| d.is_empty())
139                    .unwrap_or(true)
140            })
141            .cloned()
142            .collect()
143    }
144
145    /// Return a topological ordering of task IDs.
146    pub fn topo_sort(&self) -> Result<Vec<TaskId>> {
147        let mut in_degree: HashMap<&TaskId, usize> = HashMap::new();
148        for id in self.tasks.keys() {
149            let deps = self.dependencies.get(id).map(|d| d.len()).unwrap_or(0);
150            in_degree.insert(id, deps);
151        }
152
153        let mut queue: VecDeque<&TaskId> = in_degree
154            .iter()
155            .filter(|(_, &deg)| deg == 0)
156            .map(|(&id, _)| id)
157            .collect();
158
159        let mut order = Vec::new();
160
161        while let Some(id) = queue.pop_front() {
162            order.push(id.clone());
163            if let Some(deps) = self.dependents.get(id) {
164                for dep in deps {
165                    if let Some(degree) = in_degree.get_mut(dep) {
166                        *degree -= 1;
167                        if *degree == 0 {
168                            queue.push_back(dep);
169                        }
170                    }
171                }
172            }
173        }
174
175        if order.len() != self.tasks.len() {
176            let cycle_task = in_degree
177                .iter()
178                .find(|(_, &deg)| deg > 0)
179                .map(|(&id, _)| id.clone())
180                .unwrap_or_else(|| TaskId("unknown".into()));
181            return Err(DagError::CycleDetected(cycle_task));
182        }
183
184        Ok(order)
185    }
186
187    /// Take a task out of the graph (for dispatching to a worker thread).
188    pub fn take_task(&mut self, id: &TaskId) -> Option<Box<dyn DagTask>> {
189        self.tasks.remove(id)
190    }
191
192    /// Number of tasks.
193    pub fn len(&self) -> usize {
194        self.tasks.len()
195    }
196
197    /// Whether the graph is empty.
198    pub fn is_empty(&self) -> bool {
199        self.tasks.is_empty()
200    }
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206    use crate::dag::context::DagContext;
207    use crate::dag::types::{FnTask, TaskOutcome};
208
209    fn noop_task(name: &str) -> Box<dyn DagTask> {
210        let n = name.to_string();
211        Box::new(FnTask::new(n, |_| TaskOutcome::Success))
212    }
213
214    #[test]
215    fn test_empty_dag() {
216        let g = DagGraph::new();
217        assert!(matches!(g.validate(), Err(DagError::EmptyDag)));
218    }
219
220    #[test]
221    fn test_single_task() {
222        let mut g = DagGraph::new();
223        g.add_task(TaskId::from("a"), noop_task("a")).unwrap();
224        assert!(g.validate().is_ok());
225        assert_eq!(g.roots(), vec![TaskId::from("a")]);
226    }
227
228    #[test]
229    fn test_linear_chain() {
230        let mut g = DagGraph::new();
231        g.add_task(TaskId::from("a"), noop_task("a")).unwrap();
232        g.add_task(TaskId::from("b"), noop_task("b")).unwrap();
233        g.add_task(TaskId::from("c"), noop_task("c")).unwrap();
234        g.add_dep(&TaskId::from("b"), &TaskId::from("a")).unwrap();
235        g.add_dep(&TaskId::from("c"), &TaskId::from("b")).unwrap();
236
237        assert!(g.validate().is_ok());
238        assert_eq!(g.roots(), vec![TaskId::from("a")]);
239
240        let order = g.topo_sort().unwrap();
241        let a_pos = order.iter().position(|t| t.0 == "a").unwrap();
242        let b_pos = order.iter().position(|t| t.0 == "b").unwrap();
243        let c_pos = order.iter().position(|t| t.0 == "c").unwrap();
244        assert!(a_pos < b_pos);
245        assert!(b_pos < c_pos);
246    }
247
248    #[test]
249    fn test_diamond_dag() {
250        let mut g = DagGraph::new();
251        g.add_task(TaskId::from("a"), noop_task("a")).unwrap();
252        g.add_task(TaskId::from("b"), noop_task("b")).unwrap();
253        g.add_task(TaskId::from("c"), noop_task("c")).unwrap();
254        g.add_task(TaskId::from("d"), noop_task("d")).unwrap();
255        g.add_dep(&TaskId::from("b"), &TaskId::from("a")).unwrap();
256        g.add_dep(&TaskId::from("c"), &TaskId::from("a")).unwrap();
257        g.add_dep(&TaskId::from("d"), &TaskId::from("b")).unwrap();
258        g.add_dep(&TaskId::from("d"), &TaskId::from("c")).unwrap();
259
260        assert!(g.validate().is_ok());
261        let order = g.topo_sort().unwrap();
262        let a_pos = order.iter().position(|t| t.0 == "a").unwrap();
263        let d_pos = order.iter().position(|t| t.0 == "d").unwrap();
264        assert!(a_pos < d_pos);
265    }
266
267    #[test]
268    fn test_cycle_detection() {
269        let mut g = DagGraph::new();
270        g.add_task(TaskId::from("a"), noop_task("a")).unwrap();
271        g.add_task(TaskId::from("b"), noop_task("b")).unwrap();
272        g.add_dep(&TaskId::from("b"), &TaskId::from("a")).unwrap();
273        g.add_dep(&TaskId::from("a"), &TaskId::from("b")).unwrap();
274
275        assert!(matches!(g.validate(), Err(DagError::CycleDetected(_))));
276    }
277
278    #[test]
279    fn test_duplicate_task() {
280        let mut g = DagGraph::new();
281        g.add_task(TaskId::from("a"), noop_task("a")).unwrap();
282        assert!(matches!(
283            g.add_task(TaskId::from("a"), noop_task("a")),
284            Err(DagError::DuplicateTask(_))
285        ));
286    }
287
288    #[test]
289    fn test_unknown_dep() {
290        let mut g = DagGraph::new();
291        g.add_task(TaskId::from("a"), noop_task("a")).unwrap();
292        assert!(matches!(
293            g.add_dep(&TaskId::from("a"), &TaskId::from("missing")),
294            Err(DagError::UnknownTask(_))
295        ));
296    }
297
298    #[test]
299    fn test_self_loop_cycle() {
300        let mut g = DagGraph::new();
301        g.add_task(TaskId::from("a"), noop_task("a")).unwrap();
302        g.add_dep(&TaskId::from("a"), &TaskId::from("a")).unwrap();
303        assert!(matches!(g.validate(), Err(DagError::CycleDetected(_))));
304    }
305
306    #[test]
307    fn test_deep_cycle() {
308        let mut g = DagGraph::new();
309        for name in &["a", "b", "c", "d", "e"] {
310            g.add_task(TaskId::from(*name), noop_task(name)).unwrap();
311        }
312        g.add_dep(&TaskId::from("b"), &TaskId::from("a")).unwrap();
313        g.add_dep(&TaskId::from("c"), &TaskId::from("b")).unwrap();
314        g.add_dep(&TaskId::from("d"), &TaskId::from("c")).unwrap();
315        g.add_dep(&TaskId::from("e"), &TaskId::from("d")).unwrap();
316        g.add_dep(&TaskId::from("a"), &TaskId::from("e")).unwrap(); // closes the cycle
317        assert!(matches!(g.validate(), Err(DagError::CycleDetected(_))));
318    }
319
320    #[test]
321    fn test_large_dag_performance() {
322        let mut g = DagGraph::new();
323        // 1000 tasks in a linear chain — should validate in <100ms
324        for i in 0..1000 {
325            let id = format!("t{}", i);
326            g.add_task(TaskId::from(id.as_str()), noop_task(&id))
327                .unwrap();
328            if i > 0 {
329                let prev = format!("t{}", i - 1);
330                g.add_dep(&TaskId::from(id.as_str()), &TaskId::from(prev.as_str()))
331                    .unwrap();
332            }
333        }
334        let start = std::time::Instant::now();
335        assert!(g.validate().is_ok());
336        let elapsed = start.elapsed();
337        assert!(
338            elapsed < std::time::Duration::from_secs(1),
339            "1000-task DAG validation took {:?}",
340            elapsed
341        );
342        let order = g.topo_sort().unwrap();
343        assert_eq!(order.len(), 1000);
344        assert_eq!(order[0], TaskId::from("t0"));
345        assert_eq!(order[999], TaskId::from("t999"));
346    }
347}
consortium/dag/graph.rs

consortium/dag/
graph.rs
