rustworkx_core/traversal/bfs_visit.rs
1// Licensed under the Apache License, Version 2.0 (the "License"); you may
2// not use this file except in compliance with the License. You may obtain
3// a copy of the License at
4//
5// http://www.apache.org/licenses/LICENSE-2.0
6//
7// Unless required by applicable law or agreed to in writing, software
8// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
9// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
10// License for the specific language governing permissions and limitations
11// under the License.
12
13use petgraph::visit::{ControlFlow, EdgeRef, IntoEdges, VisitMap, Visitable};
14use std::collections::VecDeque;
15
16use super::try_control;
17
18/// A breadth first search (BFS) visitor event.
19#[derive(Copy, Clone, Debug)]
20pub enum BfsEvent<N, E> {
21 Discover(N),
22 /// An edge of the tree formed by the traversal.
23 TreeEdge(N, N, E),
24 /// An edge that does not belong to the tree.
25 NonTreeEdge(N, N, E),
26 /// For an edge *(u, v)*, if node *v* is currently in the queue
27 /// at the time of examination, then it is a gray-target edge.
28 GrayTargetEdge(N, N, E),
29 /// For an edge *(u, v)*, if node *v* has been removed from the queue
30 /// at the time of examination, then it is a black-target edge.
31 BlackTargetEdge(N, N, E),
32 /// All edges from a node have been reported.
33 Finish(N),
34}
35
36/// An iterative breadth first search.
37///
38/// Starting points are the nodes in the iterator `starts` (specify just one
39/// start vertex *x* by using `Some(x)`).
40///
41/// The traversal emits discovery and finish events for each reachable vertex,
42/// and edge classification of each reachable edge. `visitor` is called for each
43/// event, see [`BfsEvent`] for possible values.
44///
45/// The return value should implement the trait [`ControlFlow`], and can be used to change
46/// the control flow of the search.
47///
48/// [`Control`](petgraph::visit::Control) Implements [`ControlFlow`] such that `Control::Continue` resumes the search.
49/// `Control::Break` will stop the visit early, returning the contained value.
50/// `Control::Prune` will stop traversing any additional edges from the current
51/// node and proceed immediately to the `Finish` event.
52///
53/// There are implementations of [`ControlFlow`] for `()`, and [`Result<C, E>`] where
54/// `C: ControlFlow`. The implementation for `()` will continue until finished.
55/// For [`Result`], upon encountering an `E` it will break, otherwise acting the same as `C`.
56///
57/// ***Panics** if you attempt to prune a node from its `Finish` event.
58///
59/// Pseudo-code for the BFS algorithm is listed below, with the annotated event points,
60/// for which the given visitor object will be called with the appropriate method.
61///
62/// ```norust
63/// // G - graph, s - single source node
64/// BFS(G, s)
65/// let color be a mapping // color[u] - vertex u color WHITE/GRAY/BLACK
66/// for each u in G // u - vertex in G
67/// color[u] := WHITE // color all vertices as undiscovered
68/// end for
69/// let Q be a queue
70/// ENQUEUE(Q, s)
71/// color[s] := GRAY // event: Discover(s)
72/// while (Q is not empty)
73/// u := DEQUEUE(Q)
74/// for each v, w in OutEdges(G, u) // v - target vertex, w - edge weight
75/// if (WHITE = color[v]) // event: TreeEdge(u, v, w)
76/// color[v] := GRAY // event: Discover(v)
77/// ENQUEUE(Q, v)
78/// else // event: NonTreeEdge(u, v, w)
79/// if (GRAY = color[v]) // event: GrayTargetEdge(u, v, w)
80/// ...
81/// elif (BLACK = color[v]) // event: BlackTargetEdge(u, v, w)
82/// ...
83/// end for
84/// color[u] := BLACK // event: Finish(u)
85/// end while
86/// ```
87///
88/// # Example returning [`Control`](petgraph::visit::Control).
89///
90/// Find a path from vertex 0 to 5, and exit the visit as soon as we reach
91/// the goal vertex.
92///
93/// ```
94/// use rustworkx_core::petgraph::prelude::*;
95/// use rustworkx_core::petgraph::graph::node_index as n;
96/// use rustworkx_core::petgraph::visit::Control;
97///
98/// use rustworkx_core::traversal::{BfsEvent, breadth_first_search};
99///
100/// let gr: Graph<(), ()> = Graph::from_edges(&[
101/// (0, 1), (0, 2), (0, 3),
102/// (1, 3),
103/// (2, 3), (2, 4),
104/// (4, 0), (4, 5),
105/// ]);
106///
107/// // record each predecessor, mapping node → node
108/// let mut predecessor = vec![NodeIndex::end(); gr.node_count()];
109/// let start = n(0);
110/// let goal = n(5);
111/// breadth_first_search(&gr, Some(start), |event| {
112/// if let BfsEvent::TreeEdge(u, v, _) = event {
113/// predecessor[v.index()] = u;
114/// if v == goal {
115/// return Control::Break(v);
116/// }
117/// }
118/// Control::Continue
119/// });
120///
121/// let mut next = goal;
122/// let mut path = vec![next];
123/// while next != start {
124/// let pred = predecessor[next.index()];
125/// path.push(pred);
126/// next = pred;
127/// }
128/// path.reverse();
129/// assert_eq!(&path, &[n(0), n(2), n(4), n(5)]);
130/// ```
131///
132/// # Example returning a `Result`.
133/// ```
134/// use rustworkx_core::petgraph::graph::node_index as n;
135/// use rustworkx_core::petgraph::prelude::*;
136///
137/// use rustworkx_core::traversal::{BfsEvent, breadth_first_search};
138///
139/// let gr: Graph<(), ()> = Graph::from_edges(&[(0, 1), (1, 2), (1, 1), (2, 1)]);
140/// let start = n(0);
141/// let mut non_tree_edges = 0;
142///
143/// #[derive(Debug)]
144/// struct NonTreeEdgeFound {
145/// source: NodeIndex,
146/// target: NodeIndex,
147/// }
148///
149/// // Stop the search, the first time a BackEdge is encountered.
150/// let result = breadth_first_search(&gr, Some(start), |event| {
151/// match event {
152/// BfsEvent::NonTreeEdge(u, v, _) => {
153/// non_tree_edges += 1;
154/// // the implementation of ControlFlow for Result,
155/// // treats this Err value as Continue::Break
156/// Err(NonTreeEdgeFound {source: u, target: v})
157/// }
158/// // In the cases where Ok(()) is returned,
159/// // Result falls back to the implementation of Control on the value ().
160/// // In the case of (), this is to always return Control::Continue.
161/// // continuing the search.
162/// _ => Ok(()),
163/// }
164/// });
165///
166/// assert_eq!(non_tree_edges, 1);
167/// println!("number of non-tree edges encountered: {}", non_tree_edges);
168/// println!("non-tree edge: ({:?})", result.unwrap_err());
169/// ```
170pub fn breadth_first_search<G, I, F, C>(graph: G, starts: I, mut visitor: F) -> C
171where
172 G: IntoEdges + Visitable,
173 I: IntoIterator<Item = G::NodeId>,
174 F: FnMut(BfsEvent<G::NodeId, &G::EdgeWeight>) -> C,
175 C: ControlFlow,
176{
177 let discovered = &mut graph.visit_map();
178 let finished = &mut graph.visit_map();
179
180 for start in starts {
181 // `bfs_visitor` returns a "signal" to either continue or exit early
182 // but it never "prunes", so we use `unreachable`.
183 try_control!(
184 bfs_visitor(graph, start, &mut visitor, discovered, finished),
185 unreachable!()
186 );
187 }
188 C::continuing()
189}
190
191fn bfs_visitor<G, F, C>(
192 graph: G,
193 u: G::NodeId,
194 visitor: &mut F,
195 discovered: &mut G::Map,
196 finished: &mut G::Map,
197) -> C
198where
199 G: IntoEdges + Visitable,
200 F: FnMut(BfsEvent<G::NodeId, &G::EdgeWeight>) -> C,
201 C: ControlFlow,
202{
203 if !discovered.visit(u) {
204 return C::continuing();
205 }
206
207 try_control!(visitor(BfsEvent::Discover(u)), {}, {
208 let mut stack: VecDeque<G::NodeId> = VecDeque::new();
209 stack.push_front(u);
210
211 while let Some(u) = stack.pop_front() {
212 for edge in graph.edges(u) {
213 let v = edge.target();
214 if !discovered.is_visited(&v) {
215 try_control!(visitor(BfsEvent::TreeEdge(u, v, edge.weight())), continue);
216 discovered.visit(v);
217 try_control!(visitor(BfsEvent::Discover(v)), continue);
218 stack.push_back(v);
219 } else {
220 // non - tree edge.
221 try_control!(
222 visitor(BfsEvent::NonTreeEdge(u, v, edge.weight())),
223 continue
224 );
225
226 if !finished.is_visited(&v) {
227 try_control!(
228 visitor(BfsEvent::GrayTargetEdge(u, v, edge.weight())),
229 continue
230 );
231 } else {
232 try_control!(
233 visitor(BfsEvent::BlackTargetEdge(u, v, edge.weight())),
234 continue
235 );
236 }
237 }
238 }
239
240 let first_finish = finished.visit(u);
241 debug_assert!(first_finish);
242 try_control!(
243 visitor(BfsEvent::Finish(u)),
244 panic!("Pruning on the `BfsEvent::Finish` is not supported!")
245 );
246 }
247 });
248
249 C::continuing()
250}