use std::ops::ControlFlow;
use ahash::AHashSet;
use petgraph::{
algo::tarjan_scc,
visit::{GraphBase, IntoNeighbors, IntoNodeIdentifiers, NodeIndexable},
};
pub trait Cycles {
type NodeId;
fn visit_cycles<F, B>(&self, visitor: F) -> Option<B>
where
F: FnMut(&Self, &[Self::NodeId]) -> ControlFlow<B>;
fn visit_all_cycles<F>(&self, mut visitor: F)
where
F: FnMut(&Self, &[Self::NodeId]),
{
self.visit_cycles(|g, n| {
visitor(g, n);
ControlFlow::<(), ()>::Continue(())
});
}
fn cycles(&self) -> Vec<Vec<Self::NodeId>>;
}
impl<Graph: GraphBase> Cycles for Graph
where
for<'a> &'a Graph: IntoNodeIdentifiers
+ IntoNeighbors
+ NodeIndexable
+ GraphBase<NodeId = <Graph as GraphBase>::NodeId>,
{
type NodeId = <Graph as GraphBase>::NodeId;
fn visit_cycles<F, B>(&self, mut visitor: F) -> Option<B>
where
F: FnMut(&Graph, &[Self::NodeId]) -> ControlFlow<B>,
{
for component in tarjan_scc(self) {
let mut finder = CycleFinder::new(self, component);
if let ControlFlow::Break(b) = finder.visit(&mut visitor) {
return Some(b);
}
}
None
}
fn cycles(&self) -> Vec<Vec<Self::NodeId>> {
let mut cycles = Vec::new();
self.visit_cycles(|_, cycle| {
cycles.push(cycle.to_vec());
ControlFlow::<(), ()>::Continue(())
});
cycles
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
struct CycleFinder<G, N> {
graph: G,
scc: Vec<N>,
blocked: Vec<bool>,
b: Vec<AHashSet<usize>>,
stack: Vec<N>,
s: usize,
}
impl<G> CycleFinder<G, G::NodeId>
where
G: IntoNodeIdentifiers + IntoNeighbors + NodeIndexable,
{
fn new(graph: G, scc: Vec<G::NodeId>) -> Self {
let num_vertices = scc.len();
Self {
graph,
scc,
blocked: vec![false; num_vertices],
b: vec![Default::default(); num_vertices],
stack: Default::default(),
s: Default::default(),
}
}
fn visit<F, B>(&mut self, visitor: &mut F) -> ControlFlow<B>
where
F: FnMut(G, &[G::NodeId]) -> ControlFlow<B>,
{
for s in 0..self.scc.len() {
self.s = s;
self.blocked[s..].fill(false);
for b in &mut self.b[s + 1..] {
b.clear();
}
if let ControlFlow::Break(b) = self.circuit(s, visitor) {
return ControlFlow::Break(b);
}
self.blocked[s] = true;
}
ControlFlow::Continue(())
}
fn circuit<B, F>(
&mut self,
v: usize,
visitor: &mut F,
) -> ControlFlow<B, bool>
where
F: FnMut(G, &[G::NodeId]) -> ControlFlow<B>,
{
let mut f = false;
self.stack.push(self.scc[v]);
self.blocked[v] = true;
for w in self.adjacent_vertices(v) {
if w == self.s {
if let ControlFlow::Break(b) = visitor(self.graph, &self.stack)
{
return ControlFlow::Break(b);
}
f = true;
} else if !self.blocked[w]
&& matches!(
self.circuit(w, visitor),
ControlFlow::Continue(true)
)
{
f = true;
}
}
if f {
self.unblock(v)
} else {
for w in self.adjacent_vertices(v) {
self.b[w].insert(v);
}
}
self.stack.pop(); ControlFlow::Continue(f)
}
fn unblock(&mut self, v: usize) {
self.blocked[v] = false;
let tmp = self.b[v].clone();
for w in tmp {
if self.blocked[w] {
self.unblock(w)
}
}
self.b[v].clear()
}
fn adjacent_vertices(&self, v: usize) -> Vec<usize> {
self.graph
.neighbors(self.scc[v])
.filter_map(|n| self.scc.iter().position(|v| *v == n))
.collect()
}
}
#[cfg(test)]
mod tests {
use petgraph::{Directed, graph::Graph, graphmap::GraphMap};
#[test]
fn test_graph() {
let g: Graph<usize, (), Directed> = Graph::new();
crate::Cycles::cycles(&g);
}
#[test]
fn test_graph_map() {
let g: GraphMap<usize, (), Directed> = GraphMap::new();
crate::Cycles::cycles(&g);
}
}