From 53868f83c2c100dbd8b4c5f689f075699f2edab5 Mon Sep 17 00:00:00 2001 From: Matthew Treinish Date: Sun, 2 Jun 2024 16:36:22 -0400 Subject: [PATCH] Move bfs_successors and bfs_predecessors to rustworkx-core (#1209) This commit adds an implementation of the bfs_successors and bfs_predecessors functions to the rustworkx-core crate exposing the functions to rust users. The existing implementation in the rustworkx crate is removed and it is updated to call the rustworkx-core functions. The rustwork-core functions return an iterator of nodes indices and the Python crate side consumes the iterator to build the bfs view these functions previously returned. --- ...-succ-rustworkx-core-480d25a374294899.yaml | 12 ++ rustworkx-core/src/traversal/mod.rs | 171 ++++++++++++++++++ src/traversal/mod.rs | 65 ++++--- 3 files changed, 215 insertions(+), 33 deletions(-) create mode 100644 releasenotes/notes/bfs-pred-succ-rustworkx-core-480d25a374294899.yaml diff --git a/releasenotes/notes/bfs-pred-succ-rustworkx-core-480d25a374294899.yaml b/releasenotes/notes/bfs-pred-succ-rustworkx-core-480d25a374294899.yaml new file mode 100644 index 000000000..6de889660 --- /dev/null +++ b/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. diff --git a/rustworkx-core/src/traversal/mod.rs b/rustworkx-core/src/traversal/mod.rs index 9e0261acb..c8213dfb3 100644 --- a/rustworkx-core/src/traversal/mod.rs +++ b/rustworkx-core/src/traversal/mod.rs @@ -140,6 +140,177 @@ where } } +struct BFSAncestryWalker { + graph: G, + walker: Bfs, +} + +impl< + G: GraphRef + Visitable + IntoNeighborsDirected, + N: Copy + Clone + PartialEq, + VM: VisitMap, + > Iterator for BFSAncestryWalker +{ + type Item = (N, Vec); + + fn next(&mut self) -> Option { + 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)> = 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( + graph: G, + node: G::NodeId, +) -> impl Iterator)> +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)> = 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( + graph: G, + node: G::NodeId, +) -> impl Iterator)> +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)> = 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)> = 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)> = 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)> = 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}; diff --git a/src/traversal/mod.rs b/src/traversal/mod.rs index aee5b27b2..f6ce66a76 100644 --- a/src/traversal/mod.rs +++ b/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)> = Vec::with_capacity(graph.node_count()); - while let Some(nx) = bfs.next(&graph.graph) { - let successors: Vec = 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)> = Vec::with_capacity(graph.node_count()); - while let Some(nx) = bfs.next(reverse_graph) { - let predecessors: Vec = 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, }