use super::*;
pub struct Dfs<Node> {
stack: Vec<DfsEvent<Node>>,
}
impl<Node> Default for Dfs<Node> {
fn default() -> Self {
Self {
stack: Default::default(),
}
}
}
impl<Node> Dfs<Node> {
pub fn new(roots: impl IntoIterator<Item = Node>) -> Self {
let mut dfs = Self::default();
dfs.add_roots(roots);
dfs
}
pub fn add_root(&mut self, root: Node) {
self.stack.push(DfsEvent::Pre(root));
}
pub fn add_roots(&mut self, roots: impl IntoIterator<Item = Node>) {
self.stack
.extend(roots.into_iter().map(|v| DfsEvent::Pre(v)));
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DfsEvent<Node> {
Pre(Node),
AfterEdge(Node, Node),
Post(Node),
}
impl<Node> Dfs<Node>
where
Node: Copy,
{
pub fn next<G>(&mut self, graph: G, seen: impl Fn(Node) -> bool) -> Option<DfsEvent<Node>>
where
G: Graph<Node>,
{
loop {
let event = self.stack.pop()?;
if let DfsEvent::Pre(node) = event {
if seen(node) {
continue;
}
let successors = graph.successors(node);
let (min, max) = successors.size_hint();
let estimated_successors_len = max.unwrap_or_else(|| 2 * min);
self.stack.reserve(
2 * estimated_successors_len
+ 1,
);
self.stack.push(DfsEvent::Post(node));
for succ in successors {
self.stack.push(DfsEvent::AfterEdge(node, succ));
if !seen(succ) {
self.stack.push(DfsEvent::Pre(succ));
}
}
}
return Some(event);
}
}
}