Move ancestors and descendants to rustworkx-core (#1208)

This commit adds an implementation of the ancestors and descendants
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. These new functions
will be more efficient as they're not using dijkstra's algorithm to find
a path from nodes now and instead are just doing a BFS. The
rustwork-core functions also return an iterator of nodes.

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: mtreinish

releasenotes/notes/ancestors-descendants-rustworkx-core-8c4aac3364c155bf.yaml

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

rustworkx-core/src/traversal/mod.rs

@@ -17,6 +17,13 @@ mod dfs_edges;
 mod dfs_visit;
 mod dijkstra_visit;
 
+use petgraph::prelude::*;
+use petgraph::visit::GraphRef;
+use petgraph::visit::IntoNeighborsDirected;
+use petgraph::visit::Reversed;
+use petgraph::visit::VisitMap;
+use petgraph::visit::Visitable;
+
 pub use bfs_visit::{breadth_first_search, BfsEvent};
 pub use dfs_edges::dfs_edges;
 pub use dfs_visit::{depth_first_search, DfsEvent};
@@ -45,3 +52,133 @@ macro_rules! try_control {
 }
 
 use try_control;
+
+struct AncestryWalker<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 AncestryWalker<G, N, VM>
+{
+    type Item = N;
+    fn next(&mut self) -> Option<Self::Item> {
+        self.walker.next(self.graph)
+    }
+}
+
+/// Return the ancestors of a node in a graph.
+///
+/// `node` is included in the output
+///
+/// # Arguments:
+///
+/// * `node` - The node to find the ancestors of
+///
+/// # Returns
+///
+/// An iterator where each item is a node id for an ancestor of ``node``.
+/// This includes ``node`` in the returned ids.
+///
+/// # Example
+///
+/// ```rust
+/// use rustworkx_core::traversal::ancestors;
+/// 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 ancestors: Vec<usize> = ancestors(&graph, NodeIndex::new(3)).map(|x| x.index()).collect();
+/// assert_eq!(vec![3_usize, 1, 0], ancestors);
+/// ```
+pub fn ancestors<G>(graph: G, node: G::NodeId) -> impl Iterator<Item = G::NodeId>
+where
+    G: GraphRef + Visitable + IntoNeighborsDirected,
+{
+    let reversed = Reversed(graph);
+    AncestryWalker {
+        graph: reversed,
+        walker: Bfs::new(reversed, node),
+    }
+}
+
+/// Return the descendants of a node in a graph.
+///
+/// `node` is included in the output.
+/// # Arguments:
+///
+/// * `node` - The node to find the ancestors of
+///
+/// # Returns
+///
+/// An iterator where each item is a node id for an ancestor of ``node``.
+/// This includes ``node`` in the returned ids.
+///
+/// # Example
+///
+/// ```rust
+/// use rustworkx_core::traversal::descendants;
+/// 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 descendants: Vec<usize> = descendants(&graph, NodeIndex::new(3)).map(|x| x.index()).collect();
+/// assert_eq!(vec![3_usize, 4, 5], descendants);
+/// ```
+pub fn descendants<G>(graph: G, node: G::NodeId) -> impl Iterator<Item = G::NodeId>
+where
+    G: GraphRef + Visitable + IntoNeighborsDirected,
+{
+    AncestryWalker {
+        graph,
+        walker: Bfs::new(graph, node),
+    }
+}
+
+#[cfg(test)]
+mod test_ancestry {
+    use super::{ancestors, descendants};
+    use crate::petgraph::graph::DiGraph;
+    use crate::petgraph::stable_graph::{NodeIndex, StableDiGraph};
+
+    #[test]
+    fn test_ancestors_digraph() {
+        let graph: DiGraph<(), ()> =
+            DiGraph::from_edges(&[(0, 1), (1, 2), (1, 3), (2, 4), (3, 4), (4, 5)]);
+        let ancestors: Vec<usize> = ancestors(&graph, NodeIndex::new(3))
+            .map(|x| x.index())
+            .collect();
+        assert_eq!(vec![3_usize, 1, 0], ancestors);
+    }
+
+    #[test]
+    fn test_descendants() {
+        let graph: DiGraph<(), ()> =
+            DiGraph::from_edges(&[(0, 1), (1, 2), (1, 3), (2, 4), (3, 4), (4, 5)]);
+        let descendants: Vec<usize> = descendants(&graph, NodeIndex::new(3))
+            .map(|x| x.index())
+            .collect();
+        assert_eq!(vec![3_usize, 4, 5], descendants);
+    }
+
+    #[test]
+    fn test_no_ancestors() {
+        let mut graph: StableDiGraph<(), ()> = StableDiGraph::new();
+        let index = graph.add_node(());
+        let res = ancestors(&graph, index);
+        assert_eq!(vec![index], res.collect::<Vec<NodeIndex>>())
+    }
+
+    #[test]
+    fn test_no_descendants() {
+        let mut graph: StableDiGraph<(), ()> = StableDiGraph::new();
+        let index = graph.add_node(());
+        let res = descendants(&graph, index);
+        assert_eq!(vec![index], res.collect::<Vec<NodeIndex>>())
+    }
+}

src/traversal/mod.rs

@@ -19,7 +19,8 @@ use dfs_visit::{dfs_handler, PyDfsVisitor};
 use dijkstra_visit::{dijkstra_handler, PyDijkstraVisitor};
 
 use rustworkx_core::traversal::{
-    breadth_first_search, depth_first_search, dfs_edges, dijkstra_search,
+    ancestors as core_ancestors, breadth_first_search, depth_first_search,
+    descendants as core_descendants, dfs_edges, dijkstra_search,
 };
 
 use super::{digraph, graph, iterators, CostFn};
@@ -32,7 +33,6 @@ use pyo3::exceptions::PyTypeError;
 use pyo3::prelude::*;
 use pyo3::Python;
 
-use petgraph::algo;
 use petgraph::graph::NodeIndex;
 use petgraph::visit::{Bfs, NodeCount, Reversed};
 
@@ -221,16 +221,10 @@ pub fn bfs_predecessors(
 #[pyfunction]
 #[pyo3(text_signature = "(graph, node, /)")]
 pub fn ancestors(graph: &digraph::PyDiGraph, node: usize) -> HashSet<usize> {
-    let index = NodeIndex::new(node);
-    let mut out_set: HashSet<usize> = HashSet::new();
-    let reverse_graph = Reversed(&graph.graph);
-    let res = algo::dijkstra(reverse_graph, index, None, |_| 1);
-    for n in res.keys() {
-        let n_int = n.index();
-        out_set.insert(n_int);
-    }
-    out_set.remove(&node);
-    out_set
+    core_ancestors(&graph.graph, NodeIndex::new(node))
+        .map(|x| x.index())
+        .filter(|x| *x != node)
+        .collect()
 }
 
 /// Return the descendants of a node in a graph.
@@ -249,14 +243,10 @@ pub fn ancestors(graph: &digraph::PyDiGraph, node: usize) -> HashSet<usize> {
 #[pyo3(text_signature = "(graph, node, /)")]
 pub fn descendants(graph: &digraph::PyDiGraph, node: usize) -> HashSet<usize> {
     let index = NodeIndex::new(node);
-    let mut out_set: HashSet<usize> = HashSet::new();
-    let res = algo::dijkstra(&graph.graph, index, None, |_| 1);
-    for n in res.keys() {
-        let n_int = n.index();
-        out_set.insert(n_int);
-    }
-    out_set.remove(&node);
-    out_set
+    core_descendants(&graph.graph, index)
+        .map(|x| x.index())
+        .filter(|x| *x != node)
+        .collect()
 }
 
 /// Breadth-first traversal of a directed graph.