weavegraph 0.7.0

//! Graph iteration utilities and topological ordering.

use crate::types::NodeKind;
use rustc_hash::FxHashMap;
use std::collections::VecDeque;

/// Iterator over node kinds registered in a graph.
///
/// Yields each custom node kind. Virtual `Start` and `End` nodes are not
/// included — they are not stored in the node registry.
///
/// # Examples
///
/// ```
/// use weavegraph::graphs::GraphBuilder;
/// use weavegraph::types::NodeKind;
///
/// # struct MyNode;
/// # #[async_trait::async_trait]
/// # impl weavegraph::node::Node for MyNode {
/// #     async fn run(&self, _: weavegraph::state::StateSnapshot, _: weavegraph::node::NodeContext) -> Result<weavegraph::node::NodePartial, weavegraph::node::NodeError> {
/// #         Ok(weavegraph::node::NodePartial::default())
/// #     }
/// # }
/// let builder = GraphBuilder::new()
///     .add_node(NodeKind::Custom("A".into()), MyNode)
///     .add_node(NodeKind::Custom("B".into()), MyNode)
///     .add_edge(NodeKind::Start, NodeKind::Custom("A".into()))
///     .add_edge(NodeKind::Custom("A".into()), NodeKind::Custom("B".into()))
///     .add_edge(NodeKind::Custom("B".into()), NodeKind::End);
///
/// assert_eq!(builder.nodes().count(), 2);
/// ```
pub struct NodesIter<'a> {
    inner: std::collections::hash_map::Keys<'a, NodeKind, std::sync::Arc<dyn crate::node::Node>>,
}

impl<'a> NodesIter<'a> {
    pub(super) fn new(
        inner: std::collections::hash_map::Keys<
            'a,
            NodeKind,
            std::sync::Arc<dyn crate::node::Node>,
        >,
    ) -> Self {
        Self { inner }
    }
}

impl<'a> Iterator for NodesIter<'a> {
    type Item = &'a NodeKind;

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next()
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.inner.size_hint()
    }
}

impl ExactSizeIterator for NodesIter<'_> {}

/// Iterator over edges as `(source, target)` pairs.
///
/// Yields every edge including edges from/to virtual `Start` and `End`
/// nodes. Iteration order is not guaranteed to be deterministic.
///
/// # Examples
///
/// ```
/// use weavegraph::graphs::GraphBuilder;
/// use weavegraph::types::NodeKind;
///
/// # struct MyNode;
/// # #[async_trait::async_trait]
/// # impl weavegraph::node::Node for MyNode {
/// #     async fn run(&self, _: weavegraph::state::StateSnapshot, _: weavegraph::node::NodeContext) -> Result<weavegraph::node::NodePartial, weavegraph::node::NodeError> {
/// #         Ok(weavegraph::node::NodePartial::default())
/// #     }
/// # }
/// let builder = GraphBuilder::new()
///     .add_node(NodeKind::Custom("A".into()), MyNode)
///     .add_edge(NodeKind::Start, NodeKind::Custom("A".into()))
///     .add_edge(NodeKind::Custom("A".into()), NodeKind::End);
///
/// assert_eq!(builder.edges().count(), 2);
/// ```
pub struct EdgesIter<'a> {
    outer: std::collections::hash_map::Iter<'a, NodeKind, Vec<NodeKind>>,
    current_from: Option<&'a NodeKind>,
    current_targets: std::slice::Iter<'a, NodeKind>,
}

impl<'a> EdgesIter<'a> {
    pub(super) fn new(edges: &'a FxHashMap<NodeKind, Vec<NodeKind>>) -> Self {
        let mut outer = edges.iter();
        let (current_from, current_targets) = match outer.next() {
            Some((from, targets)) => (Some(from), targets.iter()),
            None => (None, [].iter()),
        };
        Self {
            outer,
            current_from,
            current_targets,
        }
    }
}

impl<'a> Iterator for EdgesIter<'a> {
    type Item = (&'a NodeKind, &'a NodeKind);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(to) = self.current_targets.next() {
                return Some((self.current_from.unwrap(), to));
            }
            let (from, targets) = self.outer.next()?;
            self.current_from = Some(from);
            self.current_targets = targets.iter();
        }
    }
}

/// Topological ordering via Kahn's algorithm.
///
/// Returns all nodes with dependencies before dependents. `Start` is always
/// first, `End` always last; ties among custom nodes are broken
/// lexicographically for determinism.
///
/// # Panics
///
/// Assumes an acyclic graph. On a cyclic graph the result is a partial
/// ordering that excludes cycle members. Use [`GraphBuilder::compile`] to
/// validate acyclicity first.
pub(super) fn topological_sort(edges: &FxHashMap<NodeKind, Vec<NodeKind>>) -> Vec<NodeKind> {
    let mut in_degree: FxHashMap<NodeKind, usize> = FxHashMap::default();
    for (from, targets) in edges {
        in_degree.entry(from.clone()).or_insert(0);
        for to in targets {
            *in_degree.entry(to.clone()).or_insert(0) += 1;
        }
    }

    let mut seeds: Vec<NodeKind> = in_degree
        .iter()
        .filter(|&(_, &deg)| deg == 0)
        .map(|(node, _)| node.clone())
        .collect();
    seeds.sort_by(node_order);
    let mut queue = VecDeque::from(seeds);

    let mut result = Vec::with_capacity(in_degree.len());
    while let Some(node) = queue.pop_front() {
        if let Some(targets) = edges.get(&node) {
            let mut newly_free: Vec<NodeKind> = targets
                .iter()
                .filter_map(|t| {
                    let deg = in_degree.get_mut(t)?;
                    *deg = deg.saturating_sub(1);
                    (*deg == 0).then(|| t.clone())
                })
                .collect();
            newly_free.sort_by(node_order);
            queue.extend(newly_free);
        }
        result.push(node);
    }

    result
}

fn node_order(a: &NodeKind, b: &NodeKind) -> std::cmp::Ordering {
    match (a, b) {
        (NodeKind::Start, _) => std::cmp::Ordering::Less,
        (_, NodeKind::Start) => std::cmp::Ordering::Greater,
        (NodeKind::End, _) => std::cmp::Ordering::Greater,
        (_, NodeKind::End) => std::cmp::Ordering::Less,
        (NodeKind::Custom(a), NodeKind::Custom(b)) => a.cmp(b),
    }
}

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

    #[test]
    fn test_topological_sort_linear() {
        let mut edges: FxHashMap<NodeKind, Vec<NodeKind>> = FxHashMap::default();
        edges.insert(NodeKind::Start, vec![NodeKind::Custom("A".into())]);
        edges.insert(
            NodeKind::Custom("A".into()),
            vec![NodeKind::Custom("B".into())],
        );
        edges.insert(NodeKind::Custom("B".into()), vec![NodeKind::End]);

        let sorted = topological_sort(&edges);

        assert_eq!(sorted[0], NodeKind::Start);
        assert_eq!(*sorted.last().unwrap(), NodeKind::End);

        let a_pos = sorted
            .iter()
            .position(|n| n == &NodeKind::Custom("A".into()))
            .unwrap();
        let b_pos = sorted
            .iter()
            .position(|n| n == &NodeKind::Custom("B".into()))
            .unwrap();
        assert!(a_pos < b_pos);
    }

    #[test]
    fn test_topological_sort_diamond() {
        let mut edges: FxHashMap<NodeKind, Vec<NodeKind>> = FxHashMap::default();
        edges.insert(
            NodeKind::Start,
            vec![NodeKind::Custom("A".into()), NodeKind::Custom("B".into())],
        );
        edges.insert(
            NodeKind::Custom("A".into()),
            vec![NodeKind::Custom("C".into())],
        );
        edges.insert(
            NodeKind::Custom("B".into()),
            vec![NodeKind::Custom("C".into())],
        );
        edges.insert(NodeKind::Custom("C".into()), vec![NodeKind::End]);

        let sorted = topological_sort(&edges);

        assert_eq!(sorted[0], NodeKind::Start);
        assert_eq!(*sorted.last().unwrap(), NodeKind::End);

        let a_pos = sorted
            .iter()
            .position(|n| n == &NodeKind::Custom("A".into()))
            .unwrap();
        let b_pos = sorted
            .iter()
            .position(|n| n == &NodeKind::Custom("B".into()))
            .unwrap();
        let c_pos = sorted
            .iter()
            .position(|n| n == &NodeKind::Custom("C".into()))
            .unwrap();
        assert!(a_pos < c_pos);
        assert!(b_pos < c_pos);
        assert!(a_pos < b_pos);
    }

    #[test]
    fn test_topological_sort_deterministic() {
        let mut edges: FxHashMap<NodeKind, Vec<NodeKind>> = FxHashMap::default();
        edges.insert(
            NodeKind::Start,
            vec![
                NodeKind::Custom("X".into()),
                NodeKind::Custom("Y".into()),
                NodeKind::Custom("Z".into()),
            ],
        );
        edges.insert(NodeKind::Custom("X".into()), vec![NodeKind::End]);
        edges.insert(NodeKind::Custom("Y".into()), vec![NodeKind::End]);
        edges.insert(NodeKind::Custom("Z".into()), vec![NodeKind::End]);

        let sorted1 = topological_sort(&edges);
        let sorted2 = topological_sort(&edges);

        assert_eq!(sorted1, sorted2);
    }
}