antecedent_graph/

dag.rs

//! Indexed DAG storage with acyclicity validation.
//!
//! SPDX-License-Identifier: MIT OR Apache-2.0

use std::sync::Arc;

use antecedent_core::VariableId;

use crate::algo::{bfs_reaches, kahn_order};
use crate::error::GraphError;
use crate::types::{DenseNodeId, MarkedEdge, NodeRef};
use crate::workspace::GraphWorkspace;

/// Static directed acyclic graph over variables.
#[derive(Clone, Debug)]
pub struct Dag {
    nodes: Vec<NodeRef>,
    /// Outgoing children per node.
    children: Vec<Vec<DenseNodeId>>,
    /// Incoming parents per node.
    parents: Vec<Vec<DenseNodeId>>,
    /// Reused by insertion-time acyclicity checks to avoid per-insert allocation.
    insert_ws: GraphWorkspace,
}

impl Dag {
    /// Empty DAG.
    #[must_use]
    pub fn empty() -> Self {
        Self {
            nodes: Vec::new(),
            children: Vec::new(),
            parents: Vec::new(),
            insert_ws: GraphWorkspace::default(),
        }
    }

    /// Build a DAG with one static node per variable `0..n`.
    #[must_use]
    pub fn with_variables(n: u32) -> Self {
        let mut g = Self::empty();
        for i in 0..n {
            let _ = g.add_node(NodeRef::Static(VariableId::from_raw(i)));
        }
        g
    }

    /// Build a DAG with one static node per schema variable (`VariableId` raw == dense id),
    /// then insert directed edges named by schema variable names.
    ///
    /// # Errors
    ///
    /// Unknown names, duplicate edges, or cycles.
    pub fn from_named_edges(
        schema: &antecedent_core::CausalSchema,
        edges: &[(&str, &str)],
    ) -> Result<Self, GraphError> {
        let n = crate::named::schema_node_count(schema)?;
        let mut g = Self::with_variables(n);
        for &(from_name, to_name) in edges {
            let (from, to) = crate::named::resolve_named_edge(schema, from_name, to_name)?;
            g.insert_directed(from, to)?;
        }
        Ok(g)
    }

    /// Number of nodes.
    #[must_use]
    pub fn node_count(&self) -> usize {
        self.nodes.len()
    }

    /// Whether empty.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.nodes.is_empty()
    }

    /// Node refs in dense order.
    #[must_use]
    pub fn nodes(&self) -> &[NodeRef] {
        &self.nodes
    }

    /// Add a static node; returns its dense id.
    ///
    /// # Errors
    ///
    /// [`GraphError::TooManyNodes`] on overflow.
    pub fn add_node(&mut self, node: NodeRef) -> Result<DenseNodeId, GraphError> {
        if !matches!(node, NodeRef::Static(_)) {
            return Err(GraphError::InvalidEndpoints { message: "Dag accepts only Static nodes" });
        }
        let id = u32::try_from(self.nodes.len()).map_err(|_| GraphError::TooManyNodes)?;
        self.nodes.push(node);
        self.children.push(Vec::new());
        self.parents.push(Vec::new());
        Ok(DenseNodeId::from_raw(id))
    }

    /// Insert a directed edge `from -> to` if it preserves acyclicity.
    ///
    /// # Errors
    ///
    /// Unknown nodes, duplicates, or cycles.
    pub fn insert_directed(
        &mut self,
        from: DenseNodeId,
        to: DenseNodeId,
    ) -> Result<(), GraphError> {
        self.validate_node(from)?;
        self.validate_node(to)?;
        if self.children[from.as_usize()].contains(&to) {
            return Err(GraphError::DuplicateEdge { from: from.raw(), to: to.raw() });
        }
        let mut ws = core::mem::take(&mut self.insert_ws);
        let cycle = self.reaches_with(to, from, &mut ws);
        self.insert_ws = ws;
        if cycle {
            return Err(GraphError::Cycle { from: from.raw(), to: to.raw() });
        }
        self.children[from.as_usize()].push(to);
        self.parents[to.as_usize()].push(from);
        Ok(())
    }

    /// Push an edge without duplicate/acyclicity checks; the caller guarantees
    /// both nodes exist and the edge preserves invariants.
    pub(crate) fn insert_directed_unchecked(&mut self, from: DenseNodeId, to: DenseNodeId) {
        self.children[from.as_usize()].push(to);
        self.parents[to.as_usize()].push(from);
    }

    /// Remove a directed edge if present.
    pub fn remove_directed(&mut self, from: DenseNodeId, to: DenseNodeId) {
        if from.as_usize() >= self.node_count() || to.as_usize() >= self.node_count() {
            return;
        }
        self.children[from.as_usize()].retain(|c| *c != to);
        self.parents[to.as_usize()].retain(|p| *p != from);
    }

    /// Children of `id`.
    #[must_use]
    pub fn children(&self, id: DenseNodeId) -> &[DenseNodeId] {
        &self.children[id.as_usize()]
    }

    /// Parents of `id`.
    #[must_use]
    pub fn parents(&self, id: DenseNodeId) -> &[DenseNodeId] {
        &self.parents[id.as_usize()]
    }

    /// Whether `from` can reach `to` via directed edges.
    #[must_use]
    pub fn reaches(&self, from: DenseNodeId, to: DenseNodeId) -> bool {
        if from == to {
            return true;
        }
        let mut ws = GraphWorkspace::default();
        self.reaches_with(from, to, &mut ws)
    }

    /// Reachability using a reusable workspace.
    pub fn reaches_with(
        &self,
        from: DenseNodeId,
        to: DenseNodeId,
        ws: &mut GraphWorkspace,
    ) -> bool {
        bfs_reaches(&self.children, from, to, ws)
    }

    /// Topological order (Kahn). Returns `None` if a cycle slipped in.
    #[must_use]
    pub fn topological_order(&self) -> Option<Vec<DenseNodeId>> {
        kahn_order(&self.parents, &self.children)
    }

    /// Validate graph invariants.
    ///
    /// # Errors
    ///
    /// Cycle detected.
    pub fn validate(&self) -> Result<(), GraphError> {
        if self.topological_order().is_none() {
            return Err(GraphError::Cycle { from: 0, to: 0 });
        }
        Ok(())
    }

    fn validate_node(&self, id: DenseNodeId) -> Result<(), GraphError> {
        if id.as_usize() >= self.node_count() {
            Err(GraphError::UnknownNode { id: id.raw() })
        } else {
            Ok(())
        }
    }

    /// Iterate directed edges as marked edges.
    pub fn edges(&self) -> impl Iterator<Item = MarkedEdge> + '_ {
        self.children.iter().enumerate().flat_map(|(i, kids)| {
            let from = DenseNodeId::from_raw(u32::try_from(i).expect("node fit"));
            kids.iter().map(move |&to| MarkedEdge::directed(from, to))
        })
    }

    /// Enumerate simple directed paths from `from` to `to` (inclusive endpoints).
    ///
    /// Bounded by `max_paths` and `max_len` (number of nodes on the path).
    ///
    /// The returned set is silently truncated when either bound binds. Callers whose
    /// correctness depends on seeing *every* path — e.g. a recanting-witness test, which
    /// concludes "no witness exists" from the absence of one — must use
    /// [`Self::directed_paths_with_budget`] and fail closed on truncation instead.
    ///
    /// # Errors
    ///
    /// Unknown nodes.
    pub fn directed_paths(
        &self,
        from: DenseNodeId,
        to: DenseNodeId,
        max_paths: usize,
        max_len: usize,
    ) -> Result<Vec<Vec<DenseNodeId>>, GraphError> {
        self.directed_paths_with_budget(from, to, max_paths, max_len).map(|(paths, _)| paths)
    }

    /// [`Self::directed_paths`] plus a flag reporting whether enumeration was cut short.
    ///
    /// The flag is `true` when `max_paths` stopped the search with candidates still
    /// pending, or when `max_len` pruned a partial path that had not yet reached `to`.
    /// It is deliberately conservative: `true` means "the path set may be incomplete",
    /// never "it is definitely incomplete".
    ///
    /// # Errors
    ///
    /// Unknown nodes.
    pub fn directed_paths_with_budget(
        &self,
        from: DenseNodeId,
        to: DenseNodeId,
        max_paths: usize,
        max_len: usize,
    ) -> Result<(Vec<Vec<DenseNodeId>>, bool), GraphError> {
        self.validate_node(from)?;
        self.validate_node(to)?;
        let mut out = Vec::new();
        if max_paths == 0 || max_len == 0 {
            // A zero budget cannot certify that no path exists.
            return Ok((out, true));
        }
        let mut truncated = false;
        let mut stack = vec![vec![from]];
        while let Some(path) = stack.pop() {
            if out.len() >= max_paths {
                truncated = true;
                break;
            }
            let last = *path.last().expect("nonempty");
            if path.len() > 1 && last == to {
                out.push(path);
                continue;
            }
            if last == to && path.len() == 1 {
                out.push(path);
                continue;
            }
            if path.len() >= max_len {
                // Pruned before reaching `to`; a longer completion may exist.
                truncated = true;
                continue;
            }
            for &c in self.children(last) {
                if path.contains(&c) {
                    continue;
                }
                let mut next = path.clone();
                next.push(c);
                stack.push(next);
            }
        }
        Ok((out, truncated))
    }
}

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

    #[test]
    fn directed_paths_reports_max_paths_truncation() {
        // t→c→y, t→w→b→y, t→w→a→y : three directed paths.
        let mut g = Dag::with_variables(6);
        for (u, v) in [(0, 4), (4, 5), (0, 1), (1, 3), (3, 5), (1, 2), (2, 5)] {
            g.insert_directed(DenseNodeId::from_raw(u), DenseNodeId::from_raw(v)).unwrap();
        }
        let (t, y) = (DenseNodeId::from_raw(0), DenseNodeId::from_raw(5));

        let (all, truncated) = g.directed_paths_with_budget(t, y, 64, 16).unwrap();
        assert_eq!(all.len(), 3);
        assert!(!truncated, "a budget that comfortably fits every path must not report truncation");

        for cap in 1..=2 {
            let (paths, truncated) = g.directed_paths_with_budget(t, y, cap, 16).unwrap();
            assert_eq!(paths.len(), cap);
            assert!(truncated, "max_paths={cap} dropped paths but reported none");
        }
    }

    #[test]
    fn directed_paths_reports_max_len_truncation() {
        // Single path 0→1→2→3 needs 4 nodes; max_len=3 prunes it before it reaches the target.
        let mut g = Dag::with_variables(4);
        for (u, v) in [(0, 1), (1, 2), (2, 3)] {
            g.insert_directed(DenseNodeId::from_raw(u), DenseNodeId::from_raw(v)).unwrap();
        }
        let (t, y) = (DenseNodeId::from_raw(0), DenseNodeId::from_raw(3));
        let (paths, truncated) = g.directed_paths_with_budget(t, y, 64, 3).unwrap();
        assert!(paths.is_empty());
        assert!(truncated, "max_len pruned the only path but reported no truncation");

        let (paths, truncated) = g.directed_paths_with_budget(t, y, 64, 4).unwrap();
        assert_eq!(paths.len(), 1);
        assert!(!truncated);
    }

    #[test]
    fn rejects_cycles() {
        let mut g = Dag::with_variables(3);
        let a = DenseNodeId::from_raw(0);
        let b = DenseNodeId::from_raw(1);
        let c = DenseNodeId::from_raw(2);
        g.insert_directed(a, b).unwrap();
        g.insert_directed(b, c).unwrap();
        assert!(matches!(g.insert_directed(c, a), Err(GraphError::Cycle { .. })));
    }

    #[test]
    fn topological_order_respects_edges() {
        let mut g = Dag::with_variables(3);
        g.insert_directed(DenseNodeId::from_raw(0), DenseNodeId::from_raw(1)).unwrap();
        g.insert_directed(DenseNodeId::from_raw(1), DenseNodeId::from_raw(2)).unwrap();
        let order = g.topological_order().unwrap();
        let pos = |id: u32| order.iter().position(|n| n.raw() == id).unwrap();
        assert!(pos(0) < pos(1) && pos(1) < pos(2));
    }

    #[test]
    fn traversal_workspace_reuses_frontier_capacity() {
        let mut dag = Dag::with_variables(1_000);
        for i in 0..999 {
            dag.insert_directed(DenseNodeId::from_raw(i), DenseNodeId::from_raw(i + 1)).unwrap();
        }
        let mut ws = GraphWorkspace::default();
        assert!(dag.reaches_with(DenseNodeId::from_raw(0), DenseNodeId::from_raw(999), &mut ws));
        let ptr = ws.frontier.as_ptr();
        let cap = ws.frontier.capacity();
        for _ in 0..50 {
            assert!(dag.reaches_with(
                DenseNodeId::from_raw(0),
                DenseNodeId::from_raw(999),
                &mut ws
            ));
            assert_eq!(ws.frontier.as_ptr(), ptr);
            assert_eq!(ws.frontier.capacity(), cap);
        }
    }
}

/// Review-required static DAG artifact (`DirectLiNGAM` and other full-DAG discovery).
#[derive(Clone, Debug)]
pub struct DagReview {
    /// Proposed discovery DAG.
    pub graph: Dag,
    /// Directed edges awaiting explicit acceptance `(from, to)` as [`VariableId`]s.
    pub pending_edges: Arc<[(VariableId, VariableId)]>,
    /// Algorithm id that produced the proposal.
    pub algorithm: Arc<str>,
}

impl DagReview {
    /// Construct a review listing all current edges as pending.
    #[must_use]
    pub fn from_dag(graph: Dag, algorithm: impl Into<Arc<str>>) -> Self {
        let mut pending = Vec::new();
        for e in graph.edges() {
            if let Some((from, to)) = e.parent_child() {
                if let (Some(fv), Some(tv)) =
                    (variable_id_of(&graph, from), variable_id_of(&graph, to))
                {
                    pending.push((fv, tv));
                }
            }
        }
        Self { graph, pending_edges: Arc::from(pending), algorithm: algorithm.into() }
    }

    /// Accept a pending directed edge (no-op if absent).
    #[must_use]
    pub fn accept_edge(mut self, from: VariableId, to: VariableId) -> Self {
        let pending: Vec<_> =
            self.pending_edges.iter().copied().filter(|e| *e != (from, to)).collect();
        self.pending_edges = Arc::from(pending);
        self
    }

    /// Accept all remaining pending edges.
    #[must_use]
    pub fn accept_all(mut self) -> Self {
        self.pending_edges = Arc::from([]);
        self
    }

    /// Whether all pending edges have been accepted.
    #[must_use]
    pub fn is_complete(&self) -> bool {
        self.pending_edges.is_empty()
    }

    /// Borrow the accepted DAG when review is complete.
    ///
    /// # Errors
    ///
    /// Incomplete review.
    pub fn try_into_dag(self) -> Result<Dag, GraphError> {
        if !self.is_complete() {
            return Err(GraphError::InvalidEndpoints {
                message: "cannot finish DagReview while pending edges remain",
            });
        }
        Ok(self.graph)
    }
}

fn variable_id_of(dag: &Dag, id: DenseNodeId) -> Option<VariableId> {
    match dag.nodes().get(id.as_usize()) {
        Some(NodeRef::Static(v)) => Some(*v),
        _ => None,
    }
}