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

consortium/dag/
graph.rs
