Skip to main content

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