gantz_core 0.3.0

//! Items related to the generation of node functions.
//!
//! In gantz, a function is generated for every unique configuration of every
//! node. That is, for each unique set of connected inputs and outputs of a
//! node, a function is generated.
//!
//! These configurations are collected by traversing from each of the push/pull
//! evaluation entrypoints. With
//! [`Config::emit_all_node_fns`][crate::compile::Config], every node's
//! all-connected configuration is included as well.

use crate::{
    Edge,
    compile::{
        NodeConf, NodeConns, RoseTree,
        error::{NestedGraphNotFound, NodeExprError, NodeFnError, NodeFnErrors},
        names::node_fn_name,
    },
    node::{self, Node},
    visit::{self, Visitor},
};
use petgraph::visit::{Data, IntoEdgesDirected, IntoNodeReferences, NodeIndexable, Visitable};
use std::collections::BTreeSet;
use steel::parser::ast::ExprKind;

/// The set of all node input/output configurations for a single graph.
///
/// These are used to determine which set of functions to generate for each
/// node.
type NodeConfs = BTreeSet<NodeConf>;

/// A visitor used to collect all node functions from a tree of nested gantz
/// graphs.
struct NodeFns<'a> {
    tree: &'a RoseTree<NodeConfs>,
    fns: Vec<(usize, ExprKind)>,
    errors: Vec<NodeFnError>,
}

impl<'a> NodeFns<'a> {
    /// Initialise the `NodeFns` visitor.
    fn new(tree: &'a RoseTree<NodeConfs>) -> Self {
        Self {
            tree,
            fns: vec![],
            errors: vec![],
        }
    }
}

impl Visitor for NodeFns<'_> {
    // We use `visit_post` so that the nested are generated before parents.
    fn visit_post(&mut self, ctx: visit::Ctx<'_, '_>, node: &dyn Node) {
        // Skip generating node functions for inlet and outlet nodes - their
        // values are handled directly by nested_expr bindings.
        let meta_ctx = node::MetaCtx::new(ctx.get_node());
        if node.inlet(meta_ctx) || node.outlet(meta_ctx) {
            return;
        }

        use std::ops::Bound::{Excluded, Included};
        let node_path = ctx.path();
        let plan_path = &node_path[..node_path.len() - 1];
        let tree = match self.tree.tree(plan_path) {
            Some(t) => t,
            None => {
                self.errors
                    .push(NestedGraphNotFound(plan_path.to_vec()).into());
                return;
            }
        };
        let id = ctx.id();
        let empty_conns = NodeConns {
            inputs: node::Conns::empty(),
            outputs: node::Conns::empty(),
        };
        let start = NodeConf {
            id,
            conns: empty_conns.clone(),
        };
        let end_id = id.checked_add(1).expect("node id out of range");
        let end = NodeConf {
            id: end_id,
            conns: empty_conns,
        };
        let range = (Included(start), Excluded(end));
        let input_confs = tree.elem.range(range);
        for conf in input_confs {
            match node_fn(ctx.get_node(), node, node_path, &conf.conns) {
                Ok(fn_expr) => self.fns.push((plan_path.len(), fn_expr)),
                Err(error) => self.errors.push(
                    NodeExprError {
                        path: node_path.to_vec(),
                        error,
                    }
                    .into(),
                ),
            }
        }
    }
}

/// Generate a function for a single node with the given set of connected inputs.
pub(crate) fn node_fn<'a>(
    get_node: node::GetNode<'a>,
    node: &dyn Node,
    node_path: &[node::Id],
    conns: &NodeConns,
) -> Result<ExprKind, node::ExprError> {
    // The binding used to receive the node's state as an argument, and whose
    // resulting value is returned from the body of the function and used to
    // update the state map.
    const STATE: &str = "state";

    fn input_name(i: usize) -> String {
        format!("input{i}")
    }

    // Create function parameters for graph state and inputs
    let mut input_args = conns
        .inputs
        .iter()
        .enumerate()
        .filter_map(|(i, b)| b.then(|| input_name(i)))
        .collect::<Vec<_>>();

    // Create input expressions for the node's expr method
    let input_exprs: Vec<Option<String>> = conns
        .inputs
        .iter()
        .enumerate()
        .map(|(i, b)| b.then(|| input_name(i)))
        .collect();

    // Get the node's expression
    let ctx = node::ExprCtx::new(get_node, node_path, &input_exprs, &conns.outputs);
    let node_expr = node.expr(ctx)?;

    // Construct the full function definition
    // FIXME: Remove this when switching to `flow::NodeConf`.
    let fn_name = node_fn_name(node_path, &conns.inputs, &conns.outputs);
    let meta_ctx = node::MetaCtx::new(get_node);
    let fn_body = if node.stateful(meta_ctx) {
        input_args.push(STATE.to_string());
        format!("(let ((output {node_expr})) (list output state))")
    } else {
        format!("{node_expr}")
    };
    let fn_args = input_args.join(" ");
    let fn_def = format!("(define ({fn_name} {fn_args}) {fn_body})");

    node::parse_expr(&fn_def)
}

/// Given a gantz graph and a rose tree with the associated node configs,
/// produce a function for every node configuration in the graph.
pub(crate) fn node_fns<'a, G>(
    get_node: node::GetNode<'a>,
    g: G,
    node_confs_tree: &RoseTree<NodeConfs>,
) -> Result<Vec<(usize, ExprKind)>, NodeFnErrors>
where
    G: Data<EdgeWeight = Edge> + IntoEdgesDirected + IntoNodeReferences + NodeIndexable + Visitable,
    G::NodeWeight: Node,
{
    let mut node_fns = NodeFns::new(node_confs_tree);
    crate::graph::visit(get_node, g, &[], &mut node_fns);
    if !node_fns.errors.is_empty() {
        return Err(NodeFnErrors(node_fns.errors));
    }
    Ok(node_fns.fns)
}