dsa64 0.1.2

//! # Directed Acyclic Graph
//!
//! A directed acyclic graph (DAG) generic over the node key type `K`.
//!
//! Forward and reverse adjacency relationships are maintained in parallel
//! `HashMap<K, HashSet<K>>` maps, providing amortised O(1) lookup of nodes,
//! edges, direct dependents, and direct dependencies.
//!
//! # Edge direction
//!
//! An edge `A -> B` means that `B` depends on `A`. Consequently:
//!
//! - `dependents(A)` returns the nodes that directly depend on `A`.
//! - `dependencies(A)` returns the nodes that `A` directly depends on.
//! - `transitive_dependents(A)` returns all nodes that directly or indirectly
//!   depend on `A`.
//! - `transitive_dependencies(A)` returns all nodes that `A` directly or
//!   indirectly depends on.
//!
//! # Cycle prevention
//!
//! `add_edge` preserves the acyclic invariant by checking whether the source
//! node is already reachable from the target before inserting an edge. If it
//! is reachable, adding the edge would create a cycle and `CycleError` is
//! returned. Self-loops are rejected directly.
//!
//! Reachability and transitive queries use breadth-first traversal.

use std::collections::{HashMap, HashSet, VecDeque};
use std::fmt;
use std::hash::Hash;

/// Error returned when adding an edge would introduce a cycle.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CycleError;

impl fmt::Display for CycleError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("edge would create a cycle")
    }
}

impl std::error::Error for CycleError {}

/// A directed acyclic graph with forward and reverse adjacency indexes.
///
/// An edge `A -> B` indicates that `B` depends on `A`. Forward adjacency
/// therefore stores dependents, while reverse adjacency stores dependencies.
///
/// Both adjacency maps contain every node in the graph, including nodes with
/// no incoming or outgoing edges. The maps are updated together to preserve
/// this invariant.
///
/// Direct node and edge lookups are amortised O(1). Adding an edge may require
/// traversing the graph to verify that the new edge would not introduce a
/// cycle. Transitive queries are linear in the portion of the graph traversed.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Dag<K: Hash + Eq + Clone> {
    /// Maps each node to its direct dependents.
    forward: HashMap<K, HashSet<K>>,

    /// Maps each node to its direct dependencies.
    reverse: HashMap<K, HashSet<K>>,

    /// Number of directed edges currently stored in the graph.
    edge_count: usize,
}

impl<K: Hash + Eq + Clone> Dag<K> {
    /// Creates an empty DAG.
    #[inline]
    pub fn new() -> Self {
        Self {
            forward: HashMap::new(),
            reverse: HashMap::new(),
            edge_count: 0,
        }
    }

    /// Returns the number of nodes in the graph.
    #[inline]
    pub fn node_count(&self) -> usize {
        self.forward.len()
    }

    /// Returns the number of directed edges in the graph.
    #[inline]
    pub fn edge_count(&self) -> usize {
        self.edge_count
    }

    /// Returns `true` if the graph contains no nodes.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.forward.is_empty()
    }

    /// Returns `true` if the graph contains `key`.
    #[inline]
    pub fn contains_node(&self, key: &K) -> bool {
        self.forward.contains_key(key)
    }

    /// Returns `true` if the graph contains the directed edge `from -> to`.
    #[inline]
    pub fn contains_edge(&self, from: &K, to: &K) -> bool {
        self.forward.get(from).is_some_and(|s| s.contains(to))
    }

    /// Inserts a node into the graph.
    ///
    /// Returns `true` if the node was inserted, or `false` if it was already
    /// present.
    #[inline]
    pub fn insert_node(&mut self, key: K) -> bool {
        if self.forward.contains_key(&key) {
            return false;
        }

        self.forward.insert(key.clone(), HashSet::new());
        self.reverse.insert(key, HashSet::new());
        true
    }

    /// Adds the directed edge `from -> to`.
    ///
    /// An edge `from -> to` indicates that `to` depends on `from`. Missing
    /// endpoint nodes are inserted automatically.
    ///
    /// Returns `Ok(true)` if the edge was inserted or `Ok(false)` if the edge
    /// already exists.
    ///
    /// # Errors
    ///
    /// Returns [`CycleError`] if the edge would introduce a cycle. Self-loops
    /// are considered cycles and are rejected.
    pub fn add_edge(&mut self, from: K, to: K) -> Result<bool, CycleError> {
        if from == to {
            return Err(CycleError);
        }

        if self.contains_edge(&from, &to) {
            return Ok(false);
        }

        // Adding `from -> to` creates a cycle exactly when `from` is already
        // reachable from `to`.
        if self.is_reachable(&to, &from) {
            return Err(CycleError);
        }

        self.insert_node(from.clone());
        self.insert_node(to.clone());

        self.forward.get_mut(&from).unwrap().insert(to.clone());
        self.reverse.get_mut(&to).unwrap().insert(from);
        self.edge_count += 1;

        Ok(true)
    }

    /// Removes a node and every edge incident to it.
    ///
    /// Returns `true` if the node existed, or `false` if it was not present.
    pub fn remove_node(&mut self, key: &K) -> bool {
        let Some(dependents) = self.forward.remove(key) else {
            return false;
        };

        // Remove each outgoing edge `key -> dependent` from the corresponding
        // reverse adjacency set.
        self.edge_count -= dependents.len();
        for dep in &dependents {
            if let Some(rev) = self.reverse.get_mut(dep) {
                rev.remove(key);
            }
        }

        // Remove each incoming edge `dependency -> key` from the corresponding
        // forward adjacency set.
        if let Some(dependencies) = self.reverse.remove(key) {
            self.edge_count -= dependencies.len();
            for dep in &dependencies {
                if let Some(fwd) = self.forward.get_mut(dep) {
                    fwd.remove(key);
                }
            }
        }

        true
    }

    /// Removes the directed edge `from -> to`.
    ///
    /// Returns `true` if the edge existed, or `false` otherwise. Nodes are
    /// retained even if removing the edge leaves them isolated.
    #[inline]
    pub fn remove_edge(&mut self, from: &K, to: &K) -> bool {
        let removed = self
            .forward
            .get_mut(from)
            .is_some_and(|s| s.remove(to));

        if removed {
            self.reverse.get_mut(to).unwrap().remove(from);
            self.edge_count -= 1;
        }

        removed
    }

    /// Returns the direct dependents of `key`.
    ///
    /// These are the nodes `D` for which an edge `key -> D` exists. Returns
    /// `None` if `key` is not present in the graph.
    #[inline]
    pub fn dependents(&self, key: &K) -> Option<&HashSet<K>> {
        self.forward.get(key)
    }

    /// Returns the direct dependencies of `key`.
    ///
    /// These are the nodes `D` for which an edge `D -> key` exists. Returns
    /// `None` if `key` is not present in the graph.
    #[inline]
    pub fn dependencies(&self, key: &K) -> Option<&HashSet<K>> {
        self.reverse.get(key)
    }

    /// Returns all direct and indirect dependents of `key`.
    ///
    /// The starting node is not included. An empty set is returned if `key`
    /// does not exist or has no dependents.
    pub fn transitive_dependents(&self, key: &K) -> HashSet<K> {
        self.walk(&self.forward, key)
    }

    /// Returns all direct and indirect dependencies of `key`.
    ///
    /// The starting node is not included. An empty set is returned if `key`
    /// does not exist or has no dependencies.
    pub fn transitive_dependencies(&self, key: &K) -> HashSet<K> {
        self.walk(&self.reverse, key)
    }

    /// Traverses an adjacency map breadth-first and returns all nodes reachable
    /// from `start`.
    ///
    /// `start` itself is not included in the returned set.
    fn walk(&self, adj: &HashMap<K, HashSet<K>>, start: &K) -> HashSet<K> {
        let mut result = HashSet::new();
        let mut queue = VecDeque::new();

        if let Some(neighbours) = adj.get(start) {
            for n in neighbours {
                if result.insert(n.clone()) {
                    queue.push_back(n);
                }
            }
        }

        while let Some(current) = queue.pop_front() {
            if let Some(neighbours) = adj.get(current) {
                for n in neighbours {
                    if result.insert(n.clone()) {
                        queue.push_back(n);
                    }
                }
            }
        }

        result
    }

    /// Returns `true` if `target` is reachable from `start` by following
    /// forward edges.
    ///
    /// The traversal stores references to node keys rather than cloning them.
    fn is_reachable(&self, start: &K, target: &K) -> bool {
        let Some(neighbours) = self.forward.get(start) else {
            return false;
        };

        let mut visited = HashSet::new();
        let mut queue = VecDeque::new();

        for n in neighbours {
            if n == target {
                return true;
            }

            if visited.insert(n) {
                queue.push_back(n);
            }
        }

        while let Some(current) = queue.pop_front() {
            if let Some(next) = self.forward.get(current) {
                for n in next {
                    if n == target {
                        return true;
                    }

                    if visited.insert(n) {
                        queue.push_back(n);
                    }
                }
            }
        }

        false
    }
}

impl<K: Hash + Eq + Clone> Default for Dag<K> {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn empty_graph() {
        let dag = Dag::<u32>::new();
        assert!(dag.is_empty());
        assert_eq!(dag.node_count(), 0);
        assert_eq!(dag.edge_count(), 0);
        assert!(!dag.contains_node(&1));
        assert!(!dag.contains_edge(&1, &2));
    }

    #[test]
    fn insert_node() {
        let mut dag = Dag::new();
        assert!(dag.insert_node(1));
        assert!(!dag.insert_node(1));
        assert_eq!(dag.node_count(), 1);
        assert!(dag.contains_node(&1));
        assert!(!dag.contains_node(&2));
    }

    #[test]
    fn add_edge() {
        let mut dag = Dag::new();
        assert_eq!(dag.add_edge(1, 2), Ok(true));
        assert_eq!(dag.add_edge(1, 2), Ok(false));
        assert_eq!(dag.node_count(), 2);
        assert_eq!(dag.edge_count(), 1);
        assert!(dag.contains_edge(&1, &2));
        assert!(!dag.contains_edge(&2, &1));
    }

    #[test]
    fn add_edge_auto_inserts_nodes() {
        let mut dag = Dag::new();
        dag.add_edge(10, 20).unwrap();
        assert!(dag.contains_node(&10));
        assert!(dag.contains_node(&20));
    }

    #[test]
    fn self_loop_rejected() {
        let mut dag = Dag::new();
        assert_eq!(dag.add_edge(1, 1), Err(CycleError));
        assert_eq!(dag.edge_count(), 0);
    }

    #[test]
    fn direct_cycle_rejected() {
        let mut dag = Dag::new();
        dag.add_edge(1, 2).unwrap();
        assert_eq!(dag.add_edge(2, 1), Err(CycleError));
        assert_eq!(dag.edge_count(), 1);
    }

    #[test]
    fn transitive_cycle_rejected() {
        let mut dag = Dag::new();
        dag.add_edge(1, 2).unwrap();
        dag.add_edge(2, 3).unwrap();
        dag.add_edge(3, 4).unwrap();
        assert_eq!(dag.add_edge(4, 1), Err(CycleError));
        assert_eq!(dag.add_edge(4, 2), Err(CycleError));
        assert_eq!(dag.edge_count(), 3);
    }

    #[test]
    fn diamond_is_valid() {
        //   1
        //  / \
        // 2   3
        //  \ /
        //   4
        let mut dag = Dag::new();
        dag.add_edge(1, 2).unwrap();
        dag.add_edge(1, 3).unwrap();
        dag.add_edge(2, 4).unwrap();
        dag.add_edge(3, 4).unwrap();
        assert_eq!(dag.node_count(), 4);
        assert_eq!(dag.edge_count(), 4);

        // Closing the diamond into a cycle is rejected.
        assert_eq!(dag.add_edge(4, 1), Err(CycleError));
    }

    #[test]
    fn remove_node() {
        let mut dag = Dag::new();
        dag.add_edge(1, 2).unwrap();
        dag.add_edge(2, 3).unwrap();
        dag.add_edge(1, 3).unwrap();

        assert!(dag.remove_node(&2));
        assert!(!dag.contains_node(&2));
        assert_eq!(dag.node_count(), 2);
        assert_eq!(dag.edge_count(), 1);
        assert!(dag.contains_edge(&1, &3));
        assert!(!dag.contains_edge(&1, &2));
        assert!(!dag.contains_edge(&2, &3));
    }

    #[test]
    fn remove_node_nonexistent() {
        let mut dag = Dag::<u32>::new();
        assert!(!dag.remove_node(&99));
    }

    #[test]
    fn remove_edge() {
        let mut dag = Dag::new();
        dag.add_edge(1, 2).unwrap();
        dag.add_edge(2, 3).unwrap();

        assert!(dag.remove_edge(&1, &2));
        assert!(!dag.contains_edge(&1, &2));
        assert_eq!(dag.edge_count(), 1);
        assert!(dag.contains_node(&1));
        assert!(dag.contains_node(&2));
    }

    #[test]
    fn remove_edge_nonexistent() {
        let mut dag = Dag::<u32>::new();
        assert!(!dag.remove_edge(&1, &2));
    }

    #[test]
    fn remove_edge_enables_previously_cyclic_edge() {
        let mut dag = Dag::new();
        dag.add_edge(1, 2).unwrap();
        dag.add_edge(2, 3).unwrap();

        // 3 -> 1 would create a cycle through 1 -> 2 -> 3 -> 1.
        assert_eq!(dag.add_edge(3, 1), Err(CycleError));

        // Removing 1 -> 2 breaks the path, so 3 -> 1 is now valid.
        dag.remove_edge(&1, &2);
        assert_eq!(dag.add_edge(3, 1), Ok(true));
    }

    #[test]
    fn direct_dependents() {
        let mut dag = Dag::new();
        dag.add_edge(1, 2).unwrap();
        dag.add_edge(1, 3).unwrap();

        let deps = dag.dependents(&1).unwrap();
        assert_eq!(deps.len(), 2);
        assert!(deps.contains(&2));
        assert!(deps.contains(&3));
    }

    #[test]
    fn direct_dependencies() {
        let mut dag = Dag::new();
        dag.add_edge(1, 3).unwrap();
        dag.add_edge(2, 3).unwrap();

        let deps = dag.dependencies(&3).unwrap();
        assert_eq!(deps.len(), 2);
        assert!(deps.contains(&1));
        assert!(deps.contains(&2));
    }

    #[test]
    fn dependents_of_unknown_node() {
        let dag = Dag::<u32>::new();
        assert!(dag.dependents(&1).is_none());
    }

    #[test]
    fn dependencies_of_unknown_node() {
        let dag = Dag::<u32>::new();
        assert!(dag.dependencies(&1).is_none());
    }

    #[test]
    fn dependents_of_leaf_node() {
        let mut dag = Dag::new();
        dag.add_edge(1, 2).unwrap();
        let deps = dag.dependents(&2).unwrap();
        assert!(deps.is_empty());
    }

    #[test]
    fn transitive_dependents_chain() {
        // 1 -> 2 -> 3 -> 4
        let mut dag = Dag::new();
        dag.add_edge(1, 2).unwrap();
        dag.add_edge(2, 3).unwrap();
        dag.add_edge(3, 4).unwrap();

        let td = dag.transitive_dependents(&1);
        assert_eq!(td, HashSet::from([2, 3, 4]));
    }

    #[test]
    fn transitive_dependencies_chain() {
        // 1 -> 2 -> 3 -> 4
        let mut dag = Dag::new();
        dag.add_edge(1, 2).unwrap();
        dag.add_edge(2, 3).unwrap();
        dag.add_edge(3, 4).unwrap();

        let td = dag.transitive_dependencies(&4);
        assert_eq!(td, HashSet::from([1, 2, 3]));
    }

    #[test]
    fn transitive_dependents_diamond() {
        //   1
        //  / \
        // 2   3
        //  \ /
        //   4
        let mut dag = Dag::new();
        dag.add_edge(1, 2).unwrap();
        dag.add_edge(1, 3).unwrap();
        dag.add_edge(2, 4).unwrap();
        dag.add_edge(3, 4).unwrap();

        let td = dag.transitive_dependents(&1);
        assert_eq!(td, HashSet::from([2, 3, 4]));

        let td2 = dag.transitive_dependents(&2);
        assert_eq!(td2, HashSet::from([4]));
    }

    #[test]
    fn transitive_dependencies_diamond() {
        //   1
        //  / \
        // 2   3
        //  \ /
        //   4
        let mut dag = Dag::new();
        dag.add_edge(1, 2).unwrap();
        dag.add_edge(1, 3).unwrap();
        dag.add_edge(2, 4).unwrap();
        dag.add_edge(3, 4).unwrap();

        let td = dag.transitive_dependencies(&4);
        assert_eq!(td, HashSet::from([1, 2, 3]));

        let td2 = dag.transitive_dependencies(&2);
        assert_eq!(td2, HashSet::from([1]));
    }

    #[test]
    fn transitive_queries_on_empty_graph() {
        let dag = Dag::<u32>::new();
        assert!(dag.transitive_dependents(&1).is_empty());
        assert!(dag.transitive_dependencies(&1).is_empty());
    }

    #[test]
    fn transitive_queries_on_isolated_node() {
        let mut dag = Dag::new();
        dag.insert_node(1);
        assert!(dag.transitive_dependents(&1).is_empty());
        assert!(dag.transitive_dependencies(&1).is_empty());
    }

    #[test]
    fn string_keys() {
        let mut dag = Dag::new();
        dag.add_edge("a".to_string(), "b".to_string()).unwrap();
        dag.add_edge("b".to_string(), "c".to_string()).unwrap();

        assert!(dag.contains_edge(&"a".to_string(), &"b".to_string()));
        let td = dag.transitive_dependents(&"a".to_string());
        assert_eq!(td.len(), 2);
        assert!(td.contains("b"));
        assert!(td.contains("c"));
    }

    #[test]
    fn wide_fan_out() {
        let mut dag = Dag::new();
        for i in 1..=100 {
            dag.add_edge(0u32, i).unwrap();
        }
        assert_eq!(dag.node_count(), 101);
        assert_eq!(dag.edge_count(), 100);

        let deps = dag.dependents(&0).unwrap();
        assert_eq!(deps.len(), 100);

        let td = dag.transitive_dependents(&0);
        assert_eq!(td.len(), 100);
    }

    #[test]
    fn wide_fan_in() {
        let mut dag = Dag::new();
        for i in 1..=100 {
            dag.add_edge(i, 0u32).unwrap();
        }

        let deps = dag.dependencies(&0).unwrap();
        assert_eq!(deps.len(), 100);

        let td = dag.transitive_dependencies(&0);
        assert_eq!(td.len(), 100);
    }

    #[test]
    fn default_is_empty() {
        let dag: Dag<u32> = Dag::default();
        assert!(dag.is_empty());
    }

    #[test]
    fn cycle_error_display() {
        assert_eq!(CycleError.to_string(), "edge would create a cycle");
    }

    #[test]
    fn equality() {
        let mut a = Dag::new();
        a.add_edge(1, 2).unwrap();
        a.add_edge(2, 3).unwrap();

        let mut b = Dag::new();
        b.add_edge(2, 3).unwrap();
        b.add_edge(1, 2).unwrap();

        assert_eq!(a, b);
    }
}