rustworkx_core/traversal/
mod.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
13//! Module for graph traversal algorithms.
14
15mod bfs_visit;
16mod dfs_edges;
17mod dfs_visit;
18mod dijkstra_visit;
19
20use petgraph::prelude::*;
21use petgraph::visit::GraphRef;
22use petgraph::visit::IntoNeighborsDirected;
23use petgraph::visit::Reversed;
24use petgraph::visit::VisitMap;
25use petgraph::visit::Visitable;
26
27pub use bfs_visit::{breadth_first_search, BfsEvent};
28pub use dfs_edges::dfs_edges;
29pub use dfs_visit::{depth_first_search, DfsEvent};
30pub use dijkstra_visit::{dijkstra_search, DijkstraEvent};
31
32/// Return if the expression is a break value, execute the provided statement
33/// if it is a prune value.
34/// https://github.com/petgraph/petgraph/blob/0.6.0/src/visit/dfsvisit.rs#L27
35macro_rules! try_control {
36    ($e:expr, $p:stmt) => {
37        try_control!($e, $p, ());
38    };
39    ($e:expr, $p:stmt, $q:stmt) => {
40        match $e {
41            x => {
42                if x.should_break() {
43                    return x;
44                } else if x.should_prune() {
45                    $p
46                } else {
47                    $q
48                }
49            }
50        }
51    };
52}
53
54use try_control;
55
56struct AncestryWalker<G, N, VM> {
57    graph: G,
58    walker: Bfs<N, VM>,
59}
60
61impl<
62        G: GraphRef + Visitable + IntoNeighborsDirected<NodeId = N>,
63        N: Copy + Clone + PartialEq,
64        VM: VisitMap<N>,
65    > Iterator for AncestryWalker<G, N, VM>
66{
67    type Item = N;
68    fn next(&mut self) -> Option<Self::Item> {
69        self.walker.next(self.graph)
70    }
71}
72
73/// Return the ancestors of a node in a graph.
74///
75/// `node` is included in the output
76///
77/// # Arguments:
78///
79/// * `node` - The node to find the ancestors of
80///
81/// # Returns
82///
83/// An iterator where each item is a node id for an ancestor of ``node``.
84/// This includes ``node`` in the returned ids.
85///
86/// # Example
87///
88/// ```rust
89/// use rustworkx_core::traversal::ancestors;
90/// use rustworkx_core::petgraph::stable_graph::{StableDiGraph, NodeIndex};
91///
92/// let graph: StableDiGraph<(), ()> = StableDiGraph::from_edges(&[
93///     (0, 1), (1, 2), (1, 3), (2, 4), (3, 4), (4, 5)
94/// ]);
95/// let ancestors: Vec<usize> = ancestors(&graph, NodeIndex::new(3)).map(|x| x.index()).collect();
96/// assert_eq!(vec![3_usize, 1, 0], ancestors);
97/// ```
98pub fn ancestors<G>(graph: G, node: G::NodeId) -> impl Iterator<Item = G::NodeId>
99where
100    G: GraphRef + Visitable + IntoNeighborsDirected,
101{
102    let reversed = Reversed(graph);
103    AncestryWalker {
104        graph: reversed,
105        walker: Bfs::new(reversed, node),
106    }
107}
108
109/// Return the descendants of a node in a graph.
110///
111/// `node` is included in the output.
112/// # Arguments:
113///
114/// * `node` - The node to find the ancestors of
115///
116/// # Returns
117///
118/// An iterator where each item is a node id for an ancestor of ``node``.
119/// This includes ``node`` in the returned ids.
120///
121/// # Example
122///
123/// ```rust
124/// use rustworkx_core::traversal::descendants;
125/// use rustworkx_core::petgraph::stable_graph::{StableDiGraph, NodeIndex};
126///
127/// let graph: StableDiGraph<(), ()> = StableDiGraph::from_edges(&[
128///     (0, 1), (1, 2), (1, 3), (2, 4), (3, 4), (4, 5)
129/// ]);
130/// let descendants: Vec<usize> = descendants(&graph, NodeIndex::new(3)).map(|x| x.index()).collect();
131/// assert_eq!(vec![3_usize, 4, 5], descendants);
132/// ```
133pub fn descendants<G>(graph: G, node: G::NodeId) -> impl Iterator<Item = G::NodeId>
134where
135    G: GraphRef + Visitable + IntoNeighborsDirected,
136{
137    AncestryWalker {
138        graph,
139        walker: Bfs::new(graph, node),
140    }
141}
142
143struct BFSAncestryWalker<G, N, VM> {
144    graph: G,
145    walker: Bfs<N, VM>,
146}
147
148impl<
149        G: GraphRef + Visitable + IntoNeighborsDirected<NodeId = N>,
150        N: Copy + Clone + PartialEq,
151        VM: VisitMap<N>,
152    > Iterator for BFSAncestryWalker<G, N, VM>
153{
154    type Item = (N, Vec<N>);
155
156    fn next(&mut self) -> Option<Self::Item> {
157        self.walker.next(self.graph).map(|nx| {
158            (
159                nx,
160                self.graph
161                    .neighbors_directed(nx, petgraph::Direction::Outgoing)
162                    .collect(),
163            )
164        })
165    }
166}
167
168/// Return the successors in a breadth-first-search from a source node
169///
170/// Each iterator step returns the node indices in a bfs order from
171/// the specified node in the form:
172///
173/// `(Parent Node, vec![children nodes])`
174///
175/// # Arguments:
176///
177/// * `graph` - The graph to search
178/// * `node` - The node to search from
179///
180/// # Returns
181///
182/// An iterator of nodes in BFS order where each item in the iterator
183/// is a tuple of the `NodeId` for the parent and a `Vec` of node ids
184/// it's successors. If a node in the bfs traversal doesn't have any
185/// successors it will still be present but contain an empty vec.
186///
187/// # Example
188///
189/// ```rust
190/// use rustworkx_core::traversal::bfs_successors;
191/// use rustworkx_core::petgraph::stable_graph::{StableDiGraph, NodeIndex};
192///
193/// let graph: StableDiGraph<(), ()> = StableDiGraph::from_edges(&[
194///     (0, 1), (1, 2), (1, 3), (2, 4), (3, 4), (4, 5)
195/// ]);
196/// let successors: Vec<(usize, Vec<usize>)> = bfs_successors(&graph, NodeIndex::new(3))
197///     .map(|(x, succ)| (x.index(), succ.iter().map(|y| y.index()).collect()))
198///     .collect();
199/// assert_eq!(vec![(3_usize, vec![4_usize]), (4, vec![5]), (5, vec![])], successors);
200/// ```
201pub fn bfs_successors<G>(
202    graph: G,
203    node: G::NodeId,
204) -> impl Iterator<Item = (G::NodeId, Vec<G::NodeId>)>
205where
206    G: GraphRef + Visitable + IntoNeighborsDirected,
207{
208    BFSAncestryWalker {
209        graph,
210        walker: Bfs::new(graph, node),
211    }
212}
213
214/// Return the predecessor in a breadth-first-search from a source node
215///
216/// Each iterator step returns the node indices in a bfs order from
217/// the specified node in the form:
218///
219/// `(Child Node, vec![parent nodes])`
220///
221/// # Arguments:
222///
223/// * `graph` - The graph to search
224/// * `node` - The node to search from
225///
226/// # Returns
227///
228/// An iterator of nodes in BFS order where each item in the iterator
229/// is a tuple of the `NodeId` for the child and a `Vec` of node ids
230/// it's predecessors. If a node in the bfs traversal doesn't have any
231/// predecessors it will still be present but contain an empty vec.
232///
233/// # Example
234///
235/// ```rust
236/// use rustworkx_core::traversal::bfs_predecessors;
237/// use rustworkx_core::petgraph::stable_graph::{StableDiGraph, NodeIndex};
238///
239/// let graph: StableDiGraph<(), ()> = StableDiGraph::from_edges(&[
240///     (0, 1), (1, 2), (1, 3), (2, 4), (3, 4), (4, 5)
241/// ]);
242/// let predecessors: Vec<(usize, Vec<usize>)> = bfs_predecessors(&graph, NodeIndex::new(3))
243///     .map(|(x, succ)| (x.index(), succ.iter().map(|y| y.index()).collect()))
244///     .collect();
245/// assert_eq!(vec![(3_usize, vec![1_usize]), (1, vec![0]), (0, vec![])], predecessors);
246/// ```
247pub fn bfs_predecessors<G>(
248    graph: G,
249    node: G::NodeId,
250) -> impl Iterator<Item = (G::NodeId, Vec<G::NodeId>)>
251where
252    G: GraphRef + Visitable + IntoNeighborsDirected,
253{
254    let reversed = Reversed(graph);
255    BFSAncestryWalker {
256        graph: reversed,
257        walker: Bfs::new(reversed, node),
258    }
259}
260
261#[cfg(test)]
262mod test_bfs_ancestry {
263    use super::{bfs_predecessors, bfs_successors};
264    use crate::petgraph::graph::DiGraph;
265    use crate::petgraph::stable_graph::{NodeIndex, StableDiGraph};
266
267    #[test]
268    fn test_bfs_predecessors_digraph() {
269        let graph: DiGraph<(), ()> =
270            DiGraph::from_edges(&[(0, 1), (1, 2), (1, 3), (2, 4), (3, 4), (4, 5)]);
271        let predecessors: Vec<(usize, Vec<usize>)> = bfs_predecessors(&graph, NodeIndex::new(3))
272            .map(|(x, succ)| (x.index(), succ.iter().map(|y| y.index()).collect()))
273            .collect();
274        assert_eq!(
275            vec![(3_usize, vec![1_usize]), (1, vec![0]), (0, vec![])],
276            predecessors
277        );
278    }
279
280    #[test]
281    fn test_bfs_successors() {
282        let graph: DiGraph<(), ()> =
283            DiGraph::from_edges(&[(0, 1), (1, 2), (1, 3), (2, 4), (3, 4), (4, 5)]);
284        let successors: Vec<(usize, Vec<usize>)> = bfs_successors(&graph, NodeIndex::new(3))
285            .map(|(x, succ)| (x.index(), succ.iter().map(|y| y.index()).collect()))
286            .collect();
287        assert_eq!(
288            vec![(3_usize, vec![4_usize]), (4, vec![5]), (5, vec![])],
289            successors
290        );
291    }
292
293    #[test]
294    fn test_no_predecessors() {
295        let graph: StableDiGraph<(), ()> =
296            StableDiGraph::from_edges(&[(0, 1), (1, 2), (1, 3), (2, 4), (3, 4), (4, 5)]);
297        let predecessors: Vec<(usize, Vec<usize>)> = bfs_predecessors(&graph, NodeIndex::new(0))
298            .map(|(x, succ)| (x.index(), succ.iter().map(|y| y.index()).collect()))
299            .collect();
300        assert_eq!(vec![(0_usize, vec![])], predecessors);
301    }
302
303    #[test]
304    fn test_no_successors() {
305        let graph: StableDiGraph<(), ()> =
306            StableDiGraph::from_edges(&[(0, 1), (1, 2), (1, 3), (2, 4), (3, 4), (4, 5)]);
307        let successors: Vec<(usize, Vec<usize>)> = bfs_successors(&graph, NodeIndex::new(5))
308            .map(|(x, succ)| (x.index(), succ.iter().map(|y| y.index()).collect()))
309            .collect();
310        assert_eq!(vec![(5_usize, vec![])], successors);
311    }
312}
313
314#[cfg(test)]
315mod test_ancestry {
316    use super::{ancestors, descendants};
317    use crate::petgraph::graph::DiGraph;
318    use crate::petgraph::stable_graph::{NodeIndex, StableDiGraph};
319
320    #[test]
321    fn test_ancestors_digraph() {
322        let graph: DiGraph<(), ()> =
323            DiGraph::from_edges(&[(0, 1), (1, 2), (1, 3), (2, 4), (3, 4), (4, 5)]);
324        let ancestors: Vec<usize> = ancestors(&graph, NodeIndex::new(3))
325            .map(|x| x.index())
326            .collect();
327        assert_eq!(vec![3_usize, 1, 0], ancestors);
328    }
329
330    #[test]
331    fn test_descendants() {
332        let graph: DiGraph<(), ()> =
333            DiGraph::from_edges(&[(0, 1), (1, 2), (1, 3), (2, 4), (3, 4), (4, 5)]);
334        let descendants: Vec<usize> = descendants(&graph, NodeIndex::new(3))
335            .map(|x| x.index())
336            .collect();
337        assert_eq!(vec![3_usize, 4, 5], descendants);
338    }
339
340    #[test]
341    fn test_no_ancestors() {
342        let mut graph: StableDiGraph<(), ()> = StableDiGraph::new();
343        let index = graph.add_node(());
344        let res = ancestors(&graph, index);
345        assert_eq!(vec![index], res.collect::<Vec<NodeIndex>>())
346    }
347
348    #[test]
349    fn test_no_descendants() {
350        let mut graph: StableDiGraph<(), ()> = StableDiGraph::new();
351        let index = graph.add_node(());
352        let res = descendants(&graph, index);
353        assert_eq!(vec![index], res.collect::<Vec<NodeIndex>>())
354    }
355}