Skip to content

Commit

Permalink
Move bfs_successors and bfs_predecessors to rustworkx-core (Qiskit#1209)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
mtreinish authored Jun 2, 2024
1 parent 4327583 commit 53868f8
Show file tree
Hide file tree
Showing 3 changed files with 215 additions and 33 deletions.
Original file line number Diff line number Diff line change
@@ -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.
171 changes: 171 additions & 0 deletions rustworkx-core/src/traversal/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down
65 changes: 32 additions & 33 deletions src/traversal/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};

Expand All @@ -34,7 +35,6 @@ use pyo3::prelude::*;
use pyo3::Python;

use petgraph::graph::NodeIndex;
use petgraph::visit::{Bfs, NodeCount, Reversed};

use crate::iterators::EdgeList;

Expand Down Expand Up @@ -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,
}
Expand All @@ -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,
}
Expand Down

0 comments on commit 53868f8

Please sign in to comment.