radicle_dag/
lib.rs

1//! Directed-acyclic graph implementation.
2#![warn(missing_docs)]
3use std::{
4    borrow::Borrow,
5    cmp::Ordering,
6    collections::{BTreeMap, BTreeSet, VecDeque},
7    fmt,
8    fmt::Write,
9    ops::{ControlFlow, Deref, Index},
10};
11
12/// A node in the graph.
13#[derive(Clone, Debug, PartialEq, Eq)]
14pub struct Node<K, V> {
15    /// The node key.
16    pub key: K,
17    /// The node value, stored by the user.
18    pub value: V,
19    /// Nodes depended on.
20    pub dependencies: BTreeSet<K>,
21    /// Nodes depending on this node.
22    pub dependents: BTreeSet<K>,
23}
24
25impl<K, V> Node<K, V> {
26    fn new(key: K, value: V) -> Self {
27        Self {
28            key,
29            value,
30            dependencies: BTreeSet::new(),
31            dependents: BTreeSet::new(),
32        }
33    }
34}
35
36impl<K, V> Borrow<V> for &Node<K, V> {
37    fn borrow(&self) -> &V {
38        &self.value
39    }
40}
41
42impl<K, V> Deref for Node<K, V> {
43    type Target = V;
44
45    fn deref(&self) -> &Self::Target {
46        &self.value
47    }
48}
49
50/// A directed acyclic graph.
51#[derive(Clone, Debug, Default, PartialEq, Eq)]
52pub struct Dag<K, V> {
53    graph: BTreeMap<K, Node<K, V>>,
54    tips: BTreeSet<K>,
55    roots: BTreeSet<K>,
56}
57
58impl<K: Ord + Copy, V> Dag<K, V> {
59    /// Create a new empty DAG.
60    pub fn new() -> Self {
61        Self {
62            graph: BTreeMap::new(),
63            tips: BTreeSet::new(),
64            roots: BTreeSet::new(),
65        }
66    }
67
68    /// Create a DAG with a root node.
69    pub fn root(key: K, value: V) -> Self {
70        Self {
71            graph: BTreeMap::from_iter([(key, Node::new(key, value))]),
72            tips: BTreeSet::from_iter([key]),
73            roots: BTreeSet::from_iter([key]),
74        }
75    }
76
77    /// Check whether there are any nodes in the graph.
78    pub fn is_empty(&self) -> bool {
79        self.graph.is_empty()
80    }
81
82    /// Return the number of nodes in the graph.
83    pub fn len(&self) -> usize {
84        self.graph.len()
85    }
86
87    /// Add a node to the graph.
88    pub fn node(&mut self, key: K, value: V) -> Option<Node<K, V>> {
89        self.tips.insert(key);
90        self.roots.insert(key);
91        self.graph.insert(
92            key,
93            Node {
94                key,
95                value,
96                dependencies: BTreeSet::new(),
97                dependents: BTreeSet::new(),
98            },
99        )
100    }
101
102    /// Add a dependency from one node to the other.
103    pub fn dependency(&mut self, from: K, to: K) {
104        if let Some(node) = self.graph.get_mut(&from) {
105            node.dependencies.insert(to);
106            self.roots.remove(&from);
107        }
108        if let Some(node) = self.graph.get_mut(&to) {
109            node.dependents.insert(from);
110            self.tips.remove(&to);
111        }
112    }
113
114    /// Check if the graph contains a node.
115    pub fn contains(&self, key: &K) -> bool {
116        self.graph.contains_key(key)
117    }
118
119    /// Get a node.
120    pub fn get(&self, key: &K) -> Option<&Node<K, V>> {
121        self.graph.get(key)
122    }
123
124    /// Check whether there is a dependency between two nodes.
125    pub fn has_dependency(&self, from: &K, to: &K) -> bool {
126        self.graph
127            .get(from)
128            .map(|n| n.dependencies.contains(to))
129            .unwrap_or_default()
130    }
131
132    /// Get the graph's root nodes, ie. nodes which don't depend on other nodes.
133    pub fn roots(&self) -> impl Iterator<Item = (&K, &Node<K, V>)> + '_ {
134        self.roots
135            .iter()
136            .filter_map(|k| self.graph.get(k).map(|n| (k, n)))
137    }
138
139    /// Get the graph's tip nodes, ie. nodes which aren't depended on by other nodes.
140    pub fn tips(&self) -> impl Iterator<Item = (&K, &Node<K, V>)> + '_ {
141        self.tips
142            .iter()
143            .filter_map(|k| self.graph.get(k).map(|n| (k, n)))
144    }
145
146    /// Merge a DAG into this one.
147    pub fn merge(&mut self, mut other: Self) {
148        let Some((root, _)) = other.roots().next() else {
149            return;
150        };
151        let mut visited = BTreeSet::new();
152        let mut queue = VecDeque::<K>::from([*root]);
153
154        while let Some(next) = queue.pop_front() {
155            if !visited.insert(next) {
156                continue;
157            }
158            if let Some(node) = other.graph.remove(&next) {
159                if !self.contains(&next) {
160                    self.node(next, node.value);
161                }
162                for k in &node.dependents {
163                    self.dependency(*k, next);
164                }
165                for k in &node.dependencies {
166                    self.dependency(next, *k);
167                }
168                queue.extend(node.dependents.iter());
169            }
170        }
171    }
172
173    /// Return a topological ordering of the graph's nodes.
174    pub fn sorted(&self) -> VecDeque<K> {
175        self.sorted_by(Ord::cmp)
176    }
177
178    /// Return a topological ordering of the graph's nodes.
179    /// Uses a comparison function to sort partially ordered nodes.
180    pub fn sorted_by<F>(&self, mut compare: F) -> VecDeque<K>
181    where
182        F: FnMut(&K, &K) -> Ordering,
183    {
184        let mut order = VecDeque::new(); // Stores the topological order.
185        let mut visited = BTreeSet::new(); // Nodes that have been visited.
186        let mut keys = self.graph.keys().collect::<Vec<_>>();
187
188        // The `visit` function builds the list in reverse order, so we counter-act
189        // that here.
190        keys.sort_by(|a, b| compare(a, b).reverse());
191
192        for node in keys {
193            self.visit(node, &mut visited, &mut order);
194        }
195        order
196    }
197
198    /// Fold over the graph in topological order, pruning branches along the way.
199    /// This is a depth-first traversal.
200    ///
201    /// To continue traversing a branch, return [`ControlFlow::Continue`] from the
202    /// filter function. To stop traversal of a branch and prune it,
203    /// return [`ControlFlow::Break`].
204    pub fn prune<F>(&mut self, roots: &[K], filter: F)
205    where
206        F: for<'r> FnMut(
207            &'r K,
208            &'r Node<K, V>,
209            Box<dyn Iterator<Item = (&'r K, &'r Node<K, V>)> + 'r>,
210        ) -> ControlFlow<()>,
211    {
212        self.prune_by(roots, filter, |(k1, _), (k2, _)| k1.cmp(k2))
213    }
214
215    /// Fold over the graph in the order provided by `order`, pruning branches
216    /// along the way.
217    /// This is a depth-first traversal.
218    ///
219    /// To continue traversing a branch, return [`ControlFlow::Continue`] from the
220    /// filter function. To stop traversal of a branch and prune it,
221    /// return [`ControlFlow::Break`].
222    pub fn prune_by<F, P>(&mut self, roots: &[K], mut filter: F, ordering: P)
223    where
224        F: for<'r> FnMut(
225            &'r K,
226            &'r Node<K, V>,
227            Box<dyn Iterator<Item = (&'r K, &'r Node<K, V>)> + 'r>,
228        ) -> ControlFlow<()>,
229        P: Fn((&K, &V), (&K, &V)) -> Ordering,
230    {
231        let mut visited = BTreeSet::new();
232        let mut result = VecDeque::new();
233
234        for root in roots {
235            self.visit_by(root, &mut visited, &mut result, &ordering);
236        }
237
238        for next in result {
239            if let Some(node) = self.graph.get(&next) {
240                let siblings = self
241                    .siblings_of(node)
242                    .filter_map(|k| self.graph.get(k))
243                    .map(|node| (&node.key, node));
244
245                match filter(&next, node, Box::new(siblings)) {
246                    ControlFlow::Continue(()) => {}
247                    ControlFlow::Break(()) => {
248                        // When pruning a node, we remove all transitive dependents on
249                        // that node.
250                        self.remove(&next);
251                    }
252                }
253            }
254        }
255    }
256
257    /// Fold over the graph in topological order, skipping certain branches.
258    /// This is a depth-first traversal.
259    ///
260    /// To continue traversing a branch, return [`ControlFlow::Continue`] from the
261    /// filter function. To stop traversal of a branch, return [`ControlFlow::Break`].
262    pub fn fold<A, F>(&self, roots: &[K], mut acc: A, mut filter: F) -> A
263    where
264        F: for<'r> FnMut(A, &'r K, &'r Node<K, V>) -> ControlFlow<A, A>,
265    {
266        let mut visited = BTreeSet::new();
267        let mut result = VecDeque::new();
268        let mut skip = BTreeSet::new();
269
270        assert!(
271            roots.windows(2).all(|w| w[0] < w[1]),
272            "The roots must be sorted in ascending order"
273        );
274
275        for root in roots.iter().rev() {
276            self.visit(root, &mut visited, &mut result);
277        }
278
279        for next in result {
280            if skip.contains(&next) {
281                continue;
282            }
283            if let Some(node) = self.graph.get(&next) {
284                match filter(acc, &next, node) {
285                    ControlFlow::Continue(a) => {
286                        acc = a;
287                    }
288                    ControlFlow::Break(a) => {
289                        // When filtering out a node, we filter out all transitive dependents on
290                        // that node by adding them to the already visited list.
291                        skip.extend(self.descendants_of(node));
292
293                        acc = a;
294                    }
295                }
296            }
297        }
298        acc
299    }
300
301    /// Remove a node from the graph, and all its dependents.
302    pub fn remove(&mut self, key: &K) -> Option<Node<K, V>> {
303        if let Some(node) = self.graph.remove(key) {
304            self.tips.remove(key);
305            self.roots.remove(key);
306
307            for k in &node.dependencies {
308                if let Some(dependency) = self.graph.get_mut(k) {
309                    dependency.dependents.remove(key);
310
311                    if dependency.dependents.is_empty() {
312                        self.tips.insert(*k);
313                    }
314                }
315            }
316            for k in &node.dependents {
317                self.remove(k);
318            }
319            Some(node)
320        } else {
321            None
322        }
323    }
324
325    fn descendants_of(&self, from: &Node<K, V>) -> Vec<K> {
326        let mut visited = BTreeSet::new();
327        let mut stack = VecDeque::new();
328        let mut nodes = Vec::new();
329
330        stack.extend(from.dependents.iter());
331
332        while let Some(key) = stack.pop_front() {
333            if let Some(node) = self.graph.get(&key) {
334                if visited.insert(key) {
335                    nodes.push(key);
336
337                    for &neighbour in &node.dependents {
338                        stack.push_back(neighbour);
339                    }
340                }
341            }
342        }
343        nodes
344    }
345
346    fn ancestors_of(&self, from: &Node<K, V>) -> Vec<K> {
347        let mut visited = BTreeSet::new();
348        let mut stack = VecDeque::new();
349        let mut nodes = Vec::new();
350
351        stack.extend(from.dependencies.iter());
352
353        while let Some(key) = stack.pop_front() {
354            if let Some(node) = self.graph.get(&key) {
355                if visited.insert(key) {
356                    nodes.push(key);
357
358                    for &neighbour in &node.dependencies {
359                        stack.push_back(neighbour);
360                    }
361                }
362            }
363        }
364        nodes
365    }
366
367    /// Get the nodes that are neither an ancestor nor a descendant of the given node.
368    fn siblings_of(&self, node: &Node<K, V>) -> impl Iterator<Item = &K> {
369        let ancestors = self.ancestors_of(node);
370        let descendants = self.descendants_of(node);
371        let key = node.key;
372
373        self.graph
374            .keys()
375            .filter(move |k| !ancestors.contains(k) && !descendants.contains(k) && **k != key)
376    }
377
378    /// Add nodes recursively to the topological order, starting from the given node.
379    fn visit(&self, key: &K, visited: &mut BTreeSet<K>, order: &mut VecDeque<K>) {
380        if visited.insert(*key) {
381            // Recursively visit all of the node's dependents.
382            if let Some(node) = self.graph.get(key) {
383                for dependent in node.dependents.iter().rev() {
384                    self.visit(dependent, visited, order);
385                }
386            }
387            // Add the node to the topological order.
388            order.push_front(*key);
389        }
390    }
391
392    /// Add nodes recursively to the provided ordering, starting from the given node.
393    fn visit_by(
394        &self,
395        key: &K,
396        visited: &mut BTreeSet<K>,
397        order: &mut VecDeque<K>,
398        ordering: &impl Fn((&K, &V), (&K, &V)) -> Ordering,
399    ) {
400        if visited.insert(*key) {
401            // Recursively visit all of the node's dependents.
402            if let Some(node) = self.graph.get(key) {
403                let mut dependents: Vec<&Node<K, V>> = node
404                    .dependents
405                    .iter()
406                    .filter_map(|k| self.get(k))
407                    .collect::<Vec<_>>();
408                dependents.sort_by(|x, y| ordering((&x.key, &x.value), (&y.key, &y.value)));
409                for dependent in dependents.iter().rev() {
410                    self.visit_by(&dependent.key, visited, order, ordering);
411                }
412            }
413            // Add the node to the topological order.
414            order.push_front(*key);
415        }
416    }
417}
418
419impl<K: Ord + Copy + fmt::Display, V> Dag<K, V> {
420    /// Return the graph in "dot" format.
421    pub fn to_dot(&self) -> String {
422        let mut output = String::new();
423
424        writeln!(output, "digraph G {{").ok();
425        for (k, v) in self.graph.iter() {
426            for d in &v.dependencies {
427                writeln!(output, "\t\"{k}\" -> \"{d}\";").ok();
428            }
429        }
430        writeln!(output, "}}").ok();
431
432        output
433    }
434}
435
436impl<K: Ord + Copy + fmt::Debug, V> Index<&K> for Dag<K, V> {
437    type Output = Node<K, V>;
438
439    fn index(&self, key: &K) -> &Self::Output {
440        self.get(key)
441            .unwrap_or_else(|| panic!("Dag::index: node {key:?} not found in graph"))
442    }
443}
444
445#[cfg(test)]
446mod tests {
447    use super::*;
448
449    #[test]
450    fn test_len() {
451        let mut dag = Dag::new();
452
453        dag.node(0, ());
454        dag.node(1, ());
455        dag.node(2, ());
456
457        assert_eq!(dag.len(), 3);
458    }
459
460    #[test]
461    fn test_is_empty() {
462        let mut dag = Dag::new();
463        assert!(dag.is_empty());
464
465        dag.node(0, ());
466        assert!(!dag.is_empty());
467    }
468
469    #[test]
470    fn test_dependencies() {
471        let mut dag = Dag::new();
472
473        dag.node(0, ());
474        dag.node(1, ());
475        dag.dependency(0, 1);
476
477        assert!(dag.has_dependency(&0, &1));
478        assert!(!dag.has_dependency(&1, &0));
479    }
480
481    #[test]
482    fn test_get() {
483        let mut dag = Dag::new();
484
485        dag.node(0, "rad");
486        dag.node(1, "dar");
487
488        assert_eq!(dag[&0].value, "rad");
489        assert_eq!(dag[&1].value, "dar");
490        assert!(dag.get(&2).is_none());
491    }
492
493    #[test]
494    fn test_cycle() {
495        let mut dag = Dag::new();
496
497        dag.node(0, ());
498        dag.node(1, ());
499
500        dag.dependency(0, 1);
501        dag.dependency(1, 0);
502
503        let mut sorted = dag.sorted();
504        let expected: &[&[i32]] = &[&[0, 1], &[1, 0]];
505
506        assert!(expected.contains(&&*sorted.make_contiguous()));
507    }
508
509    #[test]
510    fn test_merge_1() {
511        let mut a = Dag::new();
512        let mut b = Dag::new();
513        let mut c = Dag::new();
514
515        a.node(0, ());
516        a.node(1, ());
517        a.dependency(1, 0);
518
519        b.node(0, ());
520        b.node(2, ());
521        b.dependency(2, 0);
522
523        c.merge(a);
524        c.merge(b);
525
526        assert!(c.get(&0).is_some());
527        assert!(c.get(&1).is_some());
528        assert!(c.get(&2).is_some());
529        assert!(c.has_dependency(&1, &0));
530        assert!(c.has_dependency(&2, &0));
531    }
532
533    #[test]
534    fn test_merge_2() {
535        let mut a = Dag::new();
536        let mut b = Dag::new();
537
538        a.node(0, ());
539        a.node(1, ());
540        a.node(2, ());
541        a.dependency(1, 0);
542        a.dependency(2, 0);
543
544        b.node(0, ());
545        b.node(1, ());
546        b.node(2, ());
547        b.node(3, ());
548        b.node(4, ());
549        b.dependency(1, 0);
550        b.dependency(2, 0);
551        b.dependency(3, 0);
552        b.dependency(4, 2);
553
554        assert!(a.tips.contains(&2));
555
556        a.merge(b);
557
558        assert!(a.get(&0).is_some());
559        assert!(a.get(&1).is_some());
560        assert!(a.get(&2).is_some());
561        assert!(a.get(&3).is_some());
562        assert!(a.get(&4).is_some());
563        assert!(a.has_dependency(&4, &2));
564        assert!(a.get(&2).unwrap().dependents.contains(&4));
565        assert!(a.get(&0).unwrap().dependents.contains(&3));
566        assert!(a.tips.contains(&1));
567        assert!(!a.tips.contains(&2));
568        assert!(a.tips.contains(&3));
569        assert!(a.tips.contains(&4));
570        assert!(a.roots.contains(&0));
571    }
572
573    #[test]
574    fn test_diamond() {
575        let mut dag = Dag::new();
576
577        dag.node(0, ());
578        dag.node(1, ());
579        dag.node(2, ());
580        dag.node(3, ());
581
582        dag.dependency(1, 0);
583        dag.dependency(2, 0);
584        dag.dependency(3, 1);
585        dag.dependency(3, 2);
586
587        assert_eq!(dag.tips().map(|(k, _)| *k).collect::<Vec<_>>(), vec![3]);
588        assert_eq!(dag.roots().map(|(k, _)| *k).collect::<Vec<_>>(), vec![0]);
589
590        // All of the possible sort orders for the above graph.
591        let expected: &[&[i32]] = &[&[0, 1, 2, 3], &[0, 2, 1, 3]];
592        let mut actual = dag.sorted();
593
594        assert!(expected.contains(&&*actual.make_contiguous()), "{actual:?}");
595    }
596
597    #[test]
598    fn test_complex() {
599        let mut dag = Dag::new();
600
601        dag.node(0, ());
602        dag.node(1, ());
603        dag.node(2, ());
604        dag.node(3, ());
605        dag.node(4, ());
606        dag.node(5, ());
607
608        dag.dependency(3, 2);
609        dag.dependency(1, 3);
610        dag.dependency(2, 5);
611        dag.dependency(0, 5);
612        dag.dependency(0, 4);
613        dag.dependency(1, 4);
614
615        assert_eq!(
616            dag.tips().map(|(k, _)| *k).collect::<BTreeSet<_>>(),
617            BTreeSet::from_iter([1, 0])
618        );
619        assert_eq!(
620            dag.roots().map(|(k, _)| *k).collect::<BTreeSet<_>>(),
621            BTreeSet::from_iter([4, 5])
622        );
623
624        // All of the possible sort orders for the above graph.
625        let expected = &[
626            [4, 5, 0, 2, 3, 1],
627            [4, 5, 2, 0, 3, 1],
628            [4, 5, 2, 3, 0, 1],
629            [4, 5, 2, 3, 1, 0],
630            [5, 2, 3, 4, 0, 1],
631            [5, 2, 3, 4, 1, 0],
632            [5, 2, 4, 0, 3, 1],
633            [5, 2, 4, 3, 0, 1],
634            [5, 2, 4, 3, 1, 0],
635            [5, 4, 0, 2, 3, 1],
636            [5, 4, 2, 0, 3, 1],
637            [5, 4, 2, 3, 0, 1],
638            [5, 4, 2, 3, 1, 0],
639        ];
640        let mut sorts = BTreeSet::new();
641        let mut rng = fastrand::Rng::new();
642
643        while sorts.len() < expected.len() {
644            sorts.insert(
645                dag.sorted_by(|a, b| if rng.bool() { a.cmp(b) } else { b.cmp(a) })
646                    .make_contiguous()
647                    .to_vec(),
648            );
649        }
650        for e in expected {
651            assert!(sorts.remove(e.to_vec().as_slice()));
652        }
653        assert!(sorts.is_empty());
654    }
655
656    #[test]
657    fn test_fold_sorting_1() {
658        let mut dag = Dag::new();
659
660        dag.node("R", ());
661        dag.node("A1", ());
662        dag.node("A2", ());
663        dag.node("A3", ());
664        dag.node("B1", ());
665        dag.node("B2", ());
666        dag.node("B3", ());
667        dag.node("C1", ());
668
669        dag.dependency("A1", "R");
670        dag.dependency("A2", "R");
671        dag.dependency("A3", "R");
672
673        dag.dependency("B1", "A1");
674        dag.dependency("B2", "A1");
675        dag.dependency("B3", "A2");
676        dag.dependency("B3", "A3");
677
678        dag.dependency("C1", "B1");
679        dag.dependency("C1", "B2");
680        dag.dependency("C1", "B3");
681
682        let acc = dag.fold(&["R"], Vec::new(), |mut acc, key, _| {
683            acc.push(*key);
684            ControlFlow::Continue(acc)
685        });
686        assert_eq!(acc, vec!["R", "A1", "B1", "B2", "A2", "A3", "B3", "C1"]);
687    }
688
689    #[test]
690    fn test_fold_sorting_2() {
691        let mut dag = Dag::new();
692
693        dag.node("R", ());
694        dag.node("A1", ());
695        dag.node("A2", ());
696        dag.node("A3", ());
697        dag.node("B1", ());
698        dag.node("C1", ());
699        dag.node("C2", ());
700        dag.node("C3", ());
701
702        dag.dependency("A1", "R");
703        dag.dependency("A2", "A1");
704        dag.dependency("A3", "A2");
705
706        dag.dependency("B1", "R");
707
708        dag.dependency("C1", "B1");
709        dag.dependency("C1", "A3");
710        dag.dependency("C2", "B1");
711        dag.dependency("C2", "A3");
712        dag.dependency("C3", "B1");
713        dag.dependency("C3", "A3");
714
715        let acc = dag.fold(&["R"], Vec::new(), |mut acc, key, _| {
716            acc.push(*key);
717            ControlFlow::Continue(acc)
718        });
719
720        assert_eq!(acc, vec!["R", "A1", "A2", "A3", "B1", "C1", "C2", "C3"]);
721        assert_eq!(dag.sorted(), acc);
722    }
723
724    #[test]
725    fn test_fold_diamond() {
726        let mut dag = Dag::new();
727
728        dag.node("R", ());
729        dag.node("A1", ());
730        dag.node("A2", ());
731        dag.node("B", ());
732
733        dag.dependency("A1", "R");
734        dag.dependency("A2", "R");
735        dag.dependency("B", "A1");
736        dag.dependency("B", "A2");
737
738        let acc = dag.fold(&["R"], Vec::new(), |mut acc, key, _| {
739            acc.push(*key);
740            ControlFlow::Continue(acc)
741        });
742        assert_eq!(acc, vec!["R", "A1", "A2", "B"]);
743
744        let sorted = dag.sorted();
745        assert_eq!(sorted, acc);
746    }
747
748    #[test]
749    fn test_fold_multiple_roots() {
750        let mut dag = Dag::new();
751
752        dag.node("R", ());
753        dag.node("A1", ());
754        dag.node("A2", ());
755
756        dag.dependency("A1", "R");
757        dag.dependency("A2", "R");
758
759        let acc = dag.fold(&["A1", "A2"], Vec::new(), |mut acc, key, _| {
760            acc.push(*key);
761            ControlFlow::Continue(acc)
762        });
763        assert_eq!(acc, &["A1", "A2"]);
764    }
765
766    #[test]
767    fn test_fold_reject() {
768        let mut dag = Dag::new();
769
770        dag.node("R", ());
771        dag.node("A1", ());
772        dag.node("A2", ());
773        dag.node("B1", ());
774        dag.node("C1", ());
775        dag.node("D1", ());
776
777        dag.dependency("A1", "R");
778        dag.dependency("A2", "R");
779        dag.dependency("B1", "A1");
780        dag.dependency("C1", "B1");
781        dag.dependency("D1", "C1");
782        dag.dependency("D1", "A2");
783
784        let a1 = dag.get(&"A1").unwrap();
785        assert_eq!(dag.descendants_of(a1), vec!["B1", "C1", "D1"]);
786
787        let acc = dag.fold(&["R"], Vec::new(), |mut acc, key, _| {
788            if *key == "A1" {
789                ControlFlow::Break(acc)
790            } else {
791                acc.push(*key);
792                ControlFlow::Continue(acc)
793            }
794        });
795        assert_eq!(acc, vec!["R", "A2"]);
796
797        let acc = dag.fold(&["R"], Vec::new(), |mut acc, key, _| {
798            if *key == "A2" {
799                ControlFlow::Break(acc)
800            } else {
801                acc.push(*key);
802                ControlFlow::Continue(acc)
803            }
804        });
805        assert_eq!(acc, vec!["R", "A1", "B1", "C1"]);
806    }
807
808    #[test]
809    fn test_remove() {
810        let mut dag = Dag::new();
811
812        dag.node("R", ());
813        dag.node("A1", ());
814        dag.node("A2", ());
815        dag.node("A3", ());
816        dag.node("B1", ());
817        dag.node("C1", ());
818        dag.node("D1", ());
819
820        dag.dependency("A1", "R");
821        dag.dependency("A2", "R");
822        dag.dependency("A3", "A2");
823        dag.dependency("B1", "A1");
824        dag.dependency("B1", "A2");
825        dag.dependency("C1", "B1");
826        dag.dependency("C1", "A3");
827        dag.dependency("D1", "C1");
828        dag.dependency("D1", "A2");
829
830        dag.remove(&"C1");
831        assert!(dag.get(&"C1").is_none());
832        assert!(dag.get(&"D1").is_none());
833        assert!(!dag.tips.contains(&"D1"));
834        assert_eq!(dag.tips.iter().collect::<Vec<_>>(), vec![&"A3", &"B1"]);
835
836        dag.remove(&"A3");
837        assert_eq!(dag.tips.iter().collect::<Vec<_>>(), vec![&"B1"]);
838
839        dag.remove(&"A1");
840        assert!(dag.get(&"A1").is_none());
841        assert!(dag.get(&"B1").is_none());
842        assert!(dag.get(&"A2").is_some());
843        assert_eq!(dag.tips.iter().collect::<Vec<_>>(), vec![&"A2"]);
844
845        dag.remove(&"R");
846        assert!(dag.is_empty());
847        assert!(dag.tips.is_empty());
848        assert!(dag.roots.is_empty());
849    }
850
851    #[test]
852    fn test_prune_1() {
853        let mut dag = Dag::new();
854
855        dag.node("R", ());
856        dag.node("A1", ());
857        dag.node("A2", ());
858        dag.node("B1", ());
859        dag.node("C1", ());
860        dag.node("D1", ());
861
862        dag.dependency("A1", "R");
863        dag.dependency("A2", "R");
864        dag.dependency("B1", "A1");
865        dag.dependency("C1", "B1");
866        dag.dependency("D1", "C1");
867        dag.dependency("D1", "A2");
868
869        let a1 = dag.get(&"A1").unwrap();
870        assert_eq!(dag.descendants_of(a1), vec!["B1", "C1", "D1"]);
871
872        dag.prune(&["R"], |key, _, _| {
873            if key == &"B1" {
874                ControlFlow::Break(())
875            } else {
876                ControlFlow::Continue(())
877            }
878        });
879        assert_eq!(dag.sorted(), vec!["R", "A1", "A2"]);
880    }
881
882    #[test]
883    fn test_siblings() {
884        let mut dag = Dag::new();
885
886        dag.node("R", ());
887        dag.node("A1", ());
888        dag.node("A2", ());
889        dag.node("A3", ());
890        dag.node("A4", ());
891        dag.node("B1", ());
892        dag.node("C1", ());
893        dag.node("C2", ());
894        dag.node("C3", ());
895
896        dag.dependency("A1", "R");
897        dag.dependency("A2", "A1");
898        dag.dependency("A3", "A2");
899
900        dag.dependency("B1", "A2");
901
902        dag.dependency("C1", "R");
903        dag.dependency("C2", "C1");
904        dag.dependency("C3", "C2");
905
906        dag.dependency("A4", "B1");
907        dag.dependency("A4", "C3");
908        dag.dependency("A4", "A3");
909
910        let siblings: Vec<_> = dag.siblings_of(dag.get(&"A3").unwrap()).copied().collect();
911        assert_eq!(siblings, vec!["B1", "C1", "C2", "C3"]);
912
913        let siblings: Vec<_> = dag.siblings_of(dag.get(&"A4").unwrap()).copied().collect();
914        assert_eq!(siblings, Vec::<&str>::new());
915
916        let siblings: Vec<_> = dag.siblings_of(dag.get(&"C1").unwrap()).copied().collect();
917        assert_eq!(siblings, vec!["A1", "A2", "A3", "B1"]);
918
919        let siblings: Vec<_> = dag.siblings_of(dag.get(&"C2").unwrap()).copied().collect();
920        assert_eq!(siblings, vec!["A1", "A2", "A3", "B1"]);
921
922        let siblings: Vec<_> = dag.siblings_of(dag.get(&"B1").unwrap()).copied().collect();
923        assert_eq!(siblings, vec!["A3", "C1", "C2", "C3"]);
924
925        let siblings: Vec<_> = dag.siblings_of(dag.get(&"R").unwrap()).copied().collect();
926        assert_eq!(siblings, Vec::<&str>::new());
927    }
928
929    #[test]
930    fn test_prune_2() {
931        let mut dag = Dag::new();
932
933        dag.node("R", ());
934        dag.node("A1", ());
935        dag.node("A2", ());
936        dag.node("A3", ());
937        dag.node("B1", ());
938        dag.node("C1", ());
939        dag.node("C2", ());
940        dag.node("C3", ());
941
942        dag.dependency("A1", "R");
943        dag.dependency("A2", "A1");
944        dag.dependency("A3", "A2");
945
946        dag.dependency("B1", "R");
947
948        dag.dependency("C1", "B1");
949        dag.dependency("C1", "A3");
950        dag.dependency("C2", "B1");
951        dag.dependency("C2", "A3");
952        dag.dependency("C3", "B1");
953        dag.dependency("C3", "A3");
954
955        let mut order = VecDeque::new();
956
957        dag.prune(&["R"], |key, _, _| {
958            order.push_back(*key);
959            ControlFlow::Continue(())
960        });
961        assert_eq!(order, dag.sorted());
962    }
963
964    #[test]
965    fn test_prune_by_sorting() {
966        let mut dag = Dag::new();
967
968        dag.node("R", 0);
969        dag.node("A1", 1);
970        dag.node("A2", 2);
971        dag.node("A3", 3);
972        dag.node("B1", 1);
973        dag.node("B2", 2);
974        dag.node("B3", 3);
975        dag.node("C1", 1);
976
977        dag.dependency("A1", "R");
978        dag.dependency("A2", "R");
979        dag.dependency("A3", "R");
980
981        dag.dependency("B1", "A1");
982        dag.dependency("B2", "A1");
983        dag.dependency("B3", "A2");
984        dag.dependency("B3", "A3");
985
986        dag.dependency("C1", "B1");
987        dag.dependency("C1", "B2");
988        dag.dependency("C1", "B3");
989
990        let mut order = Vec::new();
991        dag.prune_by(
992            &["R"],
993            |key, _, _| {
994                order.push(*key);
995                ControlFlow::Continue(())
996            },
997            |(a, _), (b, _)| a.cmp(b),
998        );
999        assert_eq!(order, vec!["R", "A1", "B1", "B2", "A2", "A3", "B3", "C1"]);
1000
1001        let mut order = Vec::new();
1002        dag.prune_by(
1003            &["R"],
1004            |key, _, _| {
1005                order.push(*key);
1006                ControlFlow::Continue(())
1007            },
1008            |(a, _), (b, _)| a.cmp(b).reverse(),
1009        );
1010        assert_eq!(order, vec!["R", "A3", "A2", "B3", "A1", "B2", "B1", "C1"]);
1011
1012        let mut order = Vec::new();
1013        dag.prune_by(
1014            &["R"],
1015            |key, _, _| {
1016                order.push(*key);
1017                ControlFlow::Continue(())
1018            },
1019            |(_, a), (_, b)| a.cmp(b).reverse(),
1020        );
1021        assert_eq!(order, vec!["R", "A3", "A2", "B3", "A1", "B2", "B1", "C1"]);
1022    }
1023
1024    #[test]
1025    fn test_contains() {
1026        let mut dag = Dag::<u8, ()>::new();
1027
1028        assert!(!dag.contains(&0));
1029
1030        dag.node(0, ());
1031        dag.node(1, ());
1032        dag.dependency(0, 1);
1033        dag.node(2, ());
1034        dag.dependency(2, 1);
1035        dag.dependency(2, 0);
1036        dag.node(3, ());
1037
1038        for i in 0..4 {
1039            assert!(dag.contains(&i));
1040        }
1041        assert!(!dag.contains(&4));
1042    }
1043}