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/// The pseudo-code for the BFS algorithm is listed below, with the annotated
60/// event points, for which the given visitor object will be called with the
61/// appropriate method.
62///
63/// ```norust
64/// BFS(G, s)
65/// for each vertex u in V
66/// color[u] := WHITE
67/// end for
68/// color[s] := GRAY
69/// EQUEUE(Q, s) discover vertex s
70/// while (Q != Ø)
71/// u := DEQUEUE(Q)
72/// for each vertex v in Adj[u] (u,v) is a tree edge
73/// if (color[v] = WHITE)
74/// color[v] = GRAY
75/// else (u,v) is a non - tree edge
76/// if (color[v] = GRAY) (u,v) has a gray target
77/// ...
78/// else if (color[v] = BLACK) (u,v) has a black target
79/// ...
80/// end for
81/// color[u] := BLACK finish vertex u
82/// end while
83/// ```
84///
85/// # Example returning [`Control`](petgraph::visit::Control).
86///
87/// Find a path from vertex 0 to 5, and exit the visit as soon as we reach
88/// the goal vertex.
89///
90/// ```
91/// use rustworkx_core::petgraph::prelude::*;
92/// use rustworkx_core::petgraph::graph::node_index as n;
93/// use rustworkx_core::petgraph::visit::Control;
94///
95/// use rustworkx_core::traversal::{BfsEvent, breadth_first_search};
96///
97/// let gr: Graph<(), ()> = Graph::from_edges(&[
98/// (0, 1), (0, 2), (0, 3),
99/// (1, 3),
100/// (2, 3), (2, 4),
101/// (4, 0), (4, 5),
102/// ]);
103///
104/// // record each predecessor, mapping node → node
105/// let mut predecessor = vec![NodeIndex::end(); gr.node_count()];
106/// let start = n(0);
107/// let goal = n(5);
108/// breadth_first_search(&gr, Some(start), |event| {
109/// if let BfsEvent::TreeEdge(u, v, _) = event {
110/// predecessor[v.index()] = u;
111/// if v == goal {
112/// return Control::Break(v);
113/// }
114/// }
115/// Control::Continue
116/// });
117///
118/// let mut next = goal;
119/// let mut path = vec![next];
120/// while next != start {
121/// let pred = predecessor[next.index()];
122/// path.push(pred);
123/// next = pred;
124/// }
125/// path.reverse();
126/// assert_eq!(&path, &[n(0), n(2), n(4), n(5)]);
127/// ```
128///
129/// # Example returning a `Result`.
130/// ```
131/// use rustworkx_core::petgraph::graph::node_index as n;
132/// use rustworkx_core::petgraph::prelude::*;
133///
134/// use rustworkx_core::traversal::{BfsEvent, breadth_first_search};
135///
136/// let gr: Graph<(), ()> = Graph::from_edges(&[(0, 1), (1, 2), (1, 1), (2, 1)]);
137/// let start = n(0);
138/// let mut non_tree_edges = 0;
139///
140/// #[derive(Debug)]
141/// struct NonTreeEdgeFound {
142/// source: NodeIndex,
143/// target: NodeIndex,
144/// }
145///
146/// // Stop the search, the first time a BackEdge is encountered.
147/// let result = breadth_first_search(&gr, Some(start), |event| {
148/// match event {
149/// BfsEvent::NonTreeEdge(u, v, _) => {
150/// non_tree_edges += 1;
151/// // the implementation of ControlFlow for Result,
152/// // treats this Err value as Continue::Break
153/// Err(NonTreeEdgeFound {source: u, target: v})
154/// }
155/// // In the cases where Ok(()) is returned,
156/// // Result falls back to the implementation of Control on the value ().
157/// // In the case of (), this is to always return Control::Continue.
158/// // continuing the search.
159/// _ => Ok(()),
160/// }
161/// });
162///
163/// assert_eq!(non_tree_edges, 1);
164/// println!("number of non-tree edges encountered: {}", non_tree_edges);
165/// println!("non-tree edge: ({:?})", result.unwrap_err());
166/// ```
167pub fn breadth_first_search<G, I, F, C>(graph: G, starts: I, mut visitor: F) -> C
168where
169 G: IntoEdges + Visitable,
170 I: IntoIterator<Item = G::NodeId>,
171 F: FnMut(BfsEvent<G::NodeId, &G::EdgeWeight>) -> C,
172 C: ControlFlow,
173{
174 let discovered = &mut graph.visit_map();
175 let finished = &mut graph.visit_map();
176
177 for start in starts {
178 // `bfs_visitor` returns a "signal" to either continue or exit early
179 // but it never "prunes", so we use `unreachable`.
180 try_control!(
181 bfs_visitor(graph, start, &mut visitor, discovered, finished),
182 unreachable!()
183 );
184 }
185 C::continuing()
186}
187
188fn bfs_visitor<G, F, C>(
189 graph: G,
190 u: G::NodeId,
191 visitor: &mut F,
192 discovered: &mut G::Map,
193 finished: &mut G::Map,
194) -> C
195where
196 G: IntoEdges + Visitable,
197 F: FnMut(BfsEvent<G::NodeId, &G::EdgeWeight>) -> C,
198 C: ControlFlow,
199{
200 if !discovered.visit(u) {
201 return C::continuing();
202 }
203
204 try_control!(visitor(BfsEvent::Discover(u)), {}, {
205 let mut stack: VecDeque<G::NodeId> = VecDeque::new();
206 stack.push_front(u);
207
208 while let Some(u) = stack.pop_front() {
209 for edge in graph.edges(u) {
210 let v = edge.target();
211 if !discovered.is_visited(&v) {
212 try_control!(visitor(BfsEvent::TreeEdge(u, v, edge.weight())), continue);
213 discovered.visit(v);
214 try_control!(visitor(BfsEvent::Discover(v)), continue);
215 stack.push_back(v);
216 } else {
217 // non - tree edge.
218 try_control!(
219 visitor(BfsEvent::NonTreeEdge(u, v, edge.weight())),
220 continue
221 );
222
223 if !finished.is_visited(&v) {
224 try_control!(
225 visitor(BfsEvent::GrayTargetEdge(u, v, edge.weight())),
226 continue
227 );
228 } else {
229 try_control!(
230 visitor(BfsEvent::BlackTargetEdge(u, v, edge.weight())),
231 continue
232 );
233 }
234 }
235 }
236
237 let first_finish = finished.visit(u);
238 debug_assert!(first_finish);
239 try_control!(
240 visitor(BfsEvent::Finish(u)),
241 panic!("Pruning on the `BfsEvent::Finish` is not supported!")
242 );
243 }
244 });
245
246 C::continuing()
247}