gantz_core 0.4.1

//! Provides [`visit()`] and [`register()`] fns for generic gantz graphs.

use std::{
    collections::{BTreeSet, HashMap, HashSet},
    hash::{Hash, Hasher},
};

use crate::{
    Edge,
    node::{self, Node},
    visit,
};
use petgraph::visit::{
    Data, EdgeRef, IntoEdgeReferences, IntoEdgesDirected, IntoNeighborsDirected,
    IntoNodeIdentifiers, IntoNodeReferences, NodeIndexable, NodeRef, Topo, Visitable,
};
use steel::steel_vm::engine::Engine;

/// Hash the nodes, edges and their IDs of the given graph.
pub fn hash<G, H>(g: G, h: &mut H)
where
    G: Data + IntoEdgeReferences + IntoNodeReferences,
    G::EdgeId: Hash,
    G::EdgeWeight: Hash,
    G::NodeId: Hash,
    G::NodeWeight: Hash,
    H: Hasher,
{
    for n in g.node_references() {
        n.id().hash(h);
        n.weight().hash(h);
    }
    for e in g.edge_references() {
        e.id().hash(h);
        e.weight().hash(h);
    }
}

/// All of `g`'s nodes: the topological order where it exists, then any nodes on
/// a directed cycle (which [`Topo`] never yields) appended in index order.
/// Keeps acyclic traversal identical to a plain `Topo` walk while ensuring
/// cyclic nodes are still produced, deterministically.
pub(crate) fn cycle_tolerant_topo_order<G>(g: G) -> Vec<G::NodeId>
where
    G: IntoNodeIdentifiers + IntoNeighborsDirected + NodeIndexable + Visitable,
{
    let mut visited: HashSet<usize> = HashSet::new();
    let mut order: Vec<G::NodeId> = Vec::new();
    let mut topo = Topo::new(g);
    while let Some(n) = topo.next(g) {
        visited.insert(g.to_index(n));
        order.push(n);
    }
    let mut leftover: Vec<G::NodeId> = g
        .node_identifiers()
        .filter(|&n| !visited.contains(&g.to_index(n)))
        .collect();
    leftover.sort_by_key(|&n| g.to_index(n));
    order.extend(leftover);
    order
}

/// Visit all nodes in the graph in toposort order, and all nested nodes in
/// depth-first order. Tolerates cycles - see `cycle_tolerant_topo_order`
/// (order is irrelevant to the visitors, which key by path/id).
pub fn visit<'a, G>(
    get_node: node::GetNode<'a>,
    g: G,
    path: &[node::Id],
    visitor: &mut dyn node::Visitor,
) where
    G: Data<EdgeWeight = Edge> + IntoEdgesDirected + IntoNodeReferences + NodeIndexable + Visitable,
    G::NodeWeight: Node,
{
    let mut path = path.to_vec();
    for n in cycle_tolerant_topo_order(g) {
        let ix = g.to_index(n);
        path.push(ix);
        let inputs: Vec<_> = g
            .edges_directed(n, petgraph::Direction::Incoming)
            .map(|e_ref| (g.to_index(e_ref.source()), e_ref.weight().clone()))
            .collect();
        let ctx = visit::Ctx::new(get_node, &path, &inputs);

        // FIXME: index directly.
        let nref = g.node_references().find(|nref| nref.id() == n).unwrap();

        node::visit(ctx, nref.weight(), visitor);
        path.pop();
    }
}

/// Visit all nodes in the graph with a [`visit::TypedVisitor`].
pub fn visit_typed<'a, G, V>(get_node: node::GetNode<'a>, g: G, path: &[node::Id], visitor: &mut V)
where
    G: Data<EdgeWeight = Edge> + IntoEdgesDirected + IntoNodeReferences + NodeIndexable + Visitable,
    G::NodeWeight: Node,
    V: visit::TypedVisitor<G::NodeWeight>,
{
    visit(
        get_node,
        g,
        path,
        &mut visit::Typed::<&mut V, G::NodeWeight>::new(visitor),
    );
}

/// Register the given graph of nodes, including any nested nodes.
pub fn register<'a, G>(get_node: node::GetNode<'a>, g: G, path: &[node::Id], vm: &mut Engine)
where
    G: Data<EdgeWeight = Edge> + IntoEdgesDirected + IntoNodeReferences + NodeIndexable + Visitable,
    G::NodeWeight: Node,
{
    visit(get_node, g, path, &mut visit::Register(vm));
}

/// Collect all content addresses required by nodes in this graph.
pub fn required_addrs<'a, G>(get_node: node::GetNode<'a>, g: G) -> HashSet<gantz_ca::ContentAddr>
where
    G: Data<EdgeWeight = Edge> + IntoEdgesDirected + IntoNodeReferences + NodeIndexable + Visitable,
    G::NodeWeight: Node,
{
    let mut addrs = HashSet::new();
    visit(
        get_node,
        g,
        &[],
        &mut visit::RequiredAddrs { addrs: &mut addrs },
    );
    addrs
}

/// Extract a subgraph containing only the selected nodes and edges between them.
///
/// Nodes are visited in index order for determinism. Callers that need to
/// correlate old/new indices can iterate selected nodes in the same sorted
/// order alongside `subgraph.node_indices()`.
pub fn extract_subgraph<N: Clone>(
    graph: &node::graph::Graph<N>,
    selected: &HashSet<node::graph::NodeIx>,
) -> node::graph::Graph<N> {
    let mut subgraph = node::graph::Graph::default();
    let sorted: BTreeSet<_> = selected.iter().copied().collect();
    let mut ix_map = HashMap::new();
    for old_ix in &sorted {
        let weight = graph[*old_ix].clone();
        let new_ix = subgraph.add_node(weight);
        ix_map.insert(*old_ix, new_ix);
    }
    for e in graph.edge_references() {
        let src = e.source();
        let tgt = e.target();
        if let (Some(&new_src), Some(&new_tgt)) = (ix_map.get(&src), ix_map.get(&tgt)) {
            subgraph.add_edge(new_src, new_tgt, e.weight().clone());
        }
    }
    subgraph
}

/// Add all nodes and edges from `subgraph` into `target`.
///
/// Returns the new node indices in `target` in the same order as
/// `subgraph.node_indices()`, so callers can correlate them with positions.
pub fn add_subgraph<N: Clone>(
    target: &mut node::graph::Graph<N>,
    subgraph: &node::graph::Graph<N>,
) -> Vec<node::graph::NodeIx> {
    let mut ix_map = HashMap::new();
    let mut new_indices = Vec::new();
    for n in subgraph.node_indices() {
        let weight = subgraph[n].clone();
        let new_ix = target.add_node(weight);
        ix_map.insert(n, new_ix);
        new_indices.push(new_ix);
    }
    for e in subgraph.edge_references() {
        let new_src = ix_map[&e.source()];
        let new_tgt = ix_map[&e.target()];
        target.add_edge(new_src, new_tgt, e.weight().clone());
    }
    new_indices
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::Edge;

    /// Build a diamond graph: A -> B, A -> C, B -> D, C -> D.
    fn diamond_graph() -> (node::graph::Graph<&'static str>, [node::graph::NodeIx; 4]) {
        let mut g = node::graph::Graph::default();
        let a = g.add_node("A");
        let b = g.add_node("B");
        let c = g.add_node("C");
        let d = g.add_node("D");
        g.add_edge(a, b, Edge::new(0.into(), 0.into()));
        g.add_edge(a, c, Edge::new(0.into(), 0.into()));
        g.add_edge(b, d, Edge::new(0.into(), 0.into()));
        g.add_edge(c, d, Edge::new(0.into(), 1.into()));
        (g, [a, b, c, d])
    }

    #[test]
    fn extract_subgraph_basic() {
        let (g, [_a, b, c, d]) = diamond_graph();
        let selected: HashSet<_> = [b, c, d].into_iter().collect();
        let sub = extract_subgraph(&g, &selected);
        assert_eq!(sub.node_count(), 3);
        // Only edges where both endpoints are selected: B->D, C->D.
        assert_eq!(sub.edge_count(), 2);
        // Weights preserved.
        let weights: Vec<_> = sub.node_indices().map(|n| sub[n]).collect();
        assert_eq!(weights, vec!["B", "C", "D"]);
    }

    #[test]
    fn extract_subgraph_excludes_external_edges() {
        let (g, [a, _b, _c, d]) = diamond_graph();
        // Select only A and D — no direct edge between them.
        let selected: HashSet<_> = [a, d].into_iter().collect();
        let sub = extract_subgraph(&g, &selected);
        assert_eq!(sub.node_count(), 2);
        assert_eq!(sub.edge_count(), 0);
    }

    #[test]
    fn extract_subgraph_empty_selection() {
        let (g, _) = diamond_graph();
        let sub = extract_subgraph(&g, &HashSet::new());
        assert_eq!(sub.node_count(), 0);
        assert_eq!(sub.edge_count(), 0);
    }

    #[test]
    fn add_subgraph_returns_correct_indices() {
        let mut target = node::graph::Graph::<&str>::default();
        let _existing = target.add_node("X");

        let mut sub = node::graph::Graph::default();
        let sa = sub.add_node("A");
        let sb = sub.add_node("B");
        sub.add_edge(sa, sb, Edge::new(0.into(), 0.into()));

        let new = add_subgraph(&mut target, &sub);
        assert_eq!(new.len(), 2);
        assert_eq!(target.node_count(), 3);
        assert_eq!(target.edge_count(), 1);
        // New indices should be after the existing node.
        assert_eq!(target[new[0]], "A");
        assert_eq!(target[new[1]], "B");
        // Edge should connect the new nodes.
        assert!(target.find_edge(new[0], new[1]).is_some());
    }
}