tokitai-dl 0.1.2

//! Adapters that treat a tokitai-operator [`SemanticGraph`] as a god-graph
//! [`Graph`], so that god-graph's algorithm library (BFS, DFS, topological
//! sort, connected components, MST, PageRank, shortest path, etc.) can be
//! applied to operator DAGs produced by `tokitai-operator`.
//!
//! This module is gated on the `operator` feature (for the `SemanticGraph`
//! import). The graph-construction entry point additionally requires the
//! `graph` feature (for the [`Graph`] type). Pure metadata helpers
//! ([`semantic_node_count`], [`semantic_input_count`]) only need the
//! `operator` feature.
//!
//! The [`op_dag_neighbors`], [`op_dag_pagerank`], and
//! [`op_dag_shortest_path`] helpers are gated on **both** the `graph` and
//! the `operator` features, because they need to walk the
//! `SemanticGraph` *and* apply god-graph algorithms on the converted
//! `Graph`.

use tokitai_operator::ir::SemanticGraph;

#[cfg(feature = "graph")]
use god_graph::Graph;

#[cfg(feature = "graph")]
use god_graph::graph::GraphOps;

#[cfg(feature = "graph")]
use crate::BridgeError;

// ---------------------------------------------------------------------------
// Conversion
// ---------------------------------------------------------------------------

/// Convert a [`SemanticGraph`] (an operator DAG with value-id edges) into a
/// god-graph [`Graph<String, ()>`] whose node payloads are the operator
/// names (`op_name` strings) and whose directed edges follow the dataflow.
///
/// Each [`tokitai_operator::ir::SemanticNode`] in the input becomes one node in the output graph.
/// For every value id listed in `node.inputs`, the producer node is located
/// by scanning all nodes' `output_ids`, and a directed edge is added from
/// producer -> consumer.
///
/// This is the compatibility convenience API. It preserves the legacy
/// best-effort behavior: malformed or unrepresentable dependencies are
/// skipped instead of reported. Specifically, it skips self-edges, skips
/// input value ids that have no producer
/// [`tokitai_operator::ir::SemanticNode`], and drops god-graph `add_node` /
/// `add_edge` insertion errors. This keeps old callers total, but it can lose
/// information.
///
/// For strict construction that returns a diagnostic error instead of
/// dropping nodes or edges, use [`try_semantic_to_god_graph`].
///
/// This function requires both the `operator` and `graph` features.
///
/// # Example
///
/// ```rust
/// # #[cfg(all(feature = "graph", feature = "operator"))]
/// # fn demo() {
/// use god_graph::graph::traits::GraphBase;
/// use tokitai_operator::ir::SemanticGraph;
/// use tokitai_dl::op_dag_graph::semantic_to_god_graph;
///
/// // An empty SemanticGraph converts to an empty god-graph Graph.
/// let sg = SemanticGraph::new();
/// let g = semantic_to_god_graph(&sg);
/// assert_eq!(g.node_count(), 0);
/// # }
/// ```
#[cfg(feature = "graph")]
pub fn semantic_to_god_graph(sg: &SemanticGraph) -> Graph<String, ()> {
    match convert_semantic_to_god_graph(sg, ConversionMode::BestEffort) {
        Ok(g) => g,
        Err(err) => {
            #[cfg(feature = "tracing")]
            tracing::debug!("best-effort SemanticGraph conversion failed unexpectedly: {err}");
            #[cfg(all(feature = "observability", not(feature = "tracing")))]
            log::debug!("best-effort SemanticGraph conversion failed unexpectedly: {err}");
            let _ = err;
            Graph::directed()
        }
    }
}

/// Convert a [`SemanticGraph`] into a god-graph [`Graph<String, ()>`] without
/// silently dropping conversion errors.
///
/// This strict API uses the same node and edge model as
/// [`semantic_to_god_graph`]: each [`tokitai_operator::ir::SemanticNode`]
/// becomes one god-graph node, and value dependencies become directed
/// producer-to-consumer edges.
///
/// Unlike the compatibility API, this function returns [`BridgeError`] when
/// this projection would lose representable graph structure:
///
/// - god-graph `add_node` / `add_edge` failures are returned as
///   [`BridgeError::Upstream`];
/// - input value ids with no producing
///   [`tokitai_operator::ir::SemanticNode`] are returned as
///   [`BridgeError::Unsupported`], including valid external graph inputs from
///   [`SemanticGraph::add_input`] because this `Graph<String, ()>` projection
///   has no node to attach them to;
/// - self-edges are returned as [`BridgeError::Unsupported`] because they
///   indicate malformed input for this DAG projection.
///
/// "Strict" here means strict adapter accounting. It does not validate the
/// semantic correctness of upstream operators and does not prove any
/// god-graph algorithm result.
///
/// This function requires both the `operator` and `graph` features.
///
/// # Example
///
/// ```rust
/// # #[cfg(all(feature = "graph", feature = "operator"))]
/// # fn demo() -> Result<(), tokitai_dl::BridgeError> {
/// use god_graph::graph::traits::GraphBase;
/// use tokitai_operator::ir::SemanticGraph;
/// use tokitai_dl::op_dag_graph::try_semantic_to_god_graph;
///
/// let sg = SemanticGraph::new();
/// let g = try_semantic_to_god_graph(&sg)?;
/// assert_eq!(g.node_count(), 0);
/// # Ok(())
/// # }
/// ```
#[cfg(feature = "graph")]
pub fn try_semantic_to_god_graph(sg: &SemanticGraph) -> Result<Graph<String, ()>, BridgeError> {
    convert_semantic_to_god_graph(sg, ConversionMode::Strict)
}

#[cfg(feature = "graph")]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ConversionMode {
    BestEffort,
    Strict,
}

#[cfg(feature = "graph")]
impl ConversionMode {
    fn is_strict(self) -> bool {
        matches!(self, Self::Strict)
    }
}

#[cfg(feature = "graph")]
fn convert_semantic_to_god_graph(
    sg: &SemanticGraph,
    mode: ConversionMode,
) -> Result<Graph<String, ()>, BridgeError> {
    #[cfg(feature = "tracing")]
    tracing::trace!(
        "op_dag_graph::semantic_to_god_graph: entry ({} SemanticNodes)",
        sg.nodes().len()
    );
    #[cfg(all(feature = "observability", not(feature = "tracing")))]
    log::trace!(
        "op_dag_graph::semantic_to_god_graph: entry ({} SemanticNodes)",
        sg.nodes().len()
    );
    let mut g: Graph<String, ()> = Graph::directed();

    // Phase 1: add a node for every SemanticNode, keyed by node id.
    // The payload is the op_name (e.g. "add", "matmul", "relu", "input").
    // In strict mode, surface god-graph insertion failures. In best-effort
    // mode, keep the legacy behavior: leave the slot empty and drop affected
    // edges later.
    let mut node_indices: Vec<Option<god_graph::NodeIndex>> = vec![None; sg.nodes().len()];
    for (i, node) in sg.nodes().iter().enumerate() {
        let payload = node.op_name.clone();
        match g.add_node(payload) {
            Ok(ni) => node_indices[i] = Some(ni),
            Err(err) if mode.is_strict() => {
                return Err(BridgeError::Upstream(format!(
                    "failed to add god-graph node for SemanticNode {i} (op_name={:?}): {err}",
                    node.op_name
                )));
            }
            Err(_err) => {}
        }
    }

    // Phase 2: build a value-id -> producer-node-id map by scanning every
    // node's `output_ids`. Each value id in a `SemanticGraph` is produced
    // by exactly one node, so the second insert (if any) is unreachable in
    // well-formed input. Strict mode reports duplicate producers; best-effort
    // mode keeps the legacy last-writer-wins map behavior.
    let mut value_producer: std::collections::HashMap<usize, usize> =
        std::collections::HashMap::new();
    for (i, node) in sg.nodes().iter().enumerate() {
        for &out_id in &node.output_ids {
            if let Some(previous_idx) = value_producer.insert(out_id, i) {
                if mode.is_strict() {
                    return Err(BridgeError::Unsupported(format!(
                        "value id {out_id} is produced by multiple SemanticNodes: {previous_idx} and {i}"
                    )));
                }
            }
        }
    }

    // Phase 3: for each node, for each input value id, add a directed edge
    // from the producer node to this node. Best-effort mode skips self-edges,
    // missing producers, and insertion failures for compatibility. Strict
    // mode returns a diagnostic before information is lost.
    for (i, node) in sg.nodes().iter().enumerate() {
        let this_ni = match node_indices[i] {
            Some(ni) => ni,
            None if mode.is_strict() => {
                return Err(BridgeError::Upstream(format!(
                    "missing god-graph node index for SemanticNode {i} (op_name={:?})",
                    node.op_name
                )));
            }
            None => continue,
        };
        for &input_value_id in &node.inputs {
            let Some(&producer_idx) = value_producer.get(&input_value_id) else {
                if mode.is_strict() {
                    return Err(BridgeError::Unsupported(missing_producer_message(
                        sg,
                        i,
                        &node.op_name,
                        input_value_id,
                    )));
                }
                continue;
            };

            if producer_idx == i {
                if mode.is_strict() {
                    return Err(BridgeError::Unsupported(format!(
                        "SemanticNode {i} (op_name={:?}) input value id {input_value_id} is produced by the same node; strict conversion rejects self-edges",
                        node.op_name
                    )));
                }
                continue;
            }

            let producer_ni = match node_indices[producer_idx] {
                Some(ni) => ni,
                None if mode.is_strict() => {
                    return Err(BridgeError::Upstream(format!(
                        "missing god-graph node index for producer SemanticNode {producer_idx} while adding edge to SemanticNode {i}"
                    )));
                }
                None => continue,
            };

            match g.add_edge(producer_ni, this_ni, ()) {
                Ok(_edge) => {}
                Err(err) if mode.is_strict() => {
                    return Err(BridgeError::Upstream(format!(
                        "failed to add god-graph edge for input value id {input_value_id} from SemanticNode {producer_idx} to SemanticNode {i}: {err}"
                    )));
                }
                Err(_err) => {}
            }
        }
    }

    #[cfg(feature = "tracing")]
    {
        use god_graph::graph::traits::GraphBase;
        tracing::debug!(
            "op_dag_graph::semantic_to_god_graph: exit ({} nodes, {} edges)",
            g.node_count(),
            g.edge_count()
        );
    }
    #[cfg(all(feature = "observability", not(feature = "tracing")))]
    {
        use god_graph::graph::traits::GraphBase;
        log::debug!(
            "op_dag_graph::semantic_to_god_graph: exit ({} nodes, {} edges)",
            g.node_count(),
            g.edge_count()
        );
    }

    Ok(g)
}

#[cfg(feature = "graph")]
fn missing_producer_message(
    sg: &SemanticGraph,
    node_idx: usize,
    op_name: &str,
    input_value_id: usize,
) -> String {
    let value_kind = if sg.value(input_value_id).is_some() {
        "external graph input"
    } else {
        "unknown value"
    };
    format!(
        "SemanticNode {node_idx} (op_name={op_name:?}) input value id {input_value_id} has no SemanticNode producer ({value_kind}); strict conversion cannot represent this dependency in Graph<String, ()>"
    )
}

// ---------------------------------------------------------------------------
// Metadata helpers
// ---------------------------------------------------------------------------

/// Return the number of operator nodes in a [`SemanticGraph`].
///
/// This is a thin wrapper over [`SemanticGraph::nodes`]`(`)``.`len`(`)` `)`.
/// It exists so that downstream code can use tokitai-dl as a single import
/// surface without reaching into tokitai-operator directly.
///
/// # Example
///
/// ```rust
/// # #[cfg(feature = "operator")]
/// # fn demo() {
/// use tokitai_operator::ir::SemanticGraph;
/// use tokitai_dl::op_dag_graph::semantic_node_count;
///
/// let sg = SemanticGraph::new();
/// assert_eq!(semantic_node_count(&sg), 0);
/// # }
/// ```
pub fn semantic_node_count(sg: &SemanticGraph) -> usize {
    sg.nodes().len()
}

/// Return the number of graph-input nodes in a [`SemanticGraph`].
///
/// The `tokitai-operator` convention (see `op::nn::Input` and friends in the
/// upstream crate) is that input operators report their `op_name` as the
/// string `"input"`. This helper counts such nodes.
///
/// Note that [`SemanticGraph::add_input`] only appends a value to the
/// `values` arena; the corresponding `SemanticNode` with `op_name == "input"`
/// is created by `add_op` (or any operator whose `name()` is `"input"`).
///
/// # Example
///
/// ```rust
/// # #[cfg(feature = "operator")]
/// # fn demo() {
/// use tokitai_operator::ir::SemanticGraph;
/// use tokitai_dl::op_dag_graph::semantic_input_count;
///
/// let sg = SemanticGraph::new();
/// // No nodes have been added yet, so the input count is zero.
/// assert_eq!(semantic_input_count(&sg), 0);
/// # }
/// ```
pub fn semantic_input_count(sg: &SemanticGraph) -> usize {
    sg.nodes().iter().filter(|n| n.op_name == "input").count()
}

// ---------------------------------------------------------------------------
// Algorithm adapters (require both `graph` and `operator` features)
// ---------------------------------------------------------------------------

/// Find every operator node whose `op_name` equals `target_op_name`, and
/// for each match, return the list of op_names of its direct consumers
/// (one step downstream in the dataflow DAG).
///
/// The returned `Vec` has one entry `(node_id, consumer_op_names)` per
/// matching `SemanticNode`, in the same order as
/// [`SemanticGraph::nodes`]. `node_id` is the index into the
/// `SemanticGraph`'s `nodes()` slice (which matches the node id the god-graph
/// conversion assigns).
///
/// "Direct consumer" means: any `SemanticNode` whose `inputs` list
/// references a value id produced by the matching node. The consumer's
/// `op_name` (a `String`, since multiple operators can share a name) is
/// returned, deduplicated and sorted, to make test assertions easy.
///
/// This is a pure tokitai walk — it does not need the god-graph
/// `Graph<String, ()>` representation. It is provided alongside the
/// graph-algorithm adapters so users can do "neighbors" queries without
/// paying for the conversion when they only need the
/// (node, immediate consumers) shape.
///
/// # Example
///
/// ```rust
/// # #[cfg(all(feature = "graph", feature = "operator"))]
/// # fn demo() {
/// use tokitai_operator::ir::SemanticGraph;
/// use tokitai_dl::op_dag_graph::op_dag_neighbors;
///
/// // An empty graph has no matching op names, so the result is empty.
/// let sg = SemanticGraph::new();
/// let neighbors = op_dag_neighbors(&sg, "add");
/// assert!(neighbors.is_empty());
/// # }
/// ```
#[cfg(all(feature = "graph", feature = "operator"))]
pub fn op_dag_neighbors(sg: &SemanticGraph, target_op_name: &str) -> Vec<(usize, Vec<String>)> {
    // Build a value_id -> producer_node_id map (the same one used by
    // `semantic_to_god_graph`, but kept local so this function can be
    // called without paying for the conversion up front).
    let mut value_producer: std::collections::HashMap<usize, usize> =
        std::collections::HashMap::new();
    for (i, node) in sg.nodes().iter().enumerate() {
        for &out_id in &node.output_ids {
            value_producer.insert(out_id, i);
        }
    }

    let mut out: Vec<(usize, Vec<String>)> = Vec::new();
    for (i, node) in sg.nodes().iter().enumerate() {
        if node.op_name != target_op_name {
            continue;
        }
        // Collect all consumers: any node whose `inputs` references one of
        // our output value ids.
        let mut consumers: std::collections::BTreeSet<String> = std::collections::BTreeSet::new();
        for &out_id in &node.output_ids {
            for (j, other) in sg.nodes().iter().enumerate() {
                if j == i {
                    continue;
                }
                if other.inputs.contains(&out_id) {
                    consumers.insert(other.op_name.clone());
                }
            }
        }
        out.push((i, consumers.into_iter().collect()));
    }
    out
}

/// Run god-graph's [`pagerank`](god_graph::algorithms::centrality::pagerank)
/// on the operator DAG and return `op_name -> score` pairs sorted by score
/// descending.
///
/// The function converts the [`SemanticGraph`] into a god-graph
/// `Graph<String, ()>` via [`semantic_to_god_graph`] (so node payloads are
/// the operator names), runs 20 iterations of PageRank with damping
/// 0.85, zips scores with payloads, sorts, and returns the result.
///
/// A node whose payload `op_name` appears multiple times in the DAG
/// (e.g. a chain of three `relu` ops) gets the sum of its per-instance
/// PageRank scores — this matches the "importance of the operator class
/// in the whole graph" intuition.
///
/// # Example
///
/// ```rust
/// # #[cfg(all(feature = "graph", feature = "operator"))]
/// # fn demo() {
/// use tokitai_operator::ir::SemanticGraph;
/// use tokitai_dl::op_dag_graph::op_dag_pagerank;
///
/// // An empty graph has no nodes, so the PageRank result is empty.
/// let sg = SemanticGraph::new();
/// let pr = op_dag_pagerank(&sg);
/// assert!(pr.is_empty());
/// # }
/// ```
#[cfg(all(feature = "graph", feature = "operator"))]
pub fn op_dag_pagerank(sg: &SemanticGraph) -> Vec<(String, f64)> {
    use god_graph::algorithms::centrality::pagerank;
    use god_graph::graph::traits::GraphQuery;

    #[cfg(feature = "tracing")]
    tracing::trace!(
        "op_dag_graph::op_dag_pagerank: entry ({} SemanticNodes)",
        sg.nodes().len()
    );
    #[cfg(all(feature = "observability", not(feature = "tracing")))]
    log::trace!(
        "op_dag_graph::op_dag_pagerank: entry ({} SemanticNodes)",
        sg.nodes().len()
    );
    let g: Graph<String, ()> = semantic_to_god_graph(sg);
    let scores = pagerank(&g, 0.85, 20);

    // Walk node payloads in node-id order (the same order
    // `semantic_to_god_graph` inserted them) and look up the PageRank
    // score for each. If a payload appears at multiple node indices,
    // sum their scores. The result is a deterministic `Vec<(String,
    // f64)>` we can sort by score descending.
    let mut aggregated: std::collections::HashMap<String, f64> = std::collections::HashMap::new();
    for node in g.nodes() {
        let payload = node.data().clone();
        let score = scores.get(&node.index()).copied().unwrap_or(0.0);
        *aggregated.entry(payload).or_insert(0.0) += score;
    }
    let mut out: Vec<(String, f64)> = aggregated.into_iter().collect();
    out.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    out
}

/// Compute the shortest dependency chain between two operator names in the
/// dataflow DAG.
///
/// Returns `Some(Vec<String>)` of op_name labels along a shortest path from
/// `from_op_name` to `to_op_name` (inclusive of both endpoints), or `None`
/// if no path exists. "Shortest" here means fewest hops; edge weights are
/// uniform (every edge counts as 1), so the chain length equals the
/// number of operator-application steps between the two endpoints.
///
/// Implementation: convert via [`semantic_to_god_graph`], then run a
/// BFS from every node whose payload equals `from_op_name` (depth 0
/// targets) until we reach any node whose payload equals
/// `to_op_name`. The first match wins. This is implemented directly
/// in terms of the god-graph `GraphQuery` trait (BFS visitor +
/// predecessor back-pointers) because god-graph does not ship a
/// `shortest_path::bfs` convenience function with a goal predicate.
///
/// # Example
///
/// ```rust
/// # #[cfg(all(feature = "graph", feature = "operator"))]
/// # fn demo() {
/// use tokitai_operator::ir::SemanticGraph;
/// use tokitai_dl::op_dag_graph::op_dag_shortest_path;
///
/// // An empty graph has no path between any two op names.
/// let sg = SemanticGraph::new();
/// let path = op_dag_shortest_path(&sg, "add", "softmax");
/// assert!(path.is_none());
/// # }
/// ```
#[cfg(all(feature = "graph", feature = "operator"))]
pub fn op_dag_shortest_path(
    sg: &SemanticGraph,
    from_op_name: &str,
    to_op_name: &str,
) -> Option<Vec<String>> {
    use god_graph::graph::traits::{GraphBase, GraphQuery};
    use god_graph::node::NodeIndex;
    use std::collections::VecDeque;

    #[cfg(feature = "tracing")]
    tracing::trace!(
        "op_dag_graph::op_dag_shortest_path: entry (from={from_op_name:?}, to={to_op_name:?})"
    );
    #[cfg(all(feature = "observability", not(feature = "tracing")))]
    log::trace!(
        "op_dag_graph::op_dag_shortest_path: entry (from={from_op_name:?}, to={to_op_name:?})"
    );
    let g: Graph<String, ()> = semantic_to_god_graph(sg);
    if g.node_count() == 0 {
        return None;
    }

    // Find all start / goal candidates. We don't require them to be
    // unique — if multiple nodes share the same op_name we treat the
    // first reachable source-to-sink pair as the answer.
    let sources: Vec<NodeIndex> = g
        .nodes()
        .filter(|n| n.data() == from_op_name)
        .map(|n| n.index())
        .collect();
    let goals: Vec<NodeIndex> = g
        .nodes()
        .filter(|n| n.data() == to_op_name)
        .map(|n| n.index())
        .collect();
    if sources.is_empty() || goals.is_empty() {
        return None;
    }

    // BFS from every source simultaneously; we use a per-source
    // predecessor map and pick the first goal we reach.
    let n = g.node_count();
    let mut visited = vec![None::<NodeIndex>; n];
    let mut queue: VecDeque<(NodeIndex, NodeIndex)> = VecDeque::new();
    for &s in &sources {
        visited[s.index()] = Some(s);
        queue.push_back((s, s));
    }
    let mut found: Option<(NodeIndex, NodeIndex)> = None;
    while let Some((node, root)) = queue.pop_front() {
        if goals.contains(&node) {
            found = Some((node, root));
            break;
        }
        for nb in g.neighbors(node) {
            if visited[nb.index()].is_none() {
                visited[nb.index()] = Some(node);
                queue.push_back((nb, root));
            }
        }
    }

    let (goal, _root) = found?;
    // Reconstruct the path from goal back to its source root, then
    // collect op_name payloads along the way.
    let mut path: Vec<NodeIndex> = Vec::new();
    let mut cur = goal;
    loop {
        path.push(cur);
        match visited[cur.index()] {
            Some(prev) if prev != cur => cur = prev,
            _ => break, // reached a source
        }
    }
    path.reverse();
    Some(
        path.iter()
            .map(|ni| g.get_node(*ni).cloned().unwrap_or_default())
            .collect(),
    )
}