From 47c337a465b822b30e13db90e7fd52c8fcaf2b0e Mon Sep 17 00:00:00 2001 From: Matthew Treinish Date: Sun, 17 Nov 2024 09:54:13 -0500 Subject: [PATCH 1/3] Migrate Johnson's algorithm to rustworkx-core This commit adds Johnson's algorithm for computing simple cycles to rustworkx-core. This implementation returns an iterator-like struct which requires a reference to the graph for each step of the iterator. This mirrors the structure of the Bfs and similar structs in petgraph. While an implementor of Iterator could be used for this, this complicates using it from Python because we can't have a shared reference in a pyclass struct. This also has the advantage of enabling mutable references while iterating over the cycles. To facilitate this 2 traits are added `EdgeFindable` and `EdgeRemovable` which add traits for `remove_edge()` and `find_edge()` to petgraph graph types. The `EdgeRemovable` trait isn't actually used currently, it was in an earlier local draft of the PR which attempted to remove self cycles from a clone before changing direction of the interface. Since the implementation was done and relatively simple, and the potential general usefulness this was left in. `EdgeFindable` is actively used by this new function however. --- ...on-simple-cycle-core-ac0c09d6fce07f8a.yaml | 12 + .../src/connectivity/johnson_simple_cycles.rs | 532 ++++++++++++++++++ rustworkx-core/src/connectivity/mod.rs | 2 + rustworkx-core/src/graph_ext/mod.rs | 55 ++ src/connectivity/johnson_simple_cycles.rs | 313 ++--------- src/connectivity/mod.rs | 7 +- 6 files changed, 649 insertions(+), 272 deletions(-) create mode 100644 releasenotes/notes/add-johnson-simple-cycle-core-ac0c09d6fce07f8a.yaml create mode 100644 rustworkx-core/src/connectivity/johnson_simple_cycles.rs diff --git a/releasenotes/notes/add-johnson-simple-cycle-core-ac0c09d6fce07f8a.yaml b/releasenotes/notes/add-johnson-simple-cycle-core-ac0c09d6fce07f8a.yaml new file mode 100644 index 0000000000..3b978f4957 --- /dev/null +++ b/releasenotes/notes/add-johnson-simple-cycle-core-ac0c09d6fce07f8a.yaml @@ -0,0 +1,12 @@ +--- +features: + - | + Added a new function, ``johnson_simple_cycles``, to the ``rustworkx-core`` + crate. This function implements `Johnson's algorithm `__ for finding all + elementary cycles in a directed graph. + - | + Added a new trait ``EdgeFindable`` to find an ``EdgeIndex`` from a graph + given a pair of node indices. + - | + Added a new trait ``EdgeRemovable`` to remove an edge from a graph by + its ``EdgeIndex``. diff --git a/rustworkx-core/src/connectivity/johnson_simple_cycles.rs b/rustworkx-core/src/connectivity/johnson_simple_cycles.rs new file mode 100644 index 0000000000..c9a3a71b93 --- /dev/null +++ b/rustworkx-core/src/connectivity/johnson_simple_cycles.rs @@ -0,0 +1,532 @@ +// 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 ahash::RandomState; +use hashbrown::{HashMap, HashSet}; +use indexmap::IndexSet; +use std::hash::Hash; + +use petgraph::algo::kosaraju_scc; +use petgraph::stable_graph::{NodeIndex, StableDiGraph}; +use petgraph::visit::{ + EdgeCount, EdgeRef, GraphBase, GraphProp, IntoEdgeReferences, IntoNeighbors, + IntoNeighborsDirected, IntoNodeReferences, NodeCount, NodeFiltered, NodeIndexable, NodeRef, + Visitable, +}; +use petgraph::Directed; + +use crate::graph_ext::EdgeFindable; + +fn build_subgraph( + graph: G, + nodes: &[NodeIndex], +) -> (StableDiGraph<(), ()>, HashMap) +where + G: EdgeCount + NodeCount + IntoEdgeReferences + IntoNodeReferences + GraphBase + NodeIndexable, + ::NodeId: Hash + Eq, +{ + let node_set: HashSet = nodes.iter().copied().collect(); + let mut node_map: HashMap = HashMap::with_capacity(nodes.len()); + let node_filter = + |node: G::NodeId| -> bool { node_set.contains(&NodeIndex::new(graph.to_index(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::from_fn(graph, node_filter); + for node in filtered.node_references() { + let new_node = out_graph.add_node(()); + node_map.insert(NodeIndex::new(graph.to_index(node.id())), new_node); + } + for edge in filtered.edge_references() { + let new_source = *node_map + .get(&NodeIndex::new(graph.to_index(edge.source()))) + .unwrap(); + let new_target = *node_map + .get(&NodeIndex::new(graph.to_index(edge.target()))) + .unwrap(); + out_graph.add_edge( + NodeIndex::new(new_source.index()), + NodeIndex::new(new_target.index()), + (), + ); + } + (out_graph, node_map) +} + +fn unblock( + node: NodeIndex, + blocked: &mut HashSet, + block: &mut HashMap>, +) { + let mut stack: IndexSet = IndexSet::with_hasher(RandomState::default()); + 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)>, + path: &mut Vec, + closed: &mut HashSet, + blocked: &mut HashSet, + block: &mut HashMap>, + subgraph: &StableDiGraph<(), ()>, + reverse_node_map: &HashMap, +) -> Option> { + 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 = Vec::with_capacity(path.len()); + for n in path { + out_path.push(reverse_node_map[n]); + closed.insert(*n); + } + return Some(out_path); + } else if blocked.insert(next_node) { + path.push(next_node); + stack.push(( + next_node, + subgraph + .neighbors(next_node) + .collect::>(), + )); + 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 +} + +/// An iterator of simple cycles in a graph +/// +/// `SimpleCycleIter` does not itself borrow the graph, and because of this +/// you can run the algorithm while retaining mutable access to it, if you +/// use like it the following example: +/// +/// ``` +/// use rustworkx_core::petgraph::prelude::*; +/// use rustworkx_core::connectivity::johnson_simple_cycles; +/// +/// let mut graph = DiGraph::<_,()>::new(); +/// let a = graph.add_node(0); +/// +/// let mut cycle_iter = johnson_simple_cycles(&graph, None); +/// while let Some(cycle) = cycle_iter.next(&graph) { +/// // We can access `graph` mutably here still +/// graph[a] += 1; +/// } +/// +/// assert_eq!(graph[a], 0); +/// ``` +pub struct SimpleCycleIter { + scc: Vec>, + self_cycles: Option>, + path: Vec, + blocked: HashSet, + closed: HashSet, + block: HashMap>, + stack: Vec<(NodeIndex, IndexSet)>, + start_node: NodeIndex, + node_map: HashMap, + reverse_node_map: HashMap, + subgraph: StableDiGraph<(), ()>, +} + +impl SimpleCycleIter { + /// Return the next cycle found, if `None` is returned the algorithm is complete and all + /// cycles have been found. + pub fn next(&mut self, graph: G) -> Option> + where + G: IntoEdgeReferences + + IntoNodeReferences + + GraphBase + + EdgeCount + + NodeCount + + NodeIndexable, + ::NodeId: Hash + Eq, + { + if self.self_cycles.is_some() { + let self_cycles = self.self_cycles.as_mut().unwrap(); + let cycle_node = self_cycles.pop().unwrap(); + if self_cycles.is_empty() { + self.self_cycles = None; + } + return Some(vec![cycle_node]); + } + // Restore previous state if it exists + let mut stack: Vec<(NodeIndex, IndexSet)> = + std::mem::take(&mut self.stack); + let mut path: Vec = std::mem::take(&mut self.path); + let mut closed: HashSet = std::mem::take(&mut self.closed); + let mut blocked: HashSet = std::mem::take(&mut self.blocked); + let mut block: HashMap> = std::mem::take(&mut self.block); + let mut subgraph: StableDiGraph<(), ()> = std::mem::take(&mut self.subgraph); + let mut reverse_node_map: HashMap = + std::mem::take(&mut self.reverse_node_map); + let mut node_map: HashMap = std::mem::take(&mut self.node_map); + if let Some(res) = process_stack( + self.start_node, + &mut stack, + &mut path, + &mut closed, + &mut blocked, + &mut block, + &subgraph, + &reverse_node_map, + ) { + // Store internal state on yield + self.stack = stack; + self.path = path; + self.closed = closed; + self.blocked = blocked; + self.block = block; + self.subgraph = subgraph; + self.reverse_node_map = reverse_node_map; + self.node_map = node_map; + return Some(res); + } else { + subgraph.remove_node(self.start_node); + self.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::>(); + Some(res) + } else { + None + } + })); + } + while let Some(mut scc) = self.scc.pop() { + let temp = build_subgraph(graph, &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 + self.start_node = node_map[&scc.pop().unwrap()]; + path = vec![self.start_node]; + blocked = path.iter().copied().collect(); + // Nodes in cycle all + closed = HashSet::new(); + block = HashMap::new(); + stack = vec![( + self.start_node, + subgraph + .neighbors(self.start_node) + .collect::>(), + )]; + if let Some(res) = process_stack( + self.start_node, + &mut stack, + &mut path, + &mut closed, + &mut blocked, + &mut block, + &subgraph, + &reverse_node_map, + ) { + // Store internal state on yield + self.stack = stack; + self.path = path; + self.closed = closed; + self.blocked = blocked; + self.block = block; + self.subgraph = subgraph; + self.reverse_node_map = reverse_node_map; + self.node_map = node_map; + return Some(res); + } + subgraph.remove_node(self.start_node); + self.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::>(); + Some(res) + } else { + None + } + })); + } + None + } +} + +/// /// Find all simple cycles of a graph +/// +/// 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] +/// +/// To handle self cycles in a manner consistent with the NetworkX implementation you should +/// use the ``self_cycles`` argument to collect manually collected self cycle and then remove +/// the edges leading to a self cycle from the graph. If you don't do this +/// the self cycle may or may not be returned by the iterator. The underlying algorithm is not +/// able to consistently detect self cycles so it is best to handle them before calling this +/// function. The example below shows a pattern for doing this. You will need to clone the graph +/// to do this detection without modifying the graph. +/// +/// # Returns +/// +/// This function returns a `SimpleCycleIter` iterator which returns a `Vec` of `NodeIndex`. +/// Note the `NodeIndex` type is not neccesarily the same as the input graph, as it's built +/// using an internal `StableGraph` used by the algorithm. If your input `graph` uses a +/// different node index type that differs from the default `NodeIndex`/`NodeIndex` +/// you will want to convert these objects to your native `NodeIndex` type. +/// +/// The return from this function is not guaranteed to have a particular order for either the +/// cycles or the indices in each cycle. +/// +/// [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 +/// +/// # Example: +/// +/// ```rust +/// use rustworkx_core::petgraph::prelude::*; +/// use rustworkx_core::connectivity::johnson_simple_cycles; +/// +/// let mut graph = DiGraph::<(), ()>::new(); +/// graph.extend_with_edges([(0, 0), (0, 1), (0, 2), (1, 2), (2, 0), (2, 1), (2, 2)]); +/// +/// // Handle self cycles +/// let self_cycles_vec: Vec = graph +/// .node_indices() +/// .filter(|n| graph.neighbors(*n).any(|x| x == *n)) +/// .collect(); +/// for node in &self_cycles_vec { +/// while let Some(edge_index) = graph.find_edge(*node, *node) { +/// graph.remove_edge(edge_index); +/// } +/// } +/// let self_cycles = if self_cycles_vec.is_empty() { +/// None +/// } else { +/// Some(self_cycles_vec) +/// }; +/// +/// let mut cycles_iter = johnson_simple_cycles(&graph, self_cycles); +/// +/// let mut cycles = Vec::new(); +/// while let Some(mut cycle) = cycles_iter.next(&graph) { +/// cycle.sort(); +/// cycles.push(cycle); +/// } +/// +/// let expected = vec![ +/// vec![NodeIndex::new(0)], +/// vec![NodeIndex::new(2)], +/// vec![NodeIndex::new(0), NodeIndex::new(1), NodeIndex::new(2)], +/// vec![NodeIndex::new(0), NodeIndex::new(2)], +/// vec![NodeIndex::new(1), NodeIndex::new(2)], +/// ]; +/// +/// assert_eq!(expected.len(), cycles.len()); +/// for cycle in cycles { +/// assert!(expected.contains(&cycle)); +/// } +/// ``` +pub fn johnson_simple_cycles( + graph: G, + self_cycles: Option::NodeId>>, +) -> SimpleCycleIter +where + G: IntoEdgeReferences + + IntoNodeReferences + + GraphBase + + EdgeCount + + NodeCount + + NodeIndexable + + Clone + + IntoNeighbors + + IntoNeighborsDirected + + Visitable + + EdgeFindable + + GraphProp, + ::NodeId: Hash + Eq, +{ + let self_cycles = self_cycles.map(|self_cycles_vec| { + self_cycles_vec + .into_iter() + .map(|n| NodeIndex::new(graph.to_index(n))) + .collect() + }); + let strongly_connected_components: Vec> = kosaraju_scc(graph) + .into_iter() + .filter_map(|component| { + if component.len() > 1 { + Some( + component + .into_iter() + .map(|n| NodeIndex::new(graph.to_index(n))) + .collect(), + ) + } else { + None + } + }) + .collect(); + SimpleCycleIter { + 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(u32::MAX as usize), + node_map: HashMap::new(), + reverse_node_map: HashMap::new(), + subgraph: StableDiGraph::new(), + } +} + +#[cfg(test)] +mod test_longest_path { + use super::*; + use petgraph::graph::DiGraph; + use petgraph::stable_graph::NodeIndex; + use petgraph::stable_graph::StableDiGraph; + + #[test] + fn test_empty_graph() { + let graph: DiGraph<(), ()> = DiGraph::new(); + let mut result: Vec<_> = Vec::new(); + let mut cycle_iter = johnson_simple_cycles(&graph, None); + while let Some(cycle) = cycle_iter.next(&graph) { + result.push(cycle); + } + let expected: Vec> = Vec::new(); + assert_eq!(expected, result) + } + + #[test] + fn test_empty_stable_graph() { + let graph: StableDiGraph<(), ()> = StableDiGraph::new(); + let mut result: Vec<_> = Vec::new(); + let mut cycle_iter = johnson_simple_cycles(&graph, None); + while let Some(cycle) = cycle_iter.next(&graph) { + result.push(cycle); + } + let expected: Vec> = Vec::new(); + assert_eq!(expected, result) + } + + #[test] + fn test_figure_1() { + for k in 3..10 { + let mut graph: DiGraph<(), ()> = DiGraph::new(); + let mut edge_list: Vec<[usize; 2]> = Vec::new(); + for n in 2..k + 2 { + edge_list.push([1, n]); + edge_list.push([n, k + 2]); + } + edge_list.push([2 * k + 1, 1]); + for n in k + 2..2 * k + 2 { + edge_list.push([n, 2 * k + 2]); + edge_list.push([n, n + 1]); + } + edge_list.push([2 * k + 3, k + 2]); + for n in 2 * k + 3..3 * k + 3 { + edge_list.push([2 * k + 2, n]); + edge_list.push([n, 3 * k + 3]); + } + edge_list.push([3 * k + 3, 2 * k + 2]); + graph.extend_with_edges( + edge_list + .into_iter() + .map(|x| (NodeIndex::new(x[0]), NodeIndex::new(x[1]))), + ); + let mut cycles_iter = johnson_simple_cycles(&graph, None); + let mut res = 0; + while let Some(_) = cycles_iter.next(&graph) { + res += 1; + } + assert_eq!(res, 3 * k); + } + } + + #[test] + fn test_figure_1_stable_graph() { + for k in 3..10 { + let mut graph: StableDiGraph<(), ()> = StableDiGraph::new(); + let mut edge_list: Vec<[usize; 2]> = Vec::new(); + for n in 2..k + 2 { + edge_list.push([1, n]); + edge_list.push([n, k + 2]); + } + edge_list.push([2 * k + 1, 1]); + for n in k + 2..2 * k + 2 { + edge_list.push([n, 2 * k + 2]); + edge_list.push([n, n + 1]); + } + edge_list.push([2 * k + 3, k + 2]); + for n in 2 * k + 3..3 * k + 3 { + edge_list.push([2 * k + 2, n]); + edge_list.push([n, 3 * k + 3]); + } + edge_list.push([3 * k + 3, 2 * k + 2]); + graph.extend_with_edges( + edge_list + .into_iter() + .map(|x| (NodeIndex::new(x[0]), NodeIndex::new(x[1]))), + ); + let mut cycles_iter = johnson_simple_cycles(&graph, None); + let mut res = 0; + while let Some(_) = cycles_iter.next(&graph) { + res += 1; + } + assert_eq!(res, 3 * k); + } + } +} diff --git a/rustworkx-core/src/connectivity/mod.rs b/rustworkx-core/src/connectivity/mod.rs index fa236d8b61..30d42d509d 100644 --- a/rustworkx-core/src/connectivity/mod.rs +++ b/rustworkx-core/src/connectivity/mod.rs @@ -20,6 +20,7 @@ mod core_number; mod cycle_basis; mod find_cycle; mod isolates; +mod johnson_simple_cycles; mod min_cut; pub use all_simple_paths::{ @@ -35,4 +36,5 @@ pub use core_number::core_number; pub use cycle_basis::cycle_basis; pub use find_cycle::find_cycle; pub use isolates::isolates; +pub use johnson_simple_cycles::{johnson_simple_cycles, SimpleCycleIter}; pub use min_cut::stoer_wagner_min_cut; diff --git a/rustworkx-core/src/graph_ext/mod.rs b/rustworkx-core/src/graph_ext/mod.rs index 256d6ac0a5..7ca5723ae5 100644 --- a/rustworkx-core/src/graph_ext/mod.rs +++ b/rustworkx-core/src/graph_ext/mod.rs @@ -61,6 +61,8 @@ //! | HasParallelEdgesDirected | x | x | x | x | x | x | //! | HasParallelEdgesUndirected | x | x | x | x | x | x | //! | NodeRemovable | x | x | x | x | | | +//! | EdgeRemovable | x | x | | | | | +//! | EdgeFindable | x | x | | | | | use petgraph::graph::IndexType; use petgraph::graphmap::{GraphMap, NodeTrait}; @@ -130,3 +132,56 @@ impl, Ix: IndexType> NodeRemovab None } } + +/// A graph whose edge may be removed by an edge id. +pub trait EdgeRemovable: Data { + type Output; + fn remove_edge(&mut self, edge: Self::EdgeId) -> Self::Output; +} + +impl EdgeRemovable for StableGraph +where + Ty: EdgeType, + Ix: IndexType, +{ + type Output = Option; + fn remove_edge(&mut self, edge: Self::EdgeId) -> Option { + self.remove_edge(edge) + } +} + +impl EdgeRemovable for Graph +where + Ty: EdgeType, + Ix: IndexType, +{ + type Output = Option; + fn remove_edge(&mut self, edge: Self::EdgeId) -> Option { + self.remove_edge(edge) + } +} + +/// A graph that can find edges by a pair of node ids. +pub trait EdgeFindable: Data { + fn edge_find(&self, a: Self::NodeId, b: Self::NodeId) -> Option; +} + +impl EdgeFindable for &StableGraph +where + Ty: EdgeType, + Ix: IndexType, +{ + fn edge_find(&self, a: Self::NodeId, b: Self::NodeId) -> Option { + self.find_edge(a, b) + } +} + +impl EdgeFindable for &Graph +where + Ty: EdgeType, + Ix: IndexType, +{ + fn edge_find(&self, a: Self::NodeId, b: Self::NodeId) -> Option { + self.find_edge(a, b) + } +} diff --git a/src/connectivity/johnson_simple_cycles.rs b/src/connectivity/johnson_simple_cycles.rs index 555df07151..706e0b5ce3 100644 --- a/src/connectivity/johnson_simple_cycles.rs +++ b/src/connectivity/johnson_simple_cycles.rs @@ -10,301 +10,74 @@ // License for the specific language governing permissions and limitations // under the License. -use ahash::RandomState; -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 std::ops::Deref; use pyo3::prelude::*; use crate::iterators::NodeIndices; - -fn build_subgraph( - graph: &StablePyGraph, - nodes: &[NodeIndex], -) -> (StableDiGraph<(), ()>, HashMap) { - let node_set: HashSet = nodes.iter().copied().collect(); - let mut node_map: HashMap = 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) -} +use rustworkx_core::connectivity::{johnson_simple_cycles, SimpleCycleIter}; #[pyclass(module = "rustworkx")] -pub struct SimpleCycleIter { - graph_clone: StablePyGraph, - scc: Vec>, - self_cycles: Option>, - path: Vec, - blocked: HashSet, - closed: HashSet, - block: HashMap>, - stack: Vec<(NodeIndex, IndexSet)>, - start_node: NodeIndex, - node_map: HashMap, - reverse_node_map: HashMap, - subgraph: StableDiGraph<(), ()>, +pub struct PySimpleCycleIter { + graph_clone: Py, + iter: SimpleCycleIter, } -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(); - +impl PySimpleCycleIter { + pub fn new(py: Python, graph: Bound) -> PyResult { // For compatibility with networkx manually insert self cycles and filter // from Johnson's algorithm - let self_cycles_vec: Vec = graph_clone + let self_cycles_vec: Vec = graph + .borrow() + .graph .node_indices() - .filter(|n| graph_clone.neighbors(*n).any(|x| x == *n)) + .filter(|n| graph.borrow().graph.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 + if self_cycles_vec.is_empty() { + let iter = johnson_simple_cycles(&graph.borrow().graph, None); + let out_graph = graph.unbind(); + Ok(PySimpleCycleIter { + graph_clone: out_graph, + iter, + }) } else { - Some(self_cycles_vec) - }; - let strongly_connected_components: Vec> = 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(u32::MAX as usize), - node_map: HashMap::new(), - reverse_node_map: HashMap::new(), - subgraph: StableDiGraph::new(), - } - } -} - -fn unblock( - node: NodeIndex, - blocked: &mut HashSet, - block: &mut HashMap>, -) { - let mut stack: IndexSet = IndexSet::with_hasher(RandomState::default()); - 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)>, - path: &mut Vec, - closed: &mut HashSet, - blocked: &mut HashSet, - block: &mut HashMap>, - subgraph: &StableDiGraph<(), ()>, - reverse_node_map: &HashMap, -) -> Option { - 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 = Vec::with_capacity(path.len()); - for n in path { - out_path.push(reverse_node_map[n].index()); - closed.insert(*n); + // Copy graph to remove self edges before running johnson's algorithm + let mut graph_clone = graph.borrow().copy(); + for node in &self_cycles_vec { + while let Some(edge_index) = graph_clone.graph.find_edge(*node, *node) { + graph_clone.graph.remove_edge(edge_index); } - return Some(NodeIndices { nodes: out_path }); - } else if blocked.insert(next_node) { - path.push(next_node); - stack.push(( - next_node, - subgraph - .neighbors(next_node) - .collect::>(), - )); - closed.remove(&next_node); - blocked.insert(next_node); - continue; } - } - if neighbors.is_empty() { - if closed.contains(this_node) { - unblock(*this_node, blocked, block); + let self_cycles = if self_cycles_vec.is_empty() { + None } 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(); + Some(self_cycles_vec) + }; + let iter = johnson_simple_cycles(&graph_clone.graph, self_cycles); + let out_graph = Py::new(py, graph_clone)?; + Ok(PySimpleCycleIter { + graph_clone: out_graph, + iter, + }) } } - None } #[pymethods] -impl SimpleCycleIter { - fn __iter__(slf: PyRef) -> Py { +impl PySimpleCycleIter { + fn __iter__(slf: PyRef) -> Py { slf.into() } - fn __next__(mut slf: PyRefMut) -> PyResult> { - 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(Some(NodeIndices { - nodes: vec![cycle_node.index()], - })); - } - // Restore previous state if it exists - let mut stack: Vec<(NodeIndex, IndexSet)> = - std::mem::take(&mut slf.stack); - let mut path: Vec = std::mem::take(&mut slf.path); - let mut closed: HashSet = std::mem::take(&mut slf.closed); - let mut blocked: HashSet = std::mem::take(&mut slf.blocked); - let mut block: HashMap> = std::mem::take(&mut slf.block); - let mut subgraph: StableDiGraph<(), ()> = std::mem::take(&mut slf.subgraph); - let mut reverse_node_map: HashMap = - std::mem::take(&mut slf.reverse_node_map); - let mut node_map: HashMap = 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(Some(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::>(); - 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::>(), - )]; - 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(Some(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::>(); - Some(res) - } else { - None - } - })); - } - Ok(None) + fn __next__(mut slf: PyRefMut, py: Python) -> PyResult> { + let py_clone = slf.graph_clone.clone_ref(py); + let binding = py_clone.borrow(py); + let graph = binding.deref(); + let res: Option> = slf.iter.next(&graph.graph); + Ok(res.map(|cycle| NodeIndices { + nodes: cycle.into_iter().map(|x| x.index()).collect(), + })) } } diff --git a/src/connectivity/mod.rs b/src/connectivity/mod.rs index 353a57fe25..bd4490c04e 100644 --- a/src/connectivity/mod.rs +++ b/src/connectivity/mod.rs @@ -92,8 +92,11 @@ pub fn cycle_basis(graph: &graph::PyGraph, root: Option) -> Vec johnson_simple_cycles::SimpleCycleIter { - johnson_simple_cycles::SimpleCycleIter::new(graph) +pub fn simple_cycles( + graph: Bound, + py: Python, +) -> PyResult { + johnson_simple_cycles::PySimpleCycleIter::new(py, graph) } /// Compute the strongly connected components for a directed graph From 8d66d675265269c8ed13c206fa5cbf98aacb64ed Mon Sep 17 00:00:00 2001 From: Matthew Treinish Date: Sun, 17 Nov 2024 13:58:30 -0500 Subject: [PATCH 2/3] Simplify python side logic --- src/connectivity/johnson_simple_cycles.rs | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/src/connectivity/johnson_simple_cycles.rs b/src/connectivity/johnson_simple_cycles.rs index 706e0b5ce3..2780cd3707 100644 --- a/src/connectivity/johnson_simple_cycles.rs +++ b/src/connectivity/johnson_simple_cycles.rs @@ -50,12 +50,7 @@ impl PySimpleCycleIter { graph_clone.graph.remove_edge(edge_index); } } - let self_cycles = if self_cycles_vec.is_empty() { - None - } else { - Some(self_cycles_vec) - }; - let iter = johnson_simple_cycles(&graph_clone.graph, self_cycles); + let iter = johnson_simple_cycles(&graph_clone.graph, Some(self_cycles_vec)); let out_graph = Py::new(py, graph_clone)?; Ok(PySimpleCycleIter { graph_clone: out_graph, From 081aa3c76256354ea6a5547486319f90da925e2a Mon Sep 17 00:00:00 2001 From: Matthew Treinish Date: Mon, 18 Nov 2024 08:39:38 -0500 Subject: [PATCH 3/3] Set name of pyclass to SimpleCycleIter for backwards compat --- src/connectivity/johnson_simple_cycles.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/connectivity/johnson_simple_cycles.rs b/src/connectivity/johnson_simple_cycles.rs index 2780cd3707..8d0636f2f7 100644 --- a/src/connectivity/johnson_simple_cycles.rs +++ b/src/connectivity/johnson_simple_cycles.rs @@ -19,7 +19,7 @@ use pyo3::prelude::*; use crate::iterators::NodeIndices; use rustworkx_core::connectivity::{johnson_simple_cycles, SimpleCycleIter}; -#[pyclass(module = "rustworkx")] +#[pyclass(module = "rustworkx", name = "SimpleCycleIter")] pub struct PySimpleCycleIter { graph_clone: Py, iter: SimpleCycleIter,