wolf_graph/algo/
depth_first_search.rs

1use anyhow::Result;
2
3use crate::{DFSVisitor, EdgeID, Edges, NodeID, Nodes, VisitableGraph};
4
5enum State {
6    Discovered,
7    Finished,
8}
9
10/// A trait for a graph that supports depth-first search.
11pub trait DepthFirstSearch<G, O>
12where
13    G: VisitableGraph
14{
15    fn depth_first_search_opt<Visitor>(&self, visitor: &mut Visitor, roots: &Nodes, roots_only: bool, excluded_edge: Option<impl AsRef<EdgeID>>) -> Result<O>
16    where Visitor: DFSVisitor<Graph = G, Output = O>;
17
18    fn depth_first_search<Visitor>(&self, visitor: &mut Visitor) -> Result<O>
19    where Visitor: DFSVisitor<Graph = G, Output = O> {
20        self.depth_first_search_opt(visitor, &Nodes::new(), false, None::<EdgeID>)
21    }
22}
23
24impl<G, O> DepthFirstSearch<G, O> for G
25where
26    G: VisitableGraph
27{
28    fn depth_first_search_opt<Visitor>(&self, visitor: &mut Visitor, roots: &Nodes, roots_only: bool, excluded_edge: Option<impl AsRef<EdgeID>>) -> Result<O>
29    where
30        Visitor: DFSVisitor<Graph = G, Output = O>,
31    {
32        for node in self.all_nodes() {
33            if let Some(result) = visitor.init_node(self, &node)? {
34                return Ok(result);
35            }
36        }
37
38        let mut states = std::collections::HashMap::new();
39
40        let excluded_edge = excluded_edge.as_ref();
41        for root in roots.iter() {
42            if let Some(result) = search_from_root(self, visitor, &mut states, root, excluded_edge)? {
43                return Ok(result);
44            }
45        }
46        if !roots_only {
47            for node in self.all_nodes() {
48                if let Some(result) = search_from_root(self, visitor, &mut states, &node, excluded_edge)? {
49                    return Ok(result);
50                }
51            }
52        }
53
54        return visitor.finish(self);
55
56        fn search_from_root<Graph, Visitor>(
57            graph: &Graph,
58            visitor: &mut Visitor,
59            states: &mut std::collections::HashMap<NodeID, State>,
60            root: &NodeID,
61            excluded_edge: Option<impl AsRef<EdgeID>>,
62        ) -> Result<Option<Visitor::Output>>
63        where
64            Graph: VisitableGraph,
65            Visitor: DFSVisitor<Graph = Graph>,
66        {
67            if states.contains_key(root) {
68                return Ok(None);
69            }
70            if let Some(result) = visitor.start_node(graph, root)? {
71                return Ok(Some(result));
72            }
73            states.insert(root.clone(), State::Discovered);
74            if let Some(result) = visitor.discover_node(graph, root)? {
75                return Ok(Some(result));
76            }
77            let excluded_edge = excluded_edge.as_ref().map(|edge| edge.as_ref());
78            let out_edges: Edges = graph.out_edges(root)?
79                .iter()
80                .filter(|&edge| Some(edge) != excluded_edge).cloned()
81                .collect();
82            let mut stack: Vec<(NodeID, Option<EdgeID>, std::collections::BTreeSet<EdgeID>)> = Vec::new();
83            stack.push((root.clone(), None, out_edges));
84            while let Some((source, finished_edge, remaining_out_edges)) = stack.pop() {
85                let mut source = source;
86                if let Some(finished_edge) = finished_edge {
87                    if let Some(result) = visitor.finish_edge(graph, &finished_edge)? {
88                        return Ok(Some(result));
89                    }
90                }
91
92                let mut remaining_out_edges = remaining_out_edges;
93                while let Some(edge) = remaining_out_edges.pop_last() {
94                    let target = graph.target(&edge)?;
95                    if let Some(result) = visitor.examine_edge(graph, &edge)? {
96                        return Ok(Some(result));
97                    }
98                    match states.get(&target) {
99                        None => {
100                            if let Some(result) = visitor.tree_edge(graph, &edge)? {
101                                return Ok(Some(result));
102                            }
103                            stack.push((source, Some(edge), remaining_out_edges));
104                            source = target;
105                            states.insert(source.clone(), State::Discovered);
106                            if let Some(result) = visitor.discover_node(graph, &source)? {
107                                return Ok(Some(result));
108                            }
109                            remaining_out_edges = graph.out_edges(&source)?;
110                        }
111                        Some(State::Discovered) => {
112                            if let Some(result) = visitor.back_edge(graph, &edge)? {
113                                return Ok(Some(result));
114                            }
115                            if let Some(result) = visitor.finish_edge(graph, &edge)? {
116                                return Ok(Some(result));
117                            }
118                        }
119                        Some(State::Finished) => {
120                            if let Some(result) = visitor.forward_or_cross_edge(graph, &edge)? {
121                                return Ok(Some(result));
122                            }
123                            if let Some(result) = visitor.finish_edge(graph, &edge)? {
124                                return Ok(Some(result));
125                            }
126                        }
127                    }
128                }
129                states.insert(source.clone(), State::Finished);
130                if let Some(result) = visitor.finish_node(graph, &source)? {
131                    return Ok(Some(result));
132                }
133            }
134            Ok(None)
135        }
136    }
137}