gantz_core 0.2.0

//! Items related to collecting a high-level "meta" view of a gantz graph.

use crate::{
    Edge,
    compile::{
        RoseTree,
        error::{InvalidOutputIndex, MetaError, NodeConnsError, TooManyConns},
    },
    node::{self, Node},
    visit::{self, Visitor},
};
use petgraph::visit::{
    Data, EdgeRef, IntoEdgesDirected, IntoNodeReferences, NodeIndexable, NodeRef,
};
use std::collections::{BTreeMap, BTreeSet};

/// Represents a high-level representation of a gantz graph.
///
/// This is produced as the first stage of code-generation and acts as a
/// high-level overview of the gantz graph that can be used for faster
/// traversal and node metadata lookup.
#[derive(Debug, Default)]
pub struct Meta {
    pub graph: MetaGraph,
    /// The set of nodes that require branching on their outputs.
    pub branches: BTreeMap<node::Id, Vec<node::Conns>>,
    /// The set of nodes that require a push evaluation fn.
    pub push: BTreeMap<node::Id, Vec<node::Conns>>,
    /// The set of nodes that require a pull evaluation fn.
    pub pull: BTreeMap<node::Id, Vec<node::Conns>>,
    /// The set of nodes that require access to state.
    pub stateful: BTreeSet<node::Id>,
    /// The set of nodes that act as inlets (for nested graphs).
    pub inlets: BTreeSet<node::Id>,
    /// The set of nodes that act as outlets (for nested graphs).
    pub outlets: BTreeSet<node::Id>,
    /// The total number of inputs on node (whether or not they're connected).
    pub inputs: BTreeMap<node::Id, usize>,
    /// The total number of outputs on node (whether or not they're connected).
    pub outputs: BTreeMap<node::Id, usize>,
}

/// Whether an edge is known to always be traversed, or whether it is
/// conditional.
#[derive(Clone, Debug)]
pub enum EdgeKind {
    /// The edge is known to always be connected.
    Static,
    /// The edge is conditional on node branching.
    Conditional,
}

/// Represents a single flow graph.
///
/// Note that we use a `Vec<Edge>` in order to represent multiple edges
/// between the same two nodes.
pub type MetaGraph = petgraph::graphmap::DiGraphMap<node::Id, Vec<(Edge, EdgeKind)>>;

/// A rose tree of [`Meta`] with error accumulation during visitation.
///
/// Errors encountered during the visitor traversal are accumulated in the
/// `errors` field rather than panicking, allowing all errors to be reported.
#[derive(Debug, Default)]
pub(crate) struct MetaTree {
    /// The rose tree of meta information for the graph and its nested graphs.
    pub tree: RoseTree<Meta>,
    /// Errors accumulated during the visitor traversal.
    pub errors: Vec<MetaError>,
}

impl Meta {
    /// Construct a `Meta` for a single gantz graph.
    pub fn from_graph<'a, G>(get_node: node::GetNode<'a>, g: G) -> Result<Self, NodeConnsError>
    where
        G: Data<EdgeWeight = Edge> + IntoEdgesDirected + IntoNodeReferences + NodeIndexable,
        G::NodeWeight: Node,
    {
        let mut flow = Meta::default();
        let ctx = node::MetaCtx::new(get_node);
        for n_ref in g.node_references() {
            let n = n_ref.id();
            let inputs = g
                .edges_directed(n, petgraph::Direction::Incoming)
                .map(|e_ref| (g.to_index(e_ref.source()), e_ref.weight().clone()));
            let id = g.to_index(n);
            let node = n_ref.weight();
            flow.add_node(ctx, id, node, inputs)?;
        }
        Ok(flow)
    }

    /// Add the node with the given ID and inputs to the `Meta`.
    pub fn add_node(
        &mut self,
        ctx: node::MetaCtx,
        id: node::Id,
        node: &dyn Node,
        inputs: impl IntoIterator<Item = (node::Id, Edge)>,
    ) -> Result<(), NodeConnsError> {
        // Add the node.
        self.graph.add_node(id);

        // Add edges for inputs.
        for (n, edge) in inputs {
            loop {
                if let Some(edges) = self.graph.edge_weight_mut(n, id) {
                    let n_branches = self.branches.get(&n).map(|bs| &bs[..]);
                    if let Some(kind) = edge_kind(n_branches, edge.output.0 as usize)? {
                        edges.push((edge, kind));
                        break;
                    }
                }
                self.graph.add_edge(n, id, vec![]);
            }
        }

        // Register whether the node has inputs or outputs.
        let inputs = node.n_inputs(ctx);
        let outputs = node.n_outputs(ctx);
        if inputs > 0 {
            self.inputs.insert(id, inputs);
        }
        if outputs > 0 {
            self.outputs.insert(id, outputs);
        }

        // Track node branching.
        let branches = node.branches(ctx);
        if !branches.is_empty() {
            self.branches.insert(
                id,
                branches
                    .iter()
                    .map(|conf| conns_from_eval_conf(conf, outputs))
                    .collect::<Result<_, _>>()?,
            );
        }

        // Register push/pull eval for the node if necessary.
        let push_eval = node.push_eval(ctx);
        if !push_eval.is_empty() {
            self.push.insert(
                id,
                push_eval
                    .iter()
                    .map(|conf| conns_from_eval_conf(conf, outputs))
                    .collect::<Result<_, _>>()?,
            );
        }
        let pull_eval = node.pull_eval(ctx);
        if !pull_eval.is_empty() {
            self.pull.insert(
                id,
                pull_eval
                    .iter()
                    .map(|conf| conns_from_eval_conf(conf, inputs))
                    .collect::<Result<_, _>>()?,
            );
        }
        if node.inlet(ctx) {
            self.inlets.insert(id);
        }
        if node.outlet(ctx) {
            self.outlets.insert(id);
        }
        if node.stateful(ctx) {
            self.stateful.insert(id);
        }
        Ok(())
    }
}

/// Allow for constructing a rose-tree of `Meta`s (one for each graph) using
/// the `Node::visit` implementation.
impl Visitor for MetaTree {
    fn visit_pre(&mut self, ctx: visit::Ctx<'_, '_>, node: &dyn Node) {
        let node_path = ctx.path();

        // Ensure the plan for the graph owning this node exists, retrieve it.
        let tree_path = &node_path[..node_path.len() - 1];
        let tree = self.tree.tree_mut(tree_path);

        // Insert the node.
        let id = ctx.id();
        let meta_ctx = node::MetaCtx::new(ctx.get_node());
        if let Err(error) = tree
            .elem
            .add_node(meta_ctx, id, node, ctx.inputs().iter().copied())
        {
            self.errors.push(MetaError {
                path: node_path.to_vec(),
                error,
            });
        }
    }
}

impl super::Edges for Vec<(Edge, EdgeKind)> {
    fn edges(&self) -> impl Iterator<Item = Edge> {
        self.iter().map(|(e, _k)| *e)
    }
}

/// Given an eval conf and a known number of connections, convert the conf to
/// the set of conns.
fn conns_from_eval_conf(
    conf: &node::EvalConf,
    n_conns: usize,
) -> Result<node::Conns, TooManyConns> {
    match conf {
        node::EvalConf::All => node::Conns::try_from_iter((0..n_conns).map(|_| true))
            .map_err(|_| TooManyConns(n_conns)),
        node::EvalConf::Set(conns) => Ok(*conns),
    }
}

/// Given the branching of the source node and the output index of a connected
/// edge, returns the `EdgeKind` of that edge, or `None` if there is no branch
/// under which the edge can be reached.
fn edge_kind(
    confs: Option<&[node::Conns]>,
    out_ix: usize,
) -> Result<Option<EdgeKind>, InvalidOutputIndex> {
    let Some(confs) = confs else {
        return Ok(Some(EdgeKind::Static));
    };
    let mut reachable = false;
    let mut conditional = false;
    for branch in confs {
        let active = branch.get(out_ix).ok_or_else(|| InvalidOutputIndex {
            index: out_ix,
            n_outputs: branch.len(),
        })?;
        reachable |= active;
        conditional |= !active;
    }
    Ok(match (reachable, conditional) {
        (false, _) => None,
        (true, true) => Some(EdgeKind::Conditional),
        (true, false) => Some(EdgeKind::Static),
    })
}