Move bfs_successors and bfs_predecessors to rustworkx-core

releasenotes/notes/bfs-pred-succ-rustworkx-core-480d25a374294899.yaml

@@ -0,0 +1,12 @@
+---
+features:
+  - |
+    Added a new function ``bfs_predecessors()`` to the
+    ``rustworkx_core::traversal`` module. That is a generic Rust implementation
+    for the core rust library that provides the
+    :func:`.bfs_predecessors` function to Rust users.
+  - |
+    Added a new function ``bfs_successors()`` to the
+    ``rustworkx_core::traversal`` module. That is a generic Rust implementation
+    for the core rust library that provides the
+    :func:`.bfs_successors` function to Rust users.

rustworkx-core/src/traversal/mod.rs

@@ -140,6 +140,177 @@ where
     }
 }
 
+struct BFSAncestryWalker<G, N, VM> {
+    graph: G,
+    walker: Bfs<N, VM>,
+}
+
+impl<
+        G: GraphRef + Visitable + IntoNeighborsDirected<NodeId = N>,
+        N: Copy + Clone + PartialEq,
+        VM: VisitMap<N>,
+    > Iterator for BFSAncestryWalker<G, N, VM>
+{
+    type Item = (N, Vec<N>);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.walker.next(self.graph).map(|nx| {
+            (
+                nx,
+                self.graph
+                    .neighbors_directed(nx, petgraph::Direction::Outgoing)
+                    .collect(),
+            )
+        })
+    }
+}
+
+/// Return the successors in a breadth-first-search from a source node
+///
+/// Each iterator step returns the node indices in a bfs order from
+/// the specified node in the form:
+///
+/// `(Parent Node, vec![children nodes])`
+///
+/// # Arguments:
+///
+/// * `graph` - The graph to search
+/// * `node` - The node to search from
+///
+/// # Returns
+///
+/// An iterator of nodes in BFS order where each item in the iterator
+/// is a tuple of the `NodeId` for the parent and a `Vec` of node ids
+/// it's successors. If a node in the bfs traversal doesn't have any
+/// successors it will still be present but contain an empty vec.
+///
+/// # Example
+///
+/// ```rust
+/// use rustworkx_core::traversal::bfs_successors;
+/// use rustworkx_core::petgraph::stable_graph::{StableDiGraph, NodeIndex};
+///
+/// let graph: StableDiGraph<(), ()> = StableDiGraph::from_edges(&[
+///     (0, 1), (1, 2), (1, 3), (2, 4), (3, 4), (4, 5)
+/// ]);
+/// let successors: Vec<(usize, Vec<usize>)> = bfs_successors(&graph, NodeIndex::new(3))
+///     .map(|(x, succ)| (x.index(), succ.iter().map(|y| y.index()).collect()))
+///     .collect();
+/// assert_eq!(vec![(3_usize, vec![4_usize]), (4, vec![5]), (5, vec![])], successors);
+/// ```
+pub fn bfs_successors<G>(
+    graph: G,
+    node: G::NodeId,
+) -> impl Iterator<Item = (G::NodeId, Vec<G::NodeId>)>
+where
+    G: GraphRef + Visitable + IntoNeighborsDirected,
+{
+    BFSAncestryWalker {
+        graph,
+        walker: Bfs::new(graph, node),
+    }
+}
+
+/// Return the predecessor in a breadth-first-search from a source node
+///
+/// Each iterator step returns the node indices in a bfs order from
+/// the specified node in the form:
+///
+/// `(Child Node, vec![parent nodes])`
+///
+/// # Arguments:
+///
+/// * `graph` - The graph to search
+/// * `node` - The node to search from
+///
+/// # Returns
+///
+/// An iterator of nodes in BFS order where each item in the iterator
+/// is a tuple of the `NodeId` for the child and a `Vec` of node ids
+/// it's predecessors. If a node in the bfs traversal doesn't have any
+/// predecessors it will still be present but contain an empty vec.
+///
+/// # Example
+///
+/// ```rust
+/// use rustworkx_core::traversal::bfs_predecessors;
+/// use rustworkx_core::petgraph::stable_graph::{StableDiGraph, NodeIndex};
+///
+/// let graph: StableDiGraph<(), ()> = StableDiGraph::from_edges(&[
+///     (0, 1), (1, 2), (1, 3), (2, 4), (3, 4), (4, 5)
+/// ]);
+/// let predecessors: Vec<(usize, Vec<usize>)> = bfs_predecessors(&graph, NodeIndex::new(3))
+///     .map(|(x, succ)| (x.index(), succ.iter().map(|y| y.index()).collect()))
+///     .collect();
+/// assert_eq!(vec![(3_usize, vec![1_usize]), (1, vec![0]), (0, vec![])], predecessors);
+/// ```
+pub fn bfs_predecessors<G>(
+    graph: G,
+    node: G::NodeId,
+) -> impl Iterator<Item = (G::NodeId, Vec<G::NodeId>)>
+where
+    G: GraphRef + Visitable + IntoNeighborsDirected,
+{
+    let reversed = Reversed(graph);
+    BFSAncestryWalker {
+        graph: reversed,
+        walker: Bfs::new(reversed, node),
+    }
+}
+
+#[cfg(test)]
+mod test_bfs_ancestry {
+    use super::{bfs_predecessors, bfs_successors};
+    use crate::petgraph::graph::DiGraph;
+    use crate::petgraph::stable_graph::{NodeIndex, StableDiGraph};
+
+    #[test]
+    fn test_bfs_predecessors_digraph() {
+        let graph: DiGraph<(), ()> =
+            DiGraph::from_edges(&[(0, 1), (1, 2), (1, 3), (2, 4), (3, 4), (4, 5)]);
+        let predecessors: Vec<(usize, Vec<usize>)> = bfs_predecessors(&graph, NodeIndex::new(3))
+            .map(|(x, succ)| (x.index(), succ.iter().map(|y| y.index()).collect()))
+            .collect();
+        assert_eq!(
+            vec![(3_usize, vec![1_usize]), (1, vec![0]), (0, vec![])],
+            predecessors
+        );
+    }
+
+    #[test]
+    fn test_bfs_successors() {
+        let graph: DiGraph<(), ()> =
+            DiGraph::from_edges(&[(0, 1), (1, 2), (1, 3), (2, 4), (3, 4), (4, 5)]);
+        let successors: Vec<(usize, Vec<usize>)> = bfs_successors(&graph, NodeIndex::new(3))
+            .map(|(x, succ)| (x.index(), succ.iter().map(|y| y.index()).collect()))
+            .collect();
+        assert_eq!(
+            vec![(3_usize, vec![4_usize]), (4, vec![5]), (5, vec![])],
+            successors
+        );
+    }
+
+    #[test]
+    fn test_no_predecessors() {
+        let graph: StableDiGraph<(), ()> =
+            StableDiGraph::from_edges(&[(0, 1), (1, 2), (1, 3), (2, 4), (3, 4), (4, 5)]);
+        let predecessors: Vec<(usize, Vec<usize>)> = bfs_predecessors(&graph, NodeIndex::new(0))
+            .map(|(x, succ)| (x.index(), succ.iter().map(|y| y.index()).collect()))
+            .collect();
+        assert_eq!(vec![(0_usize, vec![])], predecessors);
+    }
+
+    #[test]
+    fn test_no_successors() {
+        let graph: StableDiGraph<(), ()> =
+            StableDiGraph::from_edges(&[(0, 1), (1, 2), (1, 3), (2, 4), (3, 4), (4, 5)]);
+        let successors: Vec<(usize, Vec<usize>)> = bfs_successors(&graph, NodeIndex::new(5))
+            .map(|(x, succ)| (x.index(), succ.iter().map(|y| y.index()).collect()))
+            .collect();
+        assert_eq!(vec![(5_usize, vec![])], successors);
+    }
+}
+
 #[cfg(test)]
 mod test_ancestry {
     use super::{ancestors, descendants};

src/traversal/mod.rs

@@ -19,7 +19,8 @@ use dfs_visit::{dfs_handler, PyDfsVisitor};
 use dijkstra_visit::{dijkstra_handler, PyDijkstraVisitor};
 
 use rustworkx_core::traversal::{
-    ancestors as core_ancestors, breadth_first_search, depth_first_search,
+    ancestors as core_ancestors, bfs_predecessors as core_bfs_predecessors,
+    bfs_successors as core_bfs_successors, breadth_first_search, depth_first_search,
     descendants as core_descendants, dfs_edges, dijkstra_search,
 };
 
@@ -34,7 +35,6 @@ use pyo3::prelude::*;
 use pyo3::Python;
 
 use petgraph::graph::NodeIndex;
-use petgraph::visit::{Bfs, NodeCount, Reversed};
 
 use crate::iterators::EdgeList;
 
@@ -144,21 +144,21 @@ pub fn bfs_successors(
     node: usize,
 ) -> iterators::BFSSuccessors {
     let index = NodeIndex::new(node);
-    let mut bfs = Bfs::new(&graph.graph, index);
-    let mut out_list: Vec<(PyObject, Vec<PyObject>)> = Vec::with_capacity(graph.node_count());
-    while let Some(nx) = bfs.next(&graph.graph) {
-        let successors: Vec<PyObject> = graph
-            .graph
-            .neighbors_directed(nx, petgraph::Direction::Outgoing)
-            .map(|pred| graph.graph.node_weight(pred).unwrap().clone_ref(py))
-            .collect();
-        if !successors.is_empty() {
-            out_list.push((
-                graph.graph.node_weight(nx).unwrap().clone_ref(py),
-                successors,
-            ));
-        }
-    }
+    let out_list = core_bfs_successors(&graph.graph, index)
+        .filter_map(|(nx, succ_list)| {
+            if succ_list.is_empty() {
+                None
+            } else {
+                Some((
+                    graph.graph.node_weight(nx).unwrap().clone_ref(py),
+                    succ_list
+                        .into_iter()
+                        .map(|pred| graph.graph.node_weight(pred).unwrap().clone_ref(py))
+                        .collect(),
+                ))
+            }
+        })
+        .collect();
     iterators::BFSSuccessors {
         bfs_successors: out_list,
     }
@@ -185,22 +185,21 @@ pub fn bfs_predecessors(
     node: usize,
 ) -> iterators::BFSPredecessors {
     let index = NodeIndex::new(node);
-    let reverse_graph = Reversed(&graph.graph);
-    let mut bfs = Bfs::new(reverse_graph, index);
-    let mut out_list: Vec<(PyObject, Vec<PyObject>)> = Vec::with_capacity(graph.node_count());
-    while let Some(nx) = bfs.next(reverse_graph) {
-        let predecessors: Vec<PyObject> = graph
-            .graph
-            .neighbors_directed(nx, petgraph::Direction::Incoming)
-            .map(|pred| graph.graph.node_weight(pred).unwrap().clone_ref(py))
-            .collect();
-        if !predecessors.is_empty() {
-            out_list.push((
-                graph.graph.node_weight(nx).unwrap().clone_ref(py),
-                predecessors,
-            ));
-        }
-    }
+    let out_list = core_bfs_predecessors(&graph.graph, index)
+        .filter_map(|(nx, succ_list)| {
+            if succ_list.is_empty() {
+                None
+            } else {
+                Some((
+                    graph.graph.node_weight(nx).unwrap().clone_ref(py),
+                    succ_list
+                        .into_iter()
+                        .map(|pred| graph.graph.node_weight(pred).unwrap().clone_ref(py))
+                        .collect(),
+                ))
+            }
+        })
+        .collect();
     iterators::BFSPredecessors {
         bfs_predecessors: out_list,
     }