Add Johnson's algorithm for cycles in directed graph (#633)

* Add Johnson's algorithm for cycles in directed graph

This commit adds a new function, simple_cycles(),  which is an
implementation of Johnson's algorithm [1] for finding all the simple
cycles in a directed graph. The implementation is based on the
non-recursive implementation from NetworkX. [2]

[1] https://doi.org/10.1137/0204007
[2] https://github.com/networkx/networkx/blob/networkx-2.8.4/networkx/algorithms/cycles.py#L98-L222

Closes: #622

* Remove test debugging

* Fix lint

* Add more tests

* Apply suggestions from code review

Co-authored-by: georgios-ts <45130028+georgios-ts@users.noreply.github.com>

* Fix rustfmt

* Make return from simple_cycles a iterator

This commit converts the simple_cycles function to a Python class that
implements the iterator protocol. Each step of the the iterator will
compute the next cycle using Johnson's algorithm and return a
NodeIndices object representing the cycle.

* Use IndexSet::pop() instead of set_pop()

* Fix iterator behavior

* Fix MSRV support

* Use mem::take() instead of .drain(..).collect()

Co-authored-by: georgios-ts <45130028+georgios-ts@users.noreply.github.com>
Co-authored-by: Ivan Carvalho <8753214+IvanIsCoding@users.noreply.github.com>

docs/source/api.rst

@@ -149,6 +149,7 @@ Connectivity and Cycles
    rustworkx.weakly_connected_components
    rustworkx.is_weakly_connected
    rustworkx.cycle_basis
+   rustworkx.simple_cycles
    rustworkx.digraph_find_cycle
    rustworkx.articulation_points
    rustworkx.biconnected_components

releasenotes/notes/simple_cycles-b2d64c2d02a1974a.yaml

@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    Added a new function, :func:`~.simple_cycles`, which is an implementation of
+    `Johnson's algorithm <https://doi.org/10.1137/0204007>`__ for finding all
+    elementary cycles in a directed graph.

src/connectivity/johnson_simple_cycles.rs

@@ -0,0 +1,309 @@
+// Licensed under the apache license, version 2.0 (the "license"); you may
+// not use this file except in compliance with the License. You may obtain
+// a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+// License for the specific language governing permissions and limitations
+// under the License.
+
+use hashbrown::{HashMap, HashSet};
+use indexmap::IndexSet;
+
+use crate::digraph::PyDiGraph;
+use crate::StablePyGraph;
+use petgraph::algo::kosaraju_scc;
+use petgraph::graph::NodeIndex;
+use petgraph::stable_graph::StableDiGraph;
+use petgraph::visit::EdgeRef;
+use petgraph::visit::IntoEdgeReferences;
+use petgraph::visit::IntoNodeReferences;
+use petgraph::visit::NodeFiltered;
+use petgraph::Directed;
+
+use pyo3::iter::IterNextOutput;
+use pyo3::prelude::*;
+
+use crate::iterators::NodeIndices;
+
+fn build_subgraph(
+    graph: &StablePyGraph<Directed>,
+    nodes: &[NodeIndex],
+) -> (StableDiGraph<(), ()>, HashMap<NodeIndex, NodeIndex>) {
+    let node_set: HashSet<NodeIndex> = nodes.iter().copied().collect();
+    let mut node_map: HashMap<NodeIndex, NodeIndex> = HashMap::with_capacity(nodes.len());
+    let node_filter = |node: NodeIndex| -> bool { node_set.contains(&node) };
+    // Overallocates edges, but not a big deal as this is temporary for the lifetime of the
+    // subgraph
+    let mut out_graph = StableDiGraph::<(), ()>::with_capacity(nodes.len(), graph.edge_count());
+    let filtered = NodeFiltered(&graph, node_filter);
+    for node in filtered.node_references() {
+        let new_node = out_graph.add_node(());
+        node_map.insert(node.0, new_node);
+    }
+    for edge in filtered.edge_references() {
+        let new_source = *node_map.get(&edge.source()).unwrap();
+        let new_target = *node_map.get(&edge.target()).unwrap();
+        out_graph.add_edge(new_source, new_target, ());
+    }
+    (out_graph, node_map)
+}
+
+#[pyclass(module = "rustworkx")]
+pub struct SimpleCycleIter {
+    graph_clone: StablePyGraph<Directed>,
+    scc: Vec<Vec<NodeIndex>>,
+    self_cycles: Option<Vec<NodeIndex>>,
+    path: Vec<NodeIndex>,
+    blocked: HashSet<NodeIndex>,
+    closed: HashSet<NodeIndex>,
+    block: HashMap<NodeIndex, HashSet<NodeIndex>>,
+    stack: Vec<(NodeIndex, IndexSet<NodeIndex>)>,
+    start_node: NodeIndex,
+    node_map: HashMap<NodeIndex, NodeIndex>,
+    reverse_node_map: HashMap<NodeIndex, NodeIndex>,
+    subgraph: StableDiGraph<(), ()>,
+}
+
+impl SimpleCycleIter {
+    pub fn new(graph: &PyDiGraph) -> Self {
+        // Copy graph to remove self edges before running johnson's algorithm
+        let mut graph_clone = graph.graph.clone();
+
+        // For compatibility with networkx manually insert self cycles and filter
+        // from Johnson's algorithm
+        let self_cycles_vec: Vec<NodeIndex> = graph_clone
+            .node_indices()
+            .filter(|n| graph_clone.neighbors(*n).any(|x| x == *n))
+            .collect();
+        for node in &self_cycles_vec {
+            while let Some(edge_index) = graph_clone.find_edge(*node, *node) {
+                graph_clone.remove_edge(edge_index);
+            }
+        }
+        let self_cycles = if self_cycles_vec.is_empty() {
+            None
+        } else {
+            Some(self_cycles_vec)
+        };
+        let strongly_connected_components: Vec<Vec<NodeIndex>> = kosaraju_scc(&graph_clone)
+            .into_iter()
+            .filter(|component| component.len() > 1)
+            .collect();
+        SimpleCycleIter {
+            graph_clone,
+            scc: strongly_connected_components,
+            self_cycles,
+            path: Vec::new(),
+            blocked: HashSet::new(),
+            closed: HashSet::new(),
+            block: HashMap::new(),
+            stack: Vec::new(),
+            start_node: NodeIndex::new(std::u32::MAX as usize),
+            node_map: HashMap::new(),
+            reverse_node_map: HashMap::new(),
+            subgraph: StableDiGraph::new(),
+        }
+    }
+}
+
+fn unblock(
+    node: NodeIndex,
+    blocked: &mut HashSet<NodeIndex>,
+    block: &mut HashMap<NodeIndex, HashSet<NodeIndex>>,
+) {
+    let mut stack: IndexSet<NodeIndex> = IndexSet::new();
+    stack.insert(node);
+    while let Some(stack_node) = stack.pop() {
+        if blocked.remove(&stack_node) {
+            match block.get_mut(&stack_node) {
+                // stack.update(block[stack_node]):
+                Some(block_set) => {
+                    block_set.drain().for_each(|n| {
+                        stack.insert(n);
+                    });
+                }
+                // If block doesn't have stack_node treat it as an empty set
+                // (so no updates to stack) and populate it with an empty
+                // set.
+                None => {
+                    block.insert(stack_node, HashSet::new());
+                }
+            }
+            blocked.remove(&stack_node);
+        }
+    }
+}
+
+#[allow(clippy::too_many_arguments)]
+fn process_stack(
+    start_node: NodeIndex,
+    stack: &mut Vec<(NodeIndex, IndexSet<NodeIndex>)>,
+    path: &mut Vec<NodeIndex>,
+    closed: &mut HashSet<NodeIndex>,
+    blocked: &mut HashSet<NodeIndex>,
+    block: &mut HashMap<NodeIndex, HashSet<NodeIndex>>,
+    subgraph: &StableDiGraph<(), ()>,
+    reverse_node_map: &HashMap<NodeIndex, NodeIndex>,
+) -> Option<IterNextOutput<NodeIndices, &'static str>> {
+    while let Some((this_node, neighbors)) = stack.last_mut() {
+        if let Some(next_node) = neighbors.pop() {
+            if next_node == start_node {
+                // Out path in input graph basis
+                let mut out_path: Vec<usize> = Vec::with_capacity(path.len());
+                for n in path {
+                    out_path.push(reverse_node_map[n].index());
+                    closed.insert(*n);
+                }
+                return Some(IterNextOutput::Yield(NodeIndices { nodes: out_path }));
+            } else if blocked.insert(next_node) {
+                path.push(next_node);
+                stack.push((
+                    next_node,
+                    subgraph
+                        .neighbors(next_node)
+                        .collect::<IndexSet<NodeIndex>>(),
+                ));
+                closed.remove(&next_node);
+                blocked.insert(next_node);
+                continue;
+            }
+        }
+        if neighbors.is_empty() {
+            if closed.contains(this_node) {
+                unblock(*this_node, blocked, block);
+            } else {
+                for neighbor in subgraph.neighbors(*this_node) {
+                    let block_neighbor = block.entry(neighbor).or_insert_with(HashSet::new);
+                    block_neighbor.insert(*this_node);
+                }
+            }
+            stack.pop();
+            path.pop();
+        }
+    }
+    None
+}
+
+#[pymethods]
+impl SimpleCycleIter {
+    fn __iter__(slf: PyRef<Self>) -> Py<SimpleCycleIter> {
+        slf.into()
+    }
+
+    fn __next__(mut slf: PyRefMut<Self>) -> PyResult<IterNextOutput<NodeIndices, &'static str>> {
+        if slf.self_cycles.is_some() {
+            let self_cycles = slf.self_cycles.as_mut().unwrap();
+            let cycle_node = self_cycles.pop().unwrap();
+            if self_cycles.is_empty() {
+                slf.self_cycles = None;
+            }
+            return Ok(IterNextOutput::Yield(NodeIndices {
+                nodes: vec![cycle_node.index()],
+            }));
+        }
+        // Restore previous state if it exists
+        let mut stack: Vec<(NodeIndex, IndexSet<NodeIndex>)> = std::mem::take(&mut slf.stack);
+        let mut path: Vec<NodeIndex> = std::mem::take(&mut slf.path);
+        let mut closed: HashSet<NodeIndex> = std::mem::take(&mut slf.closed);
+        let mut blocked: HashSet<NodeIndex> = std::mem::take(&mut slf.blocked);
+        let mut block: HashMap<NodeIndex, HashSet<NodeIndex>> = std::mem::take(&mut slf.block);
+        let mut subgraph: StableDiGraph<(), ()> = std::mem::take(&mut slf.subgraph);
+        let mut reverse_node_map: HashMap<NodeIndex, NodeIndex> =
+            std::mem::take(&mut slf.reverse_node_map);
+        let mut node_map: HashMap<NodeIndex, NodeIndex> = std::mem::take(&mut slf.node_map);
+
+        if let Some(res) = process_stack(
+            slf.start_node,
+            &mut stack,
+            &mut path,
+            &mut closed,
+            &mut blocked,
+            &mut block,
+            &subgraph,
+            &reverse_node_map,
+        ) {
+            // Store internal state on yield
+            slf.stack = stack;
+            slf.path = path;
+            slf.closed = closed;
+            slf.blocked = blocked;
+            slf.block = block;
+            slf.subgraph = subgraph;
+            slf.reverse_node_map = reverse_node_map;
+            slf.node_map = node_map;
+            return Ok(res);
+        } else {
+            subgraph.remove_node(slf.start_node);
+            slf.scc
+                .extend(kosaraju_scc(&subgraph).into_iter().filter_map(|scc| {
+                    if scc.len() > 1 {
+                        let res = scc
+                            .iter()
+                            .map(|n| reverse_node_map[n])
+                            .collect::<Vec<NodeIndex>>();
+                        Some(res)
+                    } else {
+                        None
+                    }
+                }));
+        }
+        while let Some(mut scc) = slf.scc.pop() {
+            let temp = build_subgraph(&slf.graph_clone, &scc);
+            subgraph = temp.0;
+            node_map = temp.1;
+            reverse_node_map = node_map.iter().map(|(k, v)| (*v, *k)).collect();
+            // start_node, path, blocked, closed, block and stack all in subgraph basis
+            slf.start_node = node_map[&scc.pop().unwrap()];
+            path = vec![slf.start_node];
+            blocked = path.iter().copied().collect();
+            // Nodes in cycle all
+            closed = HashSet::new();
+            block = HashMap::new();
+            stack = vec![(
+                slf.start_node,
+                subgraph
+                    .neighbors(slf.start_node)
+                    .collect::<IndexSet<NodeIndex>>(),
+            )];
+            if let Some(res) = process_stack(
+                slf.start_node,
+                &mut stack,
+                &mut path,
+                &mut closed,
+                &mut blocked,
+                &mut block,
+                &subgraph,
+                &reverse_node_map,
+            ) {
+                // Store internal state on yield
+                slf.stack = stack;
+                slf.path = path;
+                slf.closed = closed;
+                slf.blocked = blocked;
+                slf.block = block;
+                slf.subgraph = subgraph;
+                slf.reverse_node_map = reverse_node_map;
+                slf.node_map = node_map;
+                return Ok(res);
+            }
+            subgraph.remove_node(slf.start_node);
+            slf.scc
+                .extend(kosaraju_scc(&subgraph).into_iter().filter_map(|scc| {
+                    if scc.len() > 1 {
+                        let res = scc
+                            .iter()
+                            .map(|n| reverse_node_map[n])
+                            .collect::<Vec<NodeIndex>>();
+                        Some(res)
+                    } else {
+                        None
+                    }
+                }));
+        }
+        Ok(IterNextOutput::Return("Ended"))
+    }
+}

src/connectivity/mod.rs

@@ -14,6 +14,7 @@
 
 mod all_pairs_all_simple_paths;
 mod core_number;
+mod johnson_simple_cycles;
 
 use super::{digraph, get_edge_iter_with_weights, graph, weight_callable, InvalidNode, NullGraph};
 
@@ -126,6 +127,23 @@ pub fn cycle_basis(graph: &graph::PyGraph, root: Option<usize>) -> Vec<Vec<usize
     cycles
 }
 
+/// Find all simple cycles of a :class:`~.PyDiGraph`
+///
+/// A "simple cycle" (called an elementary circuit in [1]) is a cycle (or closed path)
+/// where no node appears more than once.
+///
+/// This function is a an implementation of Johnson's algorithm [1] also based
+/// on the non-recursive implementation found in NetworkX. [2][3]
+///
+/// [1] https://doi.org/10.1137/0204007
+/// [2] https://networkx.org/documentation/stable/reference/algorithms/generated/networkx.algorithms.cycles.simple_cycles.html
+/// [3] https://github.com/networkx/networkx/blob/networkx-2.8.4/networkx/algorithms/cycles.py#L98-L222
+#[pyfunction]
+#[pyo3(text_signature = "(graph, /)")]
+pub fn simple_cycles(graph: &digraph::PyDiGraph) -> johnson_simple_cycles::SimpleCycleIter {
+    johnson_simple_cycles::SimpleCycleIter::new(graph)
+}
+
 /// Compute the strongly connected components for a directed graph
 ///
 /// This function is implemented using Kosaraju's algorithm

src/lib.rs

@@ -424,6 +424,7 @@ fn rustworkx(py: Python<'_>, m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(undirected_gnm_random_graph))?;
     m.add_wrapped(wrap_pyfunction!(random_geometric_graph))?;
     m.add_wrapped(wrap_pyfunction!(cycle_basis))?;
+    m.add_wrapped(wrap_pyfunction!(simple_cycles))?;
     m.add_wrapped(wrap_pyfunction!(strongly_connected_components))?;
     m.add_wrapped(wrap_pyfunction!(digraph_dfs_edges))?;
     m.add_wrapped(wrap_pyfunction!(graph_dfs_edges))?;