use std::{marker::PhantomData, ops::ControlFlow};
use bitvec::vec::BitVec;
use petgraph::graph::{IndexType, NodeIndex};
use crate::graph::Graph as G;
struct BitSet<T> {
vec: BitVec,
ty: PhantomData<T>,
}
impl<T: IndexType> BitSet<T> {
fn new_empty(domain_size: usize) -> Self {
Self { vec: BitVec::repeat(false, domain_size), ty: PhantomData }
}
fn insert(&mut self, elem: T) -> bool {
let changed = !self.vec[elem.index()];
self.vec.set(elem.index(), true);
changed
}
fn contains(&mut self, elem: T) -> bool {
self.vec[elem.index()]
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum NodeStatus {
Visited,
Settled,
}
struct Event<N> {
node: N,
becomes: NodeStatus,
}
pub struct TriColorDepthFirstSearch<'graph> {
graph: &'graph G,
stack: Vec<Event<NodeIndex>>,
visited: BitSet<NodeIndex>,
settled: BitSet<NodeIndex>,
}
impl<'graph> TriColorDepthFirstSearch<'graph> {
pub fn new(graph: &'graph G) -> Self {
TriColorDepthFirstSearch {
graph,
stack: vec![],
visited: BitSet::new_empty(graph.node_count()),
settled: BitSet::new_empty(graph.node_count()),
}
}
pub fn run_from<V>(mut self, root: NodeIndex, visitor: &mut V) -> Option<Vec<NodeIndex>>
where
V: TriColorVisitor<G>,
{
use NodeStatus::{Settled, Visited};
self.stack.push(Event { node: root, becomes: Visited });
loop {
match self.stack.pop()? {
Event { node, becomes: Settled } => {
let not_previously_settled = self.settled.insert(node);
assert!(not_previously_settled, "A node should be settled exactly once");
if let ControlFlow::Break(_val) = visitor.node_settled(node) {
return None;
}
}
Event { node, becomes: Visited } => {
let not_previously_visited = self.visited.insert(node);
let prior_status = if not_previously_visited {
None
} else if self.settled.contains(node) {
Some(Settled)
} else {
Some(Visited)
};
if let ControlFlow::Break(_val) = visitor.node_examined(node, prior_status) {
return Some(identify_cycle_nodes(self.stack, node));
}
if prior_status.is_some() {
continue;
}
self.stack.push(Event { node, becomes: Settled });
for succ in self.graph.neighbors(node) {
if !visitor.ignore_edge(node, succ) {
self.stack.push(Event { node: succ, becomes: Visited });
}
}
}
}
}
}
}
pub trait TriColorVisitor<G> {
type BreakVal;
fn node_examined(
&mut self,
_node: NodeIndex,
_prior_status: Option<NodeStatus>,
) -> ControlFlow<Self::BreakVal> {
ControlFlow::Continue(())
}
fn node_settled(&mut self, _node: NodeIndex) -> ControlFlow<Self::BreakVal> {
ControlFlow::Continue(())
}
fn ignore_edge(&mut self, _source: NodeIndex, _target: NodeIndex) -> bool {
false
}
}
pub struct CycleDetector;
impl<G> TriColorVisitor<G> for CycleDetector {
type BreakVal = ();
fn node_examined(
&mut self,
_node: NodeIndex,
prior_status: Option<NodeStatus>,
) -> ControlFlow<Self::BreakVal> {
match prior_status {
Some(NodeStatus::Visited) => ControlFlow::Break(()),
_ => ControlFlow::Continue(()),
}
}
fn ignore_edge(&mut self, source: NodeIndex, target: NodeIndex) -> bool {
source == target
}
}
fn identify_cycle_nodes(mut stack: Vec<Event<NodeIndex>>, orig_node: NodeIndex) -> Vec<NodeIndex> {
let mut cycle_nodes = Vec::new();
let mut curr_node = orig_node;
while let Some(event) = stack.pop() {
if event.becomes != NodeStatus::Settled {
continue;
}
cycle_nodes.push(curr_node);
curr_node = event.node;
if curr_node == orig_node {
return cycle_nodes;
}
}
panic!("Failed to identify cycle nodes")
}