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