wolf-graph 0.1.0

use anyhow::Result;

use crate::{DFSVisitor, EdgeID, Edges, NodeID, Nodes, VisitableGraph};

enum State {
    Discovered,
    Finished,
}

/// A trait for a graph that supports depth-first search.
pub trait DepthFirstSearch<G, O>
where
    G: VisitableGraph
{
    fn depth_first_search_opt<Visitor>(&self, visitor: &mut Visitor, roots: &Nodes, roots_only: bool, excluded_edge: Option<impl AsRef<EdgeID>>) -> Result<O>
    where Visitor: DFSVisitor<Graph = G, Output = O>;

    fn depth_first_search<Visitor>(&self, visitor: &mut Visitor) -> Result<O>
    where Visitor: DFSVisitor<Graph = G, Output = O> {
        self.depth_first_search_opt(visitor, &Nodes::new(), false, None::<EdgeID>)
    }
}

impl<G, O> DepthFirstSearch<G, O> for G
where
    G: VisitableGraph
{
    fn depth_first_search_opt<Visitor>(&self, visitor: &mut Visitor, roots: &Nodes, roots_only: bool, excluded_edge: Option<impl AsRef<EdgeID>>) -> Result<O>
    where
        Visitor: DFSVisitor<Graph = G, Output = O>,
    {
        for node in self.all_nodes() {
            if let Some(result) = visitor.init_node(self, &node)? {
                return Ok(result);
            }
        }

        let mut states = std::collections::HashMap::new();

        let excluded_edge = excluded_edge.as_ref();
        for root in roots.iter() {
            if let Some(result) = search_from_root(self, visitor, &mut states, root, excluded_edge)? {
                return Ok(result);
            }
        }
        if !roots_only {
            for node in self.all_nodes() {
                if let Some(result) = search_from_root(self, visitor, &mut states, &node, excluded_edge)? {
                    return Ok(result);
                }
            }
        }

        return visitor.finish(self);

        fn search_from_root<Graph, Visitor>(
            graph: &Graph,
            visitor: &mut Visitor,
            states: &mut std::collections::HashMap<NodeID, State>,
            root: &NodeID,
            excluded_edge: Option<impl AsRef<EdgeID>>,
        ) -> Result<Option<Visitor::Output>>
        where
            Graph: VisitableGraph,
            Visitor: DFSVisitor<Graph = Graph>,
        {
            if states.contains_key(root) {
                return Ok(None);
            }
            if let Some(result) = visitor.start_node(graph, root)? {
                return Ok(Some(result));
            }
            states.insert(root.clone(), State::Discovered);
            if let Some(result) = visitor.discover_node(graph, root)? {
                return Ok(Some(result));
            }
            let excluded_edge = excluded_edge.as_ref().map(|edge| edge.as_ref());
            let out_edges: Edges = graph.out_edges(root)?
                .iter()
                .filter(|&edge| Some(edge) != excluded_edge).cloned()
                .collect();
            let mut stack: Vec<(NodeID, Option<EdgeID>, std::collections::BTreeSet<EdgeID>)> = Vec::new();
            stack.push((root.clone(), None, out_edges));
            while let Some((source, finished_edge, remaining_out_edges)) = stack.pop() {
                let mut source = source;
                if let Some(finished_edge) = finished_edge {
                    if let Some(result) = visitor.finish_edge(graph, &finished_edge)? {
                        return Ok(Some(result));
                    }
                }

                let mut remaining_out_edges = remaining_out_edges;
                while let Some(edge) = remaining_out_edges.pop_last() {
                    let target = graph.target(&edge)?;
                    if let Some(result) = visitor.examine_edge(graph, &edge)? {
                        return Ok(Some(result));
                    }
                    match states.get(&target) {
                        None => {
                            if let Some(result) = visitor.tree_edge(graph, &edge)? {
                                return Ok(Some(result));
                            }
                            stack.push((source, Some(edge), remaining_out_edges));
                            source = target;
                            states.insert(source.clone(), State::Discovered);
                            if let Some(result) = visitor.discover_node(graph, &source)? {
                                return Ok(Some(result));
                            }
                            remaining_out_edges = graph.out_edges(&source)?;
                        }
                        Some(State::Discovered) => {
                            if let Some(result) = visitor.back_edge(graph, &edge)? {
                                return Ok(Some(result));
                            }
                            if let Some(result) = visitor.finish_edge(graph, &edge)? {
                                return Ok(Some(result));
                            }
                        }
                        Some(State::Finished) => {
                            if let Some(result) = visitor.forward_or_cross_edge(graph, &edge)? {
                                return Ok(Some(result));
                            }
                            if let Some(result) = visitor.finish_edge(graph, &edge)? {
                                return Ok(Some(result));
                            }
                        }
                    }
                }
                states.insert(source.clone(), State::Finished);
                if let Some(result) = visitor.finish_node(graph, &source)? {
                    return Ok(Some(result));
                }
            }
            Ok(None)
        }
    }
}