jj_lib/
dag_walk.rs

1// Copyright 2020 The Jujutsu Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// https://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! General-purpose DAG algorithms.
16
17use std::collections::BinaryHeap;
18use std::collections::HashMap;
19use std::collections::HashSet;
20use std::convert::Infallible;
21use std::hash::Hash;
22use std::iter;
23use std::mem;
24
25use itertools::Itertools as _;
26use smallvec::SmallVec;
27use smallvec::smallvec_inline;
28
29/// Traverses nodes from `start` in depth-first order.
30pub fn dfs<T, ID, II, NI>(
31    start: II,
32    id_fn: impl Fn(&T) -> ID,
33    mut neighbors_fn: impl FnMut(&T) -> NI,
34) -> impl Iterator<Item = T>
35where
36    ID: Hash + Eq,
37    II: IntoIterator<Item = T>,
38    NI: IntoIterator<Item = T>,
39{
40    let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node));
41    dfs_ok(to_ok_iter(start), id_fn, neighbors_fn).map(Result::unwrap)
42}
43
44/// Traverses nodes from `start` in depth-first order.
45///
46/// An `Err` is emitted as a node with no neighbors. Caller may decide to
47/// short-circuit on it.
48pub fn dfs_ok<T, ID, E, II, NI>(
49    start: II,
50    id_fn: impl Fn(&T) -> ID,
51    mut neighbors_fn: impl FnMut(&T) -> NI,
52) -> impl Iterator<Item = Result<T, E>>
53where
54    ID: Hash + Eq,
55    II: IntoIterator<Item = Result<T, E>>,
56    NI: IntoIterator<Item = Result<T, E>>,
57{
58    let mut work: Vec<Result<T, E>> = start.into_iter().collect();
59    let mut visited: HashSet<ID> = HashSet::new();
60    iter::from_fn(move || {
61        loop {
62            let c = match work.pop() {
63                Some(Ok(c)) => c,
64                r @ (Some(Err(_)) | None) => return r,
65            };
66            let id = id_fn(&c);
67            if visited.contains(&id) {
68                continue;
69            }
70            for p in neighbors_fn(&c) {
71                work.push(p);
72            }
73            visited.insert(id);
74            return Some(Ok(c));
75        }
76    })
77}
78
79/// Builds a list of nodes reachable from the `start` where neighbors come
80/// before the node itself.
81///
82/// Panics if the graph has cycle.
83pub fn topo_order_forward<T, ID, II, NI>(
84    start: II,
85    id_fn: impl Fn(&T) -> ID,
86    mut neighbors_fn: impl FnMut(&T) -> NI,
87) -> Vec<T>
88where
89    ID: Hash + Eq + Clone,
90    II: IntoIterator<Item = T>,
91    NI: IntoIterator<Item = T>,
92{
93    let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node));
94    let cycle_fn = |_| panic!("graph has cycle");
95    topo_order_forward_ok(to_ok_iter(start), id_fn, neighbors_fn, cycle_fn).unwrap()
96}
97
98/// Builds a list of `Ok` nodes reachable from the `start` where neighbors come
99/// before the node itself.
100///
101/// If `start` or `neighbors_fn()` yields an `Err`, this function terminates and
102/// returns the error. If the graph has cycle, `cycle_fn()` is called with one
103/// of the nodes involved in the cycle.
104pub fn topo_order_forward_ok<T, ID, E, II, NI>(
105    start: II,
106    id_fn: impl Fn(&T) -> ID,
107    mut neighbors_fn: impl FnMut(&T) -> NI,
108    cycle_fn: impl FnOnce(T) -> E,
109) -> Result<Vec<T>, E>
110where
111    ID: Hash + Eq + Clone,
112    II: IntoIterator<Item = Result<T, E>>,
113    NI: IntoIterator<Item = Result<T, E>>,
114{
115    let mut stack: Vec<(T, bool)> = start.into_iter().map(|r| Ok((r?, false))).try_collect()?;
116    let mut visiting = HashSet::new();
117    let mut emitted = HashSet::new();
118    let mut result = vec![];
119    while let Some((node, neighbors_visited)) = stack.pop() {
120        let id = id_fn(&node);
121        if emitted.contains(&id) {
122            continue;
123        }
124        if !neighbors_visited {
125            if !visiting.insert(id.clone()) {
126                return Err(cycle_fn(node));
127            }
128            let neighbors_iter = neighbors_fn(&node).into_iter();
129            stack.reserve(neighbors_iter.size_hint().0 + 1);
130            stack.push((node, true));
131            for neighbor in neighbors_iter {
132                stack.push((neighbor?, false));
133            }
134        } else {
135            visiting.remove(&id);
136            emitted.insert(id);
137            result.push(node);
138        }
139    }
140    Ok(result)
141}
142
143/// Builds a list of nodes reachable from the `start` where neighbors come after
144/// the node itself.
145///
146/// Panics if the graph has cycle.
147pub fn topo_order_reverse<T, ID, II, NI>(
148    start: II,
149    id_fn: impl Fn(&T) -> ID,
150    mut neighbors_fn: impl FnMut(&T) -> NI,
151) -> Vec<T>
152where
153    ID: Hash + Eq + Clone,
154    II: IntoIterator<Item = T>,
155    NI: IntoIterator<Item = T>,
156{
157    let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node));
158    let cycle_fn = |_| panic!("graph has cycle");
159    topo_order_reverse_ok(to_ok_iter(start), id_fn, neighbors_fn, cycle_fn).unwrap()
160}
161
162/// Builds a list of `Ok` nodes reachable from the `start` where neighbors come
163/// after the node itself.
164///
165/// If `start` or `neighbors_fn()` yields an `Err`, this function terminates and
166/// returns the error. If the graph has cycle, `cycle_fn()` is called with one
167/// of the nodes involved in the cycle.
168pub fn topo_order_reverse_ok<T, ID, E, II, NI>(
169    start: II,
170    id_fn: impl Fn(&T) -> ID,
171    neighbors_fn: impl FnMut(&T) -> NI,
172    cycle_fn: impl FnOnce(T) -> E,
173) -> Result<Vec<T>, E>
174where
175    ID: Hash + Eq + Clone,
176    II: IntoIterator<Item = Result<T, E>>,
177    NI: IntoIterator<Item = Result<T, E>>,
178{
179    let mut result = topo_order_forward_ok(start, id_fn, neighbors_fn, cycle_fn)?;
180    result.reverse();
181    Ok(result)
182}
183
184/// Like `topo_order_reverse()`, but can iterate linear DAG lazily.
185///
186/// The DAG is supposed to be (mostly) topologically ordered by `T: Ord`.
187/// For example, topological order of chronological data should respect
188/// timestamp (except a few outliers caused by clock skew.)
189///
190/// Use `topo_order_reverse()` if the DAG is heavily branched. This can
191/// only process linear part lazily.
192///
193/// Panics if the graph has cycle.
194pub fn topo_order_reverse_lazy<T, ID, II, NI>(
195    start: II,
196    id_fn: impl Fn(&T) -> ID,
197    mut neighbors_fn: impl FnMut(&T) -> NI,
198) -> impl Iterator<Item = T>
199where
200    T: Ord,
201    ID: Hash + Eq + Clone,
202    II: IntoIterator<Item = T>,
203    NI: IntoIterator<Item = T>,
204{
205    let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node));
206    let cycle_fn = |_| panic!("graph has cycle");
207    topo_order_reverse_lazy_ok(to_ok_iter(start), id_fn, neighbors_fn, cycle_fn).map(Result::unwrap)
208}
209
210/// Like `topo_order_reverse_ok()`, but can iterate linear DAG lazily.
211///
212/// The returned iterator short-circuits at an `Err`. Pending non-linear nodes
213/// before the `Err` will be discarded. If the graph has cycle, `cycle_fn()` is
214/// called with one of the nodes involved in the cycle.
215pub fn topo_order_reverse_lazy_ok<T, ID, E, II, NI>(
216    start: II,
217    id_fn: impl Fn(&T) -> ID,
218    mut neighbors_fn: impl FnMut(&T) -> NI,
219    mut cycle_fn: impl FnMut(T) -> E,
220) -> impl Iterator<Item = Result<T, E>>
221where
222    T: Ord,
223    ID: Hash + Eq + Clone,
224    II: IntoIterator<Item = Result<T, E>>,
225    NI: IntoIterator<Item = Result<T, E>>,
226{
227    let mut inner = TopoOrderReverseLazyInner::empty();
228    inner.extend(start);
229    iter::from_fn(move || inner.next(&id_fn, &mut neighbors_fn, &mut cycle_fn))
230}
231
232#[derive(Clone, Debug)]
233struct TopoOrderReverseLazyInner<T, ID, E> {
234    start: Vec<T>,
235    result: Vec<Result<T, E>>,
236    emitted: HashSet<ID>,
237}
238
239impl<T: Ord, ID: Hash + Eq + Clone, E> TopoOrderReverseLazyInner<T, ID, E> {
240    fn empty() -> Self {
241        Self {
242            start: Vec::new(),
243            result: Vec::new(),
244            emitted: HashSet::new(),
245        }
246    }
247
248    fn extend(&mut self, iter: impl IntoIterator<Item = Result<T, E>>) {
249        let iter = iter.into_iter();
250        self.start.reserve(iter.size_hint().0);
251        for res in iter {
252            if let Ok(node) = res {
253                self.start.push(node);
254            } else {
255                // Emit the error and terminate
256                self.start.clear();
257                self.result.insert(0, res);
258                return;
259            }
260        }
261    }
262
263    fn next<NI: IntoIterator<Item = Result<T, E>>>(
264        &mut self,
265        id_fn: impl Fn(&T) -> ID,
266        mut neighbors_fn: impl FnMut(&T) -> NI,
267        mut cycle_fn: impl FnMut(T) -> E,
268    ) -> Option<Result<T, E>> {
269        if let Some(res) = self.result.pop() {
270            return Some(res);
271        }
272
273        // Fast path for linear DAG
274        if self.start.len() <= 1 {
275            let node = self.start.pop()?;
276            self.extend(neighbors_fn(&node));
277            if self.emitted.insert(id_fn(&node)) {
278                return Some(Ok(node));
279            } else {
280                return Some(Err(cycle_fn(node)));
281            }
282        }
283
284        // Extract graph nodes based on T's order, and sort them by using ids
285        // (because we wouldn't want to clone T itself)
286        let start_ids = self.start.iter().map(&id_fn).collect_vec();
287        match look_ahead_sub_graph(mem::take(&mut self.start), &id_fn, &mut neighbors_fn) {
288            Ok((mut node_map, neighbor_ids_map, remainder)) => {
289                self.start = remainder;
290                let sorted_ids = match topo_order_forward_ok(
291                    start_ids.iter().map(Ok),
292                    |id| *id,
293                    |id| neighbor_ids_map[id].iter().map(Ok),
294                    |id| cycle_fn(node_map.remove(id).unwrap()),
295                ) {
296                    Ok(ids) => ids,
297                    Err(err) => return Some(Err(err)),
298                };
299                self.result.reserve(sorted_ids.len());
300                for id in sorted_ids {
301                    let (id, node) = node_map.remove_entry(id).unwrap();
302                    if self.emitted.insert(id) {
303                        self.result.push(Ok(node));
304                    } else {
305                        self.result.push(Err(cycle_fn(node)));
306                    }
307                }
308                self.result.pop()
309            }
310            Err(err) => Some(Err(err)),
311        }
312    }
313}
314
315/// Splits DAG at the first single fork point, and builds a list of nodes
316/// reachable from the `start` where neighbors come after the node itself.
317///
318/// This is a building block for lazy DAG iterators similar to
319/// [`topo_order_reverse_lazy()`]. The `start` list will be updated to include
320/// the next nodes to visit.
321///
322/// If the split chunk of the graph has cycle, `cycle_fn()` is called with one
323/// of the nodes involved in the cycle.
324pub fn topo_order_reverse_chunked<T, ID, E, NI>(
325    start: &mut Vec<T>,
326    id_fn: impl Fn(&T) -> ID,
327    mut neighbors_fn: impl FnMut(&T) -> NI,
328    mut cycle_fn: impl FnMut(T) -> E,
329) -> Result<SmallVec<[T; 1]>, E>
330where
331    T: Ord,
332    ID: Hash + Eq + Clone,
333    NI: IntoIterator<Item = Result<T, E>>,
334{
335    // Fast path for linear DAG
336    if start.len() <= 1 {
337        let Some(node) = start.pop() else {
338            return Ok(SmallVec::new());
339        };
340        let neighbors_iter = neighbors_fn(&node).into_iter();
341        start.reserve(neighbors_iter.size_hint().0);
342        for neighbor in neighbors_iter {
343            start.push(neighbor?);
344        }
345        return Ok(smallvec_inline![node]);
346    }
347
348    // Extract graph nodes based on T's order, and sort them by using ids
349    // (because we wouldn't want to clone T itself)
350    let start_ids = start.iter().map(&id_fn).collect_vec();
351    let (mut node_map, neighbor_ids_map, remainder) =
352        look_ahead_sub_graph(mem::take(start), &id_fn, &mut neighbors_fn)?;
353    *start = remainder;
354    let sorted_ids = topo_order_forward_ok(
355        start_ids.iter().map(Ok),
356        |id| *id,
357        |id| neighbor_ids_map[id].iter().map(Ok),
358        |id| cycle_fn(node_map.remove(id).unwrap()),
359    )?;
360    let sorted_nodes = sorted_ids
361        .iter()
362        .rev()
363        .map(|&id| node_map.remove(id).unwrap())
364        .collect();
365    Ok(sorted_nodes)
366}
367
368/// Splits DAG at single fork point, and extracts branchy part as sub graph.
369///
370/// ```text
371///  o | C
372///  | o B
373///  |/ <---- split here (A->B or A->C would create cycle)
374///  o A
375/// ```
376///
377/// If a branch reached to root (empty neighbors), the graph can't be split
378/// anymore because the other branch may be connected to a descendant of
379/// the rooted branch.
380///
381/// ```text
382///  o | C
383///  | o B
384///  |  <---- can't split here (there may be edge A->B)
385///  o A
386/// ```
387///
388/// We assume the graph is (mostly) topologically ordered by `T: Ord`.
389#[expect(clippy::type_complexity)]
390fn look_ahead_sub_graph<T, ID, E, NI>(
391    start: Vec<T>,
392    id_fn: impl Fn(&T) -> ID,
393    mut neighbors_fn: impl FnMut(&T) -> NI,
394) -> Result<(HashMap<ID, T>, HashMap<ID, Vec<ID>>, Vec<T>), E>
395where
396    T: Ord,
397    ID: Hash + Eq + Clone,
398    NI: IntoIterator<Item = Result<T, E>>,
399{
400    let mut queue: BinaryHeap<T> = start.into();
401    // Build separate node/neighbors maps since lifetime is different at caller
402    let mut node_map: HashMap<ID, T> = HashMap::new();
403    let mut neighbor_ids_map: HashMap<ID, Vec<ID>> = HashMap::new();
404    let mut has_reached_root = false;
405    while queue.len() > 1 || node_map.is_empty() || has_reached_root {
406        let Some(node) = queue.pop() else {
407            break;
408        };
409        let node_id = id_fn(&node);
410        if node_map.contains_key(&node_id) {
411            continue;
412        }
413
414        let mut neighbor_ids = Vec::new();
415        let mut neighbors_iter = neighbors_fn(&node).into_iter().peekable();
416        has_reached_root |= neighbors_iter.peek().is_none();
417        for neighbor in neighbors_iter {
418            let neighbor = neighbor?;
419            neighbor_ids.push(id_fn(&neighbor));
420            queue.push(neighbor);
421        }
422        node_map.insert(node_id.clone(), node);
423        neighbor_ids_map.insert(node_id, neighbor_ids);
424    }
425
426    assert!(queue.len() <= 1, "order of remainder shouldn't matter");
427    let remainder = queue.into_vec();
428
429    // Omit unvisited neighbors
430    if let Some(unvisited_id) = remainder.first().map(&id_fn) {
431        for neighbor_ids in neighbor_ids_map.values_mut() {
432            neighbor_ids.retain(|id| *id != unvisited_id);
433        }
434    }
435
436    Ok((node_map, neighbor_ids_map, remainder))
437}
438
439/// Builds a list of nodes reachable from the `start` where neighbors come after
440/// the node itself.
441///
442/// Unlike `topo_order_reverse()`, nodes are sorted in reverse `T: Ord` order so
443/// long as they can respect the topological requirement.
444///
445/// Panics if the graph has cycle.
446pub fn topo_order_reverse_ord<T, ID, II, NI>(
447    start: II,
448    id_fn: impl Fn(&T) -> ID,
449    mut neighbors_fn: impl FnMut(&T) -> NI,
450) -> Vec<T>
451where
452    T: Ord,
453    ID: Hash + Eq + Clone,
454    II: IntoIterator<Item = T>,
455    NI: IntoIterator<Item = T>,
456{
457    let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node));
458    let cycle_fn = |_| panic!("graph has cycle");
459    topo_order_reverse_ord_ok(to_ok_iter(start), id_fn, neighbors_fn, cycle_fn).unwrap()
460}
461
462/// Builds a list of `Ok` nodes reachable from the `start` where neighbors come
463/// after the node itself.
464///
465/// Unlike `topo_order_reverse_ok()`, nodes are sorted in reverse `T: Ord` order
466/// so long as they can respect the topological requirement.
467///
468/// If `start` or `neighbors_fn()` yields an `Err`, this function terminates and
469/// returns the error. If the graph has cycle, `cycle_fn()` is called with one
470/// of the nodes involved in the cycle.
471pub fn topo_order_reverse_ord_ok<T, ID, E, II, NI>(
472    start: II,
473    id_fn: impl Fn(&T) -> ID,
474    mut neighbors_fn: impl FnMut(&T) -> NI,
475    cycle_fn: impl FnOnce(T) -> E,
476) -> Result<Vec<T>, E>
477where
478    T: Ord,
479    ID: Hash + Eq + Clone,
480    II: IntoIterator<Item = Result<T, E>>,
481    NI: IntoIterator<Item = Result<T, E>>,
482{
483    struct InnerNode<T> {
484        node: Option<T>,
485        indegree: usize,
486    }
487
488    // DFS to accumulate incoming edges
489    let mut stack: Vec<T> = start.into_iter().try_collect()?;
490    let mut head_node_map: HashMap<ID, T> = HashMap::new();
491    let mut inner_node_map: HashMap<ID, InnerNode<T>> = HashMap::new();
492    let mut neighbor_ids_map: HashMap<ID, Vec<ID>> = HashMap::new();
493    while let Some(node) = stack.pop() {
494        let node_id = id_fn(&node);
495        if neighbor_ids_map.contains_key(&node_id) {
496            continue; // Already visited
497        }
498
499        let neighbors_iter = neighbors_fn(&node).into_iter();
500        let pos = stack.len();
501        stack.reserve(neighbors_iter.size_hint().0);
502        for neighbor in neighbors_iter {
503            stack.push(neighbor?);
504        }
505        let neighbor_ids = stack[pos..].iter().map(&id_fn).collect_vec();
506        if let Some(inner) = inner_node_map.get_mut(&node_id) {
507            inner.node = Some(node);
508        } else {
509            head_node_map.insert(node_id.clone(), node);
510        }
511
512        for neighbor_id in &neighbor_ids {
513            if let Some(inner) = inner_node_map.get_mut(neighbor_id) {
514                inner.indegree += 1;
515            } else {
516                let inner = InnerNode {
517                    node: head_node_map.remove(neighbor_id),
518                    indegree: 1,
519                };
520                inner_node_map.insert(neighbor_id.clone(), inner);
521            }
522        }
523        neighbor_ids_map.insert(node_id, neighbor_ids);
524    }
525
526    debug_assert!(
527        head_node_map
528            .keys()
529            .all(|id| !inner_node_map.contains_key(id))
530    );
531    debug_assert!(inner_node_map.values().all(|inner| inner.node.is_some()));
532    debug_assert!(inner_node_map.values().all(|inner| inner.indegree > 0));
533
534    // Using Kahn's algorithm
535    let mut queue: BinaryHeap<T> = head_node_map.into_values().collect();
536    let mut result = Vec::new();
537    while let Some(node) = queue.pop() {
538        let node_id = id_fn(&node);
539        result.push(node);
540        for neighbor_id in neighbor_ids_map.remove(&node_id).unwrap() {
541            let inner = inner_node_map.get_mut(&neighbor_id).unwrap();
542            inner.indegree -= 1;
543            if inner.indegree == 0 {
544                queue.push(inner.node.take().unwrap());
545                inner_node_map.remove(&neighbor_id);
546            }
547        }
548    }
549
550    if let Some(inner) = inner_node_map.into_values().next() {
551        Err(cycle_fn(inner.node.unwrap()))
552    } else {
553        Ok(result)
554    }
555}
556
557/// Find nodes in the start set that are not reachable from other nodes in the
558/// start set.
559pub fn heads<T, ID, II, NI>(
560    start: II,
561    id_fn: impl Fn(&T) -> ID,
562    mut neighbors_fn: impl FnMut(&T) -> NI,
563) -> HashSet<T>
564where
565    T: Hash + Eq + Clone,
566    ID: Hash + Eq,
567    II: IntoIterator<Item = T>,
568    NI: IntoIterator<Item = T>,
569{
570    let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node));
571    heads_ok(to_ok_iter(start), id_fn, neighbors_fn).unwrap()
572}
573
574/// Finds `Ok` nodes in the start set that are not reachable from other nodes in
575/// the start set.
576///
577/// If `start` or `neighbors_fn()` yields an `Err`, this function terminates and
578/// returns the error.
579pub fn heads_ok<T, ID, E, II, NI>(
580    start: II,
581    id_fn: impl Fn(&T) -> ID,
582    mut neighbors_fn: impl FnMut(&T) -> NI,
583) -> Result<HashSet<T>, E>
584where
585    T: Hash + Eq + Clone,
586    ID: Hash + Eq,
587    II: IntoIterator<Item = Result<T, E>>,
588    NI: IntoIterator<Item = Result<T, E>>,
589{
590    let mut heads: HashSet<T> = start.into_iter().try_collect()?;
591    // Do a BFS until we have only one item left in the frontier. That frontier must
592    // have originated from one of the heads, and since there can't be cycles,
593    // it won't be able to eliminate any other heads.
594    let mut frontier: Vec<T> = heads.iter().cloned().collect();
595    let mut visited: HashSet<ID> = heads.iter().map(&id_fn).collect();
596    let mut root_reached = false;
597    while frontier.len() > 1 || (!frontier.is_empty() && root_reached) {
598        frontier = frontier
599            .iter()
600            .flat_map(|node| {
601                let neighbors = neighbors_fn(node).into_iter().collect_vec();
602                if neighbors.is_empty() {
603                    root_reached = true;
604                }
605                neighbors
606            })
607            .try_collect()?;
608        for node in &frontier {
609            heads.remove(node);
610        }
611        frontier.retain(|node| visited.insert(id_fn(node)));
612    }
613    Ok(heads)
614}
615
616/// Finds the closest common neighbor among the `set1` and `set2`.
617pub fn closest_common_node<T, ID, II1, II2, NI>(
618    set1: II1,
619    set2: II2,
620    id_fn: impl Fn(&T) -> ID,
621    mut neighbors_fn: impl FnMut(&T) -> NI,
622) -> Option<T>
623where
624    ID: Hash + Eq,
625    II1: IntoIterator<Item = T>,
626    II2: IntoIterator<Item = T>,
627    NI: IntoIterator<Item = T>,
628{
629    let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node));
630    closest_common_node_ok(to_ok_iter(set1), to_ok_iter(set2), id_fn, neighbors_fn).unwrap()
631}
632
633/// Finds the closest common `Ok` neighbor among the `set1` and `set2`.
634///
635/// If the traverse reached to an `Err`, this function terminates and returns
636/// the error.
637pub fn closest_common_node_ok<T, ID, E, II1, II2, NI>(
638    set1: II1,
639    set2: II2,
640    id_fn: impl Fn(&T) -> ID,
641    mut neighbors_fn: impl FnMut(&T) -> NI,
642) -> Result<Option<T>, E>
643where
644    ID: Hash + Eq,
645    II1: IntoIterator<Item = Result<T, E>>,
646    II2: IntoIterator<Item = Result<T, E>>,
647    NI: IntoIterator<Item = Result<T, E>>,
648{
649    let mut visited1 = HashSet::new();
650    let mut visited2 = HashSet::new();
651
652    // TODO: might be better to leave an Err so long as the work contains at
653    // least one Ok node. If a work1 node is included in visited2, it should be
654    // the closest node even if work2 had previously contained an Err.
655    let mut work1: Vec<Result<T, E>> = set1.into_iter().collect();
656    let mut work2: Vec<Result<T, E>> = set2.into_iter().collect();
657    while !work1.is_empty() || !work2.is_empty() {
658        let mut new_work1 = vec![];
659        for node in work1 {
660            let node = node?;
661            let id: ID = id_fn(&node);
662            if visited2.contains(&id) {
663                return Ok(Some(node));
664            }
665            if visited1.insert(id) {
666                for neighbor in neighbors_fn(&node) {
667                    new_work1.push(neighbor);
668                }
669            }
670        }
671        work1 = new_work1;
672
673        let mut new_work2 = vec![];
674        for node in work2 {
675            let node = node?;
676            let id: ID = id_fn(&node);
677            if visited1.contains(&id) {
678                return Ok(Some(node));
679            }
680            if visited2.insert(id) {
681                for neighbor in neighbors_fn(&node) {
682                    new_work2.push(neighbor);
683                }
684            }
685        }
686        work2 = new_work2;
687    }
688    Ok(None)
689}
690
691fn to_ok_iter<T>(iter: impl IntoIterator<Item = T>) -> impl Iterator<Item = Result<T, Infallible>> {
692    iter.into_iter().map(Ok)
693}
694
695#[cfg(test)]
696mod tests {
697    use std::panic;
698
699    use assert_matches::assert_matches;
700    use maplit::hashmap;
701    use maplit::hashset;
702
703    use super::*;
704
705    #[test]
706    fn test_dfs_ok() {
707        let neighbors = hashmap! {
708            'A' => vec![],
709            'B' => vec![Ok('A'), Err('X')],
710            'C' => vec![Ok('B')],
711        };
712        let id_fn = |node: &char| *node;
713        let neighbors_fn = |node: &char| neighbors[node].clone();
714
715        // Self and neighbor nodes shouldn't be lost at the error.
716        let nodes = dfs_ok([Ok('C')], id_fn, neighbors_fn).collect_vec();
717        assert_eq!(nodes, [Ok('C'), Ok('B'), Err('X'), Ok('A')]);
718    }
719
720    #[test]
721    fn test_topo_order_reverse_linear() {
722        // This graph:
723        //  o C
724        //  o B
725        //  o A
726
727        let neighbors = hashmap! {
728            'A' => vec![],
729            'B' => vec!['A'],
730            'C' => vec!['B'],
731        };
732        let id_fn = |node: &char| *node;
733        let neighbors_fn = |node: &char| neighbors[node].clone();
734
735        let common = topo_order_reverse(vec!['C'], id_fn, neighbors_fn);
736        assert_eq!(common, vec!['C', 'B', 'A']);
737        let common = topo_order_reverse(vec!['C', 'B'], id_fn, neighbors_fn);
738        assert_eq!(common, vec!['C', 'B', 'A']);
739        let common = topo_order_reverse(vec!['B', 'C'], id_fn, neighbors_fn);
740        assert_eq!(common, vec!['C', 'B', 'A']);
741
742        let common = topo_order_reverse_lazy(vec!['C'], id_fn, neighbors_fn).collect_vec();
743        assert_eq!(common, vec!['C', 'B', 'A']);
744        let common = topo_order_reverse_lazy(vec!['C', 'B'], id_fn, neighbors_fn).collect_vec();
745        assert_eq!(common, vec!['C', 'B', 'A']);
746        let common = topo_order_reverse_lazy(vec!['B', 'C'], id_fn, neighbors_fn).collect_vec();
747        assert_eq!(common, vec!['C', 'B', 'A']);
748
749        let common = topo_order_reverse_ord(vec!['C'], id_fn, neighbors_fn);
750        assert_eq!(common, vec!['C', 'B', 'A']);
751        let common = topo_order_reverse_ord(vec!['C', 'B'], id_fn, neighbors_fn);
752        assert_eq!(common, vec!['C', 'B', 'A']);
753        let common = topo_order_reverse_ord(vec!['B', 'C'], id_fn, neighbors_fn);
754        assert_eq!(common, vec!['C', 'B', 'A']);
755    }
756
757    #[test]
758    fn test_topo_order_reverse_merge() {
759        // This graph:
760        //  o F
761        //  |\
762        //  o | E
763        //  | o D
764        //  | o C
765        //  | o B
766        //  |/
767        //  o A
768
769        let neighbors = hashmap! {
770            'A' => vec![],
771            'B' => vec!['A'],
772            'C' => vec!['B'],
773            'D' => vec!['C'],
774            'E' => vec!['A'],
775            'F' => vec!['E', 'D'],
776        };
777        let id_fn = |node: &char| *node;
778        let neighbors_fn = |node: &char| neighbors[node].clone();
779
780        let common = topo_order_reverse(vec!['F'], id_fn, neighbors_fn);
781        assert_eq!(common, vec!['F', 'E', 'D', 'C', 'B', 'A']);
782        let common = topo_order_reverse(vec!['F', 'E', 'C'], id_fn, neighbors_fn);
783        assert_eq!(common, vec!['F', 'D', 'E', 'C', 'B', 'A']);
784        let common = topo_order_reverse(vec!['F', 'D', 'E'], id_fn, neighbors_fn);
785        assert_eq!(common, vec!['F', 'D', 'C', 'B', 'E', 'A']);
786
787        let common = topo_order_reverse_lazy(vec!['F'], id_fn, neighbors_fn).collect_vec();
788        assert_eq!(common, vec!['F', 'E', 'D', 'C', 'B', 'A']);
789        let common =
790            topo_order_reverse_lazy(vec!['F', 'E', 'C'], id_fn, neighbors_fn).collect_vec();
791        assert_eq!(common, vec!['F', 'D', 'E', 'C', 'B', 'A']);
792        let common =
793            topo_order_reverse_lazy(vec!['F', 'D', 'E'], id_fn, neighbors_fn).collect_vec();
794        assert_eq!(common, vec!['F', 'D', 'C', 'B', 'E', 'A']);
795
796        let common = topo_order_reverse_ord(vec!['F'], id_fn, neighbors_fn);
797        assert_eq!(common, vec!['F', 'E', 'D', 'C', 'B', 'A']);
798        let common = topo_order_reverse_ord(vec!['F', 'E', 'C'], id_fn, neighbors_fn);
799        assert_eq!(common, vec!['F', 'E', 'D', 'C', 'B', 'A']);
800        let common = topo_order_reverse_ord(vec!['F', 'D', 'E'], id_fn, neighbors_fn);
801        assert_eq!(common, vec!['F', 'E', 'D', 'C', 'B', 'A']);
802    }
803
804    #[test]
805    fn test_topo_order_reverse_nested_merges() {
806        // This graph:
807        //  o I
808        //  |\
809        //  | o H
810        //  | |\
811        //  | | o G
812        //  | o | F
813        //  | | o E
814        //  o |/ D
815        //  | o C
816        //  o | B
817        //  |/
818        //  o A
819
820        let neighbors = hashmap! {
821            'A' => vec![],
822            'B' => vec!['A'],
823            'C' => vec!['A'],
824            'D' => vec!['B'],
825            'E' => vec!['C'],
826            'F' => vec!['C'],
827            'G' => vec!['E'],
828            'H' => vec!['F', 'G'],
829            'I' => vec!['D', 'H'],
830        };
831        let id_fn = |node: &char| *node;
832        let neighbors_fn = |node: &char| neighbors[node].clone();
833
834        let common = topo_order_reverse(vec!['I'], id_fn, neighbors_fn);
835        assert_eq!(common, vec!['I', 'D', 'B', 'H', 'F', 'G', 'E', 'C', 'A']);
836
837        let common = topo_order_reverse_lazy(vec!['I'], id_fn, neighbors_fn).collect_vec();
838        assert_eq!(common, vec!['I', 'D', 'B', 'H', 'F', 'G', 'E', 'C', 'A']);
839
840        let common = topo_order_reverse_ord(vec!['I'], id_fn, neighbors_fn);
841        assert_eq!(common, vec!['I', 'H', 'G', 'F', 'E', 'D', 'C', 'B', 'A']);
842    }
843
844    #[test]
845    fn test_topo_order_reverse_nested_merges_bad_order() {
846        // This graph:
847        //  o I
848        //  |\
849        //  | |\
850        //  | | |\
851        //  | | | o h (h > I)
852        //  | | |/|
853        //  | | o | G
854        //  | |/| o f
855        //  | o |/ e (e > I, G)
856        //  |/| o D
857        //  o |/ C
858        //  | o b (b > D)
859        //  |/
860        //  o A
861
862        let neighbors = hashmap! {
863            'A' => vec![],
864            'b' => vec!['A'],
865            'C' => vec!['A'],
866            'D' => vec!['b'],
867            'e' => vec!['C', 'b'],
868            'f' => vec!['D'],
869            'G' => vec!['e', 'D'],
870            'h' => vec!['G', 'f'],
871            'I' => vec!['C', 'e', 'G', 'h'],
872        };
873        let id_fn = |node: &char| *node;
874        let neighbors_fn = |node: &char| neighbors[node].clone();
875
876        let common = topo_order_reverse(vec!['I'], id_fn, neighbors_fn);
877        assert_eq!(common, vec!['I', 'h', 'G', 'e', 'C', 'f', 'D', 'b', 'A']);
878
879        let common = topo_order_reverse_lazy(vec!['I'], id_fn, neighbors_fn).collect_vec();
880        assert_eq!(common, vec!['I', 'h', 'G', 'e', 'C', 'f', 'D', 'b', 'A']);
881
882        let common = topo_order_reverse_ord(vec!['I'], id_fn, neighbors_fn);
883        assert_eq!(common, vec!['I', 'h', 'f', 'G', 'e', 'D', 'b', 'C', 'A']);
884    }
885
886    #[test]
887    fn test_topo_order_reverse_merge_bad_fork_order_at_root() {
888        // This graph:
889        //  o E
890        //  |\
891        //  o | D
892        //  | o C
893        //  | o B
894        //  |/
895        //  o a (a > D, B)
896
897        let neighbors = hashmap! {
898            'a' => vec![],
899            'B' => vec!['a'],
900            'C' => vec!['B'],
901            'D' => vec!['a'],
902            'E' => vec!['D', 'C'],
903        };
904        let id_fn = |node: &char| *node;
905        let neighbors_fn = |node: &char| neighbors[node].clone();
906
907        let common = topo_order_reverse(vec!['E'], id_fn, neighbors_fn);
908        assert_eq!(common, vec!['E', 'D', 'C', 'B', 'a']);
909
910        // The root node 'a' is visited before 'C'. If the graph were split there,
911        // the branch 'C->B->a' would be orphaned.
912        let common = topo_order_reverse_lazy(vec!['E'], id_fn, neighbors_fn).collect_vec();
913        assert_eq!(common, vec!['E', 'D', 'C', 'B', 'a']);
914
915        let common = topo_order_reverse_ord(vec!['E'], id_fn, neighbors_fn);
916        assert_eq!(common, vec!['E', 'D', 'C', 'B', 'a']);
917    }
918
919    #[test]
920    fn test_topo_order_reverse_merge_and_linear() {
921        // This graph:
922        //  o G
923        //  |\
924        //  | o F
925        //  o | E
926        //  | o D
927        //  |/
928        //  o C
929        //  o B
930        //  o A
931
932        let neighbors = hashmap! {
933            'A' => vec![],
934            'B' => vec!['A'],
935            'C' => vec!['B'],
936            'D' => vec!['C'],
937            'E' => vec!['C'],
938            'F' => vec!['D'],
939            'G' => vec!['E', 'F'],
940        };
941        let id_fn = |node: &char| *node;
942        let neighbors_fn = |node: &char| neighbors[node].clone();
943
944        let common = topo_order_reverse(vec!['G'], id_fn, neighbors_fn);
945        assert_eq!(common, vec!['G', 'E', 'F', 'D', 'C', 'B', 'A']);
946
947        let common = topo_order_reverse_lazy(vec!['G'], id_fn, neighbors_fn).collect_vec();
948        assert_eq!(common, vec!['G', 'E', 'F', 'D', 'C', 'B', 'A']);
949
950        let common = topo_order_reverse_ord(vec!['G'], id_fn, neighbors_fn);
951        assert_eq!(common, vec!['G', 'F', 'E', 'D', 'C', 'B', 'A']);
952
953        // Iterator can be lazy for linear chunks.
954        let neighbors_fn = |node: &char| to_ok_iter(neighbors[node].iter().copied());
955        let cycle_fn = |_| panic!("graph has cycle");
956        let mut inner_iter = TopoOrderReverseLazyInner::empty();
957        inner_iter.extend([Ok('G')]);
958        assert_eq!(
959            inner_iter.next(id_fn, neighbors_fn, cycle_fn),
960            Some(Ok('G'))
961        );
962        assert!(!inner_iter.start.is_empty());
963        assert!(inner_iter.result.is_empty());
964        assert_eq!(
965            iter::from_fn(|| inner_iter.next(id_fn, neighbors_fn, cycle_fn))
966                .take(4)
967                .collect_vec(),
968            ['E', 'F', 'D', 'C'].map(Ok),
969        );
970        assert!(!inner_iter.start.is_empty());
971        assert!(inner_iter.result.is_empty());
972
973        // Run each step of lazy iterator by using low-level function.
974        let mut start = vec!['G'];
975        let next = |start: &mut Vec<char>| {
976            topo_order_reverse_chunked(start, id_fn, neighbors_fn, cycle_fn).unwrap()
977        };
978        assert_eq!(next(&mut start), ['G'].into());
979        assert_eq!(start, ['E', 'F']);
980        assert_eq!(next(&mut start), ['E', 'F', 'D', 'C'].into());
981        assert_eq!(start, ['B']);
982        assert_eq!(next(&mut start), ['B'].into());
983        assert_eq!(start, ['A']);
984        assert_eq!(next(&mut start), ['A'].into());
985        assert!(start.is_empty());
986        assert!(next(&mut start).is_empty());
987        assert!(start.is_empty());
988    }
989
990    #[test]
991    fn test_topo_order_reverse_merge_and_linear_bad_fork_order() {
992        // This graph:
993        //  o G
994        //  |\
995        //  o | F
996        //  o | E
997        //  | o D
998        //  |/
999        //  o c (c > E, D)
1000        //  o B
1001        //  o A
1002
1003        let neighbors = hashmap! {
1004            'A' => vec![],
1005            'B' => vec!['A'],
1006            'c' => vec!['B'],
1007            'D' => vec!['c'],
1008            'E' => vec!['c'],
1009            'F' => vec!['E'],
1010            'G' => vec!['F', 'D'],
1011        };
1012        let id_fn = |node: &char| *node;
1013        let neighbors_fn = |node: &char| neighbors[node].clone();
1014
1015        let common = topo_order_reverse(vec!['G'], id_fn, neighbors_fn);
1016        assert_eq!(common, vec!['G', 'F', 'E', 'D', 'c', 'B', 'A']);
1017
1018        let common = topo_order_reverse_lazy(vec!['G'], id_fn, neighbors_fn).collect_vec();
1019        assert_eq!(common, vec!['G', 'F', 'E', 'D', 'c', 'B', 'A']);
1020
1021        let common = topo_order_reverse_ord(vec!['G'], id_fn, neighbors_fn);
1022        assert_eq!(common, vec!['G', 'F', 'E', 'D', 'c', 'B', 'A']);
1023
1024        // Iterator can be lazy for linear chunks. The node 'c' is visited before 'D',
1025        // but it will be processed lazily.
1026        let neighbors_fn = |node: &char| to_ok_iter(neighbors[node].iter().copied());
1027        let cycle_fn = |_| panic!("graph has cycle");
1028        let mut inner_iter = TopoOrderReverseLazyInner::empty();
1029        inner_iter.extend([Ok('G')]);
1030        assert_eq!(
1031            inner_iter.next(id_fn, neighbors_fn, cycle_fn),
1032            Some(Ok('G'))
1033        );
1034        assert!(!inner_iter.start.is_empty());
1035        assert!(inner_iter.result.is_empty());
1036        assert_eq!(
1037            iter::from_fn(|| inner_iter.next(id_fn, neighbors_fn, cycle_fn))
1038                .take(4)
1039                .collect_vec(),
1040            ['F', 'E', 'D', 'c'].map(Ok),
1041        );
1042        assert!(!inner_iter.start.is_empty());
1043        assert!(inner_iter.result.is_empty());
1044
1045        // Run each step of lazy iterator by using low-level function.
1046        let mut start = vec!['G'];
1047        let next = |start: &mut Vec<char>| {
1048            topo_order_reverse_chunked(start, id_fn, neighbors_fn, cycle_fn).unwrap()
1049        };
1050        assert_eq!(next(&mut start), ['G'].into());
1051        assert_eq!(start, ['F', 'D']);
1052        assert_eq!(next(&mut start), ['F', 'E', 'D', 'c'].into());
1053        assert_eq!(start, ['B']);
1054        assert_eq!(next(&mut start), ['B'].into());
1055        assert_eq!(start, ['A']);
1056        assert_eq!(next(&mut start), ['A'].into());
1057        assert!(start.is_empty());
1058        assert!(next(&mut start).is_empty());
1059        assert!(start.is_empty());
1060    }
1061
1062    #[test]
1063    fn test_topo_order_reverse_merge_and_linear_bad_merge_order() {
1064        // This graph:
1065        //  o G
1066        //  |\
1067        //  o | f (f > G)
1068        //  o | e
1069        //  | o d (d > G)
1070        //  |/
1071        //  o C
1072        //  o B
1073        //  o A
1074
1075        let neighbors = hashmap! {
1076            'A' => vec![],
1077            'B' => vec!['A'],
1078            'C' => vec!['B'],
1079            'd' => vec!['C'],
1080            'e' => vec!['C'],
1081            'f' => vec!['e'],
1082            'G' => vec!['f', 'd'],
1083        };
1084        let id_fn = |node: &char| *node;
1085        let neighbors_fn = |node: &char| neighbors[node].clone();
1086
1087        let common = topo_order_reverse(vec!['G'], id_fn, neighbors_fn);
1088        assert_eq!(common, vec!['G', 'f', 'e', 'd', 'C', 'B', 'A']);
1089
1090        let common = topo_order_reverse_lazy(vec!['G'], id_fn, neighbors_fn).collect_vec();
1091        assert_eq!(common, vec!['G', 'f', 'e', 'd', 'C', 'B', 'A']);
1092
1093        let common = topo_order_reverse_ord(vec!['G'], id_fn, neighbors_fn);
1094        assert_eq!(common, vec!['G', 'f', 'e', 'd', 'C', 'B', 'A']);
1095
1096        // Iterator can be lazy for linear chunks.
1097        let neighbors_fn = |node: &char| to_ok_iter(neighbors[node].iter().copied());
1098        let cycle_fn = |_| panic!("graph has cycle");
1099        let mut inner_iter = TopoOrderReverseLazyInner::empty();
1100        inner_iter.extend([Ok('G')]);
1101        assert_eq!(
1102            inner_iter.next(id_fn, neighbors_fn, cycle_fn),
1103            Some(Ok('G'))
1104        );
1105        assert!(!inner_iter.start.is_empty());
1106        assert!(inner_iter.result.is_empty());
1107        assert_eq!(
1108            iter::from_fn(|| inner_iter.next(id_fn, neighbors_fn, cycle_fn))
1109                .take(4)
1110                .collect_vec(),
1111            ['f', 'e', 'd', 'C'].map(Ok),
1112        );
1113        assert!(!inner_iter.start.is_empty());
1114        assert!(inner_iter.result.is_empty());
1115
1116        // Run each step of lazy iterator by using low-level function.
1117        let mut start = vec!['G'];
1118        let next = |start: &mut Vec<char>| {
1119            topo_order_reverse_chunked(start, id_fn, neighbors_fn, cycle_fn).unwrap()
1120        };
1121        assert_eq!(next(&mut start), ['G'].into());
1122        assert_eq!(start, ['f', 'd']);
1123        assert_eq!(next(&mut start), ['f', 'e', 'd', 'C'].into());
1124        assert_eq!(start, ['B']);
1125        assert_eq!(next(&mut start), ['B'].into());
1126        assert_eq!(start, ['A']);
1127        assert_eq!(next(&mut start), ['A'].into());
1128        assert!(start.is_empty());
1129        assert!(next(&mut start).is_empty());
1130        assert!(start.is_empty());
1131    }
1132
1133    #[test]
1134    fn test_topo_order_reverse_multiple_heads() {
1135        // This graph:
1136        //  o F
1137        //  |\
1138        //  o | E
1139        //  | o D
1140        //  | | o C
1141        //  | | |
1142        //  | | o B
1143        //  | |/
1144        //  |/
1145        //  o A
1146
1147        let neighbors = hashmap! {
1148            'A' => vec![],
1149            'B' => vec!['A'],
1150            'C' => vec!['B'],
1151            'D' => vec!['A'],
1152            'E' => vec!['A'],
1153            'F' => vec!['E', 'D'],
1154        };
1155        let id_fn = |node: &char| *node;
1156        let neighbors_fn = |node: &char| neighbors[node].clone();
1157
1158        let common = topo_order_reverse(vec!['F', 'C'], id_fn, neighbors_fn);
1159        assert_eq!(common, vec!['F', 'E', 'D', 'C', 'B', 'A']);
1160
1161        let common = topo_order_reverse_lazy(vec!['F', 'C'], id_fn, neighbors_fn).collect_vec();
1162        assert_eq!(common, vec!['F', 'E', 'D', 'C', 'B', 'A']);
1163
1164        let common = topo_order_reverse_ord(vec!['F', 'C'], id_fn, neighbors_fn);
1165        assert_eq!(common, vec!['F', 'E', 'D', 'C', 'B', 'A']);
1166    }
1167
1168    #[test]
1169    fn test_topo_order_reverse_multiple_roots() {
1170        // This graph:
1171        //  o D
1172        //  | \
1173        //  o | C
1174        //    o B
1175        //    o A
1176
1177        let neighbors = hashmap! {
1178            'A' => vec![],
1179            'B' => vec!['A'],
1180            'C' => vec![],
1181            'D' => vec!['C', 'B'],
1182        };
1183        let id_fn = |node: &char| *node;
1184        let neighbors_fn = |node: &char| neighbors[node].clone();
1185
1186        let common = topo_order_reverse(vec!['D'], id_fn, neighbors_fn);
1187        assert_eq!(common, vec!['D', 'C', 'B', 'A']);
1188
1189        let common = topo_order_reverse_lazy(vec!['D'], id_fn, neighbors_fn).collect_vec();
1190        assert_eq!(common, vec!['D', 'C', 'B', 'A']);
1191
1192        let common = topo_order_reverse_ord(vec!['D'], id_fn, neighbors_fn);
1193        assert_eq!(common, vec!['D', 'C', 'B', 'A']);
1194    }
1195
1196    #[test]
1197    fn test_topo_order_reverse_cycle_linear() {
1198        // This graph:
1199        //  o C
1200        //  o B
1201        //  o A (to C)
1202
1203        let neighbors = hashmap! {
1204            'A' => vec!['C'],
1205            'B' => vec!['A'],
1206            'C' => vec!['B'],
1207        };
1208        let id_fn = |node: &char| *node;
1209        let neighbors_fn = |node: &char| neighbors[node].clone();
1210
1211        let result = panic::catch_unwind(|| topo_order_reverse(vec!['C'], id_fn, neighbors_fn));
1212        assert!(result.is_err());
1213
1214        topo_order_reverse_lazy(vec!['C'], id_fn, neighbors_fn)
1215            .take(3)
1216            .collect_vec(); // sanity check
1217        let result = panic::catch_unwind(|| {
1218            topo_order_reverse_lazy(vec!['C'], id_fn, neighbors_fn)
1219                .take(4)
1220                .collect_vec()
1221        });
1222        assert!(result.is_err());
1223
1224        let result = panic::catch_unwind(|| topo_order_reverse_ord(vec!['C'], id_fn, neighbors_fn));
1225        assert!(result.is_err());
1226
1227        // Try again with non-panicking cycle handler
1228        let neighbors_fn = |node: &char| neighbors[node].iter().copied().map(Ok);
1229        let cycle_fn = |node: char| node;
1230        assert_matches!(
1231            topo_order_reverse_ok([Ok('C')], id_fn, neighbors_fn, cycle_fn),
1232            Err('C' | 'B' | 'A')
1233        );
1234        assert_matches!(
1235            topo_order_reverse_lazy_ok([Ok('C')], id_fn, neighbors_fn, cycle_fn).nth(3),
1236            Some(Err('C' | 'B' | 'A'))
1237        );
1238        assert_matches!(
1239            topo_order_reverse_ord_ok([Ok('C')], id_fn, neighbors_fn, cycle_fn),
1240            Err('C' | 'B' | 'A')
1241        );
1242    }
1243
1244    #[test]
1245    fn test_topo_order_reverse_cycle_to_branchy_sub_graph() {
1246        // This graph:
1247        //  o D
1248        //  |\
1249        //  | o C
1250        //  |/
1251        //  o B
1252        //  o A (to C)
1253
1254        let neighbors = hashmap! {
1255            'A' => vec!['C'],
1256            'B' => vec!['A'],
1257            'C' => vec!['B'],
1258            'D' => vec!['B', 'C'],
1259        };
1260        let id_fn = |node: &char| *node;
1261        let neighbors_fn = |node: &char| neighbors[node].clone();
1262
1263        let result = panic::catch_unwind(|| topo_order_reverse(vec!['D'], id_fn, neighbors_fn));
1264        assert!(result.is_err());
1265
1266        topo_order_reverse_lazy(vec!['D'], id_fn, neighbors_fn)
1267            .take(4)
1268            .collect_vec(); // sanity check
1269        let result = panic::catch_unwind(|| {
1270            topo_order_reverse_lazy(vec!['D'], id_fn, neighbors_fn)
1271                .take(5)
1272                .collect_vec()
1273        });
1274        assert!(result.is_err());
1275
1276        let result = panic::catch_unwind(|| topo_order_reverse_ord(vec!['D'], id_fn, neighbors_fn));
1277        assert!(result.is_err());
1278
1279        // Try again with non-panicking cycle handler
1280        let neighbors_fn = |node: &char| neighbors[node].iter().copied().map(Ok);
1281        let cycle_fn = |node: char| node;
1282        assert_matches!(
1283            topo_order_reverse_ok([Ok('D')], id_fn, neighbors_fn, cycle_fn),
1284            Err('C' | 'B' | 'A')
1285        );
1286        assert_matches!(
1287            topo_order_reverse_lazy_ok([Ok('D')], id_fn, neighbors_fn, cycle_fn).nth(4),
1288            Some(Err('C' | 'B' | 'A'))
1289        );
1290        assert_matches!(
1291            topo_order_reverse_ord_ok([Ok('D')], id_fn, neighbors_fn, cycle_fn),
1292            Err('C' | 'B' | 'A')
1293        );
1294    }
1295
1296    #[test]
1297    fn test_topo_order_reverse_lazy_cycle_within_branchy_sub_graph() {
1298        // This graph:
1299        //  o D
1300        //  |\
1301        //  | o C
1302        //  |/
1303        //  o B (to C)
1304        //  o A
1305
1306        let neighbors = hashmap! {
1307            'A' => vec![],
1308            'B' => vec!['A', 'C'],
1309            'C' => vec!['B'],
1310            'D' => vec!['B', 'C'],
1311        };
1312        let id_fn = |node: &char| *node;
1313        let neighbors_fn = |node: &char| neighbors[node].clone();
1314
1315        let result =
1316            panic::catch_unwind(|| topo_order_reverse_lazy(['D'], id_fn, neighbors_fn).nth(1));
1317        assert!(result.is_err());
1318
1319        // Try again with non-panicking cycle handler
1320        let neighbors_fn = |node: &char| neighbors[node].iter().copied().map(Ok);
1321        let cycle_fn = |node: char| node;
1322        assert_matches!(
1323            topo_order_reverse_lazy_ok([Ok('D')], id_fn, neighbors_fn, cycle_fn).nth(1),
1324            Some(Err('C' | 'B'))
1325        );
1326
1327        // Try again with low-level function
1328        let mut start = vec!['D'];
1329        topo_order_reverse_chunked(&mut start, id_fn, neighbors_fn, cycle_fn).unwrap();
1330        assert_matches!(
1331            topo_order_reverse_chunked(&mut start, id_fn, neighbors_fn, cycle_fn),
1332            Err('C' | 'B')
1333        );
1334    }
1335
1336    #[test]
1337    fn test_topo_order_ok() {
1338        let neighbors = hashmap! {
1339            'A' => vec![Err('Y')],
1340            'B' => vec![Ok('A'), Err('X')],
1341            'C' => vec![Ok('B')],
1342        };
1343        let id_fn = |node: &char| *node;
1344        let neighbors_fn = |node: &char| neighbors[node].clone();
1345        let cycle_fn = |_| panic!("graph has cycle");
1346
1347        // Terminates at Err('X') no matter if the sorting order is forward or
1348        // reverse. The visiting order matters.
1349        let result = topo_order_forward_ok([Ok('C')], id_fn, neighbors_fn, cycle_fn);
1350        assert_eq!(result, Err('X'));
1351        let result = topo_order_reverse_ok([Ok('C')], id_fn, neighbors_fn, cycle_fn);
1352        assert_eq!(result, Err('X'));
1353        let nodes =
1354            topo_order_reverse_lazy_ok([Ok('C')], id_fn, neighbors_fn, cycle_fn).collect_vec();
1355        assert_eq!(nodes, [Ok('C'), Ok('B'), Err('X')]);
1356        let result = topo_order_reverse_ord_ok([Ok('C')], id_fn, neighbors_fn, cycle_fn);
1357        assert_eq!(result, Err('X'));
1358    }
1359
1360    #[test]
1361    fn test_closest_common_node_tricky() {
1362        // Test this case where A is the shortest distance away, but we still want the
1363        // result to be B because A is an ancestor of B. In other words, we want
1364        // to minimize the longest distance.
1365        //
1366        //  E       H
1367        //  |\     /|
1368        //  | D   G |
1369        //  | C   F |
1370        //   \ \ / /
1371        //    \ B /
1372        //     \|/
1373        //      A
1374
1375        let neighbors = hashmap! {
1376            'A' => vec![],
1377            'B' => vec!['A'],
1378            'C' => vec!['B'],
1379            'D' => vec!['C'],
1380            'E' => vec!['A','D'],
1381            'F' => vec!['B'],
1382            'G' => vec!['F'],
1383            'H' => vec!['A', 'G'],
1384        };
1385
1386        let common = closest_common_node(
1387            vec!['E'],
1388            vec!['H'],
1389            |node| *node,
1390            |node| neighbors[node].clone(),
1391        );
1392
1393        // TODO: fix the implementation to return B
1394        assert_eq!(common, Some('A'));
1395    }
1396
1397    #[test]
1398    fn test_closest_common_node_ok() {
1399        let neighbors = hashmap! {
1400            'A' => vec![Err('Y')],
1401            'B' => vec![Ok('A')],
1402            'C' => vec![Ok('A')],
1403            'D' => vec![Err('X')],
1404        };
1405        let id_fn = |node: &char| *node;
1406        let neighbors_fn = |node: &char| neighbors[node].clone();
1407
1408        let result = closest_common_node_ok([Ok('B')], [Ok('C')], id_fn, neighbors_fn);
1409        assert_eq!(result, Ok(Some('A')));
1410        let result = closest_common_node_ok([Ok('C')], [Ok('D')], id_fn, neighbors_fn);
1411        assert_eq!(result, Err('X'));
1412        let result = closest_common_node_ok([Ok('C')], [Err('Z')], id_fn, neighbors_fn);
1413        assert_eq!(result, Err('Z'));
1414    }
1415
1416    #[test]
1417    fn test_heads_mixed() {
1418        // Test the uppercase letters are in the start set
1419        //
1420        //  D F
1421        //  |/|
1422        //  C e
1423        //  |/
1424        //  b
1425        //  |
1426        //  A
1427
1428        let neighbors = hashmap! {
1429            'A' => vec![],
1430            'b' => vec!['A'],
1431            'C' => vec!['b'],
1432            'D' => vec!['C'],
1433            'e' => vec!['b'],
1434            'F' => vec!['C', 'e'],
1435        };
1436
1437        let actual = heads(
1438            vec!['A', 'C', 'D', 'F'],
1439            |node| *node,
1440            |node| neighbors[node].clone(),
1441        );
1442        assert_eq!(actual, hashset!['D', 'F']);
1443
1444        // Check with a different order in the start set
1445        let actual = heads(
1446            vec!['F', 'D', 'C', 'A'],
1447            |node| *node,
1448            |node| neighbors[node].clone(),
1449        );
1450        assert_eq!(actual, hashset!['D', 'F']);
1451    }
1452
1453    #[test]
1454    fn test_heads_ok() {
1455        let neighbors = hashmap! {
1456            'A' => vec![],
1457            'B' => vec![Ok('A'), Err('X')],
1458            'C' => vec![Ok('B')],
1459        };
1460        let id_fn = |node: &char| *node;
1461        let neighbors_fn = |node: &char| neighbors[node].clone();
1462
1463        let result = heads_ok([Ok('C')], id_fn, neighbors_fn);
1464        assert_eq!(result, Ok(hashset! {'C'}));
1465        let result = heads_ok([Ok('B')], id_fn, neighbors_fn);
1466        assert_eq!(result, Ok(hashset! {'B'}));
1467        let result = heads_ok([Ok('A')], id_fn, neighbors_fn);
1468        assert_eq!(result, Ok(hashset! {'A'}));
1469        let result = heads_ok([Ok('C'), Ok('B')], id_fn, neighbors_fn);
1470        assert_eq!(result, Err('X'));
1471        let result = heads_ok([Ok('C'), Ok('A')], id_fn, neighbors_fn);
1472        assert_eq!(result, Err('X'));
1473    }
1474}