tierkreis_runtime/operations/

mod.rs

//! Implements operations as processes each acting asynchronously on a stream
//! of inputs to produce a stream of outputs.
use self::eval::run_eval;
use self::function::run_fn;
use self::graph::GraphOperation;
use self::r#box::run_box;
use crate::util::JoinHandleWithDrop;
use crate::workers::EscapeHatch;
use crate::Runtime;
use anyhow::{anyhow, bail};
use futures::future::{self, AbortHandle, BoxFuture};
use futures::stream::BoxStream;
use futures::{Future, FutureExt, Stream, StreamExt};
use std::collections::HashSet;
use std::pin::Pin;
use std::{collections::HashMap, sync::Arc};
use thiserror::Error;
use tierkreis_core::graph::{Edge, Graph, GraphBuilder, Node, Value};
use tierkreis_core::prelude::TryInto;
use tierkreis_core::symbol::{Label, Location, SymbolError};
use tierkreis_proto::messages::{Callback, Completed, GraphTrace, Status};
use tokio::sync::mpsc;
use tokio::sync::watch;
use tokio_stream::wrappers::UnboundedReceiverStream;
use tracing::Instrument;

pub(crate) mod r#box;
pub(crate) mod eval;
pub(crate) mod function;
pub(crate) mod graph;
pub(crate) mod variant;

pub use graph::checkpoint_client::CheckpointClient;

/// A stream of inputs (one per wire) arriving at an operation,
/// plus a notification after the last input (wire) has arrived.
pub struct OperationInputs(BoxStream<'static, Input>);

impl Stream for OperationInputs {
    type Item = Input;

    fn poll_next(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        Pin::new(&mut self.as_mut().0).poll_next(cx)
    }
}

/// Stream of outputs generated by a running operation.
///
/// When this stream is dropped the operation will be cancelled.
pub struct OperationOutputs {
    stream: BoxStream<'static, Output>,
    abort_input: AbortHandle,
}

impl OperationOutputs {
    /// Views the operation as a process
    pub fn into_task(self) -> TaskHandle {
        TaskHandle::new(self)
    }
}

impl Drop for OperationOutputs {
    fn drop(&mut self) {
        self.abort_input.abort()
    }
}

impl Stream for OperationOutputs {
    type Item = Output;

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        Pin::new(&mut self.as_mut().stream).poll_next(cx)
    }
}

/// A single operation that can be started (executed asynchronously) at runtime,
/// given some inputs and a context including a stream of outputs. Does not
/// necessarily have to be related to a specific node in a Tierkreis graph,
/// although generally it is.
pub struct RuntimeOperation {
    start: Box<
        dyn FnOnce(OperationContext, OperationInputs) -> BoxFuture<'static, anyhow::Result<()>>
            + Send,
    >,
}

impl RuntimeOperation {
    fn new<F, FF>(f: F) -> Self
    where
        F: FnOnce(OperationContext, OperationInputs) -> FF + Send + 'static,
        FF: Future<Output = anyhow::Result<()>> + Send + 'static,
    {
        Self {
            start: Box::new(|ctx, inputs| f(ctx, inputs).boxed()),
        }
    }

    /// Creates a RuntimeOperation that just produces a constant value (given).
    pub fn new_const(value: Value) -> RuntimeOperation {
        operation_const(value)
    }

    /// Creates a RuntimeOperation that runs a graph (found in a box) at some location,
    /// i.e. perhaps remotely.
    /// Note: if `loc` is [Location::local] then this is largely equivalent to [Self::new_graph].
    pub fn new_box(loc: Location, graph: Graph) -> RuntimeOperation {
        RuntimeOperation::new(|ctx, inputs| run_box(loc, graph, ctx, inputs))
    }

    /// Creates a RuntimeOperation that runs a graph (in the local runtime, interleaved
    /// with other nodes/etc. executing). Inputs to the operation will become outputs of
    /// the graph Input node; the operation outputs are those received by the graph Output node.    
    pub fn new_graph(graph: Graph) -> RuntimeOperation {
        RuntimeOperation::new(move |ctx, inputs| GraphOperation::new(graph, ctx, inputs).run())
    }

    /// Creates a RuntimeOperation that runs a `match` operation, i.e. expects
    /// to be given a variant value and a Graph for each *possible* variant, and
    /// runs the appropriate one of those Graphs.
    pub(crate) fn new_match() -> RuntimeOperation {
        RuntimeOperation::new(variant::run_match)
    }

    /// Creates a RuntimeOperation that tags a value to make a [Value::Variant]
    pub(crate) fn new_tag(tag: Label) -> RuntimeOperation {
        RuntimeOperation::new(move |ctx, inputs| variant::run_tag(tag, ctx, inputs))
    }

    /// Create a new runtime operation from a simple function.
    ///
    /// Do not use this for blocking or expensive CPU bound functions.
    pub fn new_fn_simple<F>(f: F) -> Self
    where
        F: FnOnce(HashMap<Label, Value>, OperationContext) -> anyhow::Result<HashMap<Label, Value>>
            + Send
            + 'static,
    {
        let f = |inputs, ctx| futures::future::ready(f(inputs, ctx));
        RuntimeOperation::new(move |ctx, inputs| run_fn(f, ctx, inputs))
    }

    /// Create a new runtime operation from an async function.
    ///
    /// Do not use this for blocking or expensive CPU bound functions.
    pub fn new_fn_async<F, FF>(f: F) -> Self
    where
        F: FnOnce(HashMap<Label, Value>, OperationContext) -> FF + Send + 'static,
        FF: Future<Output = anyhow::Result<HashMap<Label, Value>>> + Send + 'static,
    {
        let f = |inputs, ctx| {
            let span = tracing::Span::current();
            let handle = tokio::spawn(f(inputs, ctx).instrument(span));
            JoinHandleWithDrop::from(handle).map(|r| r.unwrap_or_else(|e| Err(e.into())))
        };
        RuntimeOperation::new(move |ctx, inputs| run_fn(f, ctx, inputs))
    }

    /// Create a new runtime operation from a blocking function.
    ///
    /// The function is run via [`tokio::task::spawn_blocking`].
    pub fn new_fn_blocking<F>(f: F) -> Self
    where
        F: FnOnce(HashMap<Label, Value>, OperationContext) -> anyhow::Result<HashMap<Label, Value>>
            + Send
            + 'static,
    {
        let f = |inputs, ctx| {
            let span = tracing::Span::current();
            let handle = tokio::task::spawn_blocking(move || span.in_scope(|| f(inputs, ctx)));
            JoinHandleWithDrop::from(handle).map(|r| r.unwrap_or_else(|e| Err(e.into())))
        };
        RuntimeOperation::new(move |ctx, inputs| run_fn(f, ctx, inputs))
    }

    /// Starts the operation i.e. so it will execute asynchronously (using `tokio::spawn`).
    /// Returns the stream of outputs.
    pub fn run<S>(
        self,
        runtime: Runtime,
        callback: Callback,
        escape: EscapeHatch,
        inputs: S,
        stack_trace: GraphTrace,
        checkpoint_client: Option<CheckpointClient>,
    ) -> OperationOutputs
    where
        S: Stream<Item = Input> + Send + 'static,
    {
        let inputs = inputs.chain(futures::stream::pending());
        let (inputs, abort_handle) = futures::stream::abortable(inputs);

        let (output_tx, output_rx) = mpsc::unbounded_channel();

        let context = OperationContext {
            output: output_tx.clone(),
            callback,
            escape,
            runtime,
            graph_trace: stack_trace,
            checkpoint_client,
        };

        let span = tracing::Span::current();

        // TODO: Use the join handle to react to panics
        tokio::spawn(
            async move {
                let result = (self.start)(context, OperationInputs(inputs.boxed())).await;

                let _ = match result {
                    Ok(()) => output_tx.send(Output::Success),
                    Err(error) => output_tx.send(Output::Failure { error }),
                };
            }
            .instrument(span),
        );

        OperationOutputs {
            stream: UnboundedReceiverStream::new(output_rx).boxed(),
            abort_input: abort_handle,
        }
    }

    /// Like [`RuntimeOperation::run`] but the inputs are available immediately.
    pub fn run_simple<I>(
        self,
        runtime: Runtime,
        callback: Callback,
        escape: EscapeHatch,
        inputs: I,
        stack_trace: GraphTrace,
        checkpoint_client: Option<CheckpointClient>,
    ) -> OperationOutputs
    where
        I: IntoIterator<Item = (Label, Value)>,
    {
        let inputs: HashMap<_, _> = inputs.into_iter().collect();

        let inputs_stream = futures::stream::iter(inputs)
            .map(|(port, value)| Input::Input { port, value })
            .chain(tokio_stream::once(Input::Complete));

        self.run(
            runtime,
            callback,
            escape,
            inputs_stream,
            stack_trace,
            checkpoint_client,
        )
    }
}

/// A task-level view of a process, allowing only to poll or block for completion
/// and to request the task to abort, rather than a stream of results.
#[derive(Clone)]
pub struct TaskHandle {
    status: watch::Receiver<Status>,
    abort: mpsc::UnboundedSender<()>,
}

impl std::fmt::Debug for TaskHandle {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TaskHandle").finish()
    }
}

impl TaskHandle {
    fn new(mut output_stream: OperationOutputs) -> Self {
        let (status_tx, status_rx) = watch::channel(Status::Running);
        let (abort_tx, mut abort_rx) = mpsc::unbounded_channel();

        tokio::spawn(async move {
            let mut outputs = HashMap::new();

            loop {
                let output = tokio::select! { biased;
                    _ = abort_rx.recv() => break,
                    msg = output_stream.next() => {
                        match msg {
                            Some(msg) => msg,
                            None => break,
                        }
                    }
                };

                match output {
                    Output::Output { port, value } => {
                        outputs.insert(port, value);
                    }
                    Output::Success => {
                        let _ = status_tx.send(Status::Completed(Ok(Arc::new(outputs))));

                        return;
                    }
                    Output::Failure { error } => {
                        let _ = status_tx.send(Status::Completed(Err(Arc::new(error))));
                        return;
                    }
                }
            }

            let error = anyhow!("task was cancelled");
            let _ = status_tx.send(Status::Completed(Err(Arc::new(error))));
        });

        Self {
            status: status_rx,
            abort: abort_tx,
        }
    }

    /// Polls the current status of the process.
    /// (Note this copies the entire result map, if there is one.)
    pub fn status(&self) -> Status {
        self.status.borrow().clone()
    }

    /// Cancels a running task even when there are other copies of the task handle.
    pub fn cancel(&self) {
        let _ = self.abort.send(());
    }

    /// Waits for the task to complete and returns the result (outputs or error)
    pub async fn complete(&mut self) -> Completed {
        loop {
            if let Status::Completed(result) = self.status() {
                match result {
                    Ok(_) => tracing::debug!("complete"),
                    Err(_) => tracing::warn!("complete with error"),
                }
                return result;
            }
            tracing::debug!("still running");

            let event = self.status.changed().await;

            if event.is_err() {
                unreachable!("watch was closed before setting a completed status");
            }
        }
    }
}

/// An event or signal sent to a process about its inputs
#[derive(Debug)]
pub enum Input {
    /// A value is available at an input port.
    #[allow(missing_docs)]
    Input { port: Label, value: Value },
    /// Values for all input ports have been sent.
    Complete,
}

/// An event or signal sent from a process about its outputs
#[derive(Debug)]
pub enum Output {
    /// A value is available at an output port.
    #[allow(missing_docs)]
    Output { port: Label, value: Value },
    /// The operation has finished processing successfully.
    Success,
    /// The operation has finished processing with an error.
    #[allow(missing_docs)]
    Failure { error: anyhow::Error },
}

impl Output {
    /// Adds a line to the error trace that will be displayed if this is an [Output::Failure]
    pub fn context<C>(self, context: C) -> Self
    where
        C: std::fmt::Display + Send + Sync + 'static,
    {
        match self {
            Self::Output { port, value } => Self::Output { port, value },
            Self::Success => Self::Success,
            Self::Failure { error } => {
                let error = error.context(context);
                Self::Failure { error }
            }
        }
    }
}

/// The context available to an actor running an operation.
#[derive(Clone)]
pub struct OperationContext {
    output: mpsc::UnboundedSender<Output>,
    /// Runtime within which the operation executes; allows access to features
    /// outside the particular graph being run, such as connected workers.
    pub runtime: Runtime,
    /// FunctionWorkers (executing run_function requests) can use this to make run_graph
    /// calls (also infer_type and other RuntimeWorker operations, but not run_function).
    /// The Callback will route requests back up the tree of runtimes to the
    /// closest explicitly-enclosing Scope specified by the user's Graph.
    pub callback: Callback,
    /// Runtimes use this to run any function(-name)s they don't recognize themselves;
    /// it will route requests back up to the original root  of the client request.
    pub escape: EscapeHatch,
    /// The trace of evaluations by which we came to be executing this operation (node *or* graph)
    pub graph_trace: GraphTrace,
    /// The checkpointing client for the current job.
    pub checkpoint_client: Option<CheckpointClient>,
}

impl OperationContext {
    /// Notify the graph actor that a value is available at one of the node's output ports.
    pub fn set_output(&self, port: impl Into<Label>, value: Value) {
        let _ = self.output.send(Output::Output {
            port: port.into(),
            value,
        });
    }

    // check if the graph is the outermost graph and has a checkpoint client
    fn outer_graph_checkpoint(&mut self) -> Option<&mut CheckpointClient> {
        if self.graph_trace == GraphTrace::Root {
            self.checkpoint_client.as_mut()
        } else {
            None
        }
    }
}

pub(crate) fn operation_eval() -> RuntimeOperation {
    RuntimeOperation::new(run_eval)
}

pub(crate) fn operation_const(value: Value) -> RuntimeOperation {
    RuntimeOperation::new_fn_simple(move |_inputs, _context| {
        let mut outputs = HashMap::new();
        outputs.insert(Label::value(), value);
        Ok(outputs)
    })
}

pub(crate) fn operation_id() -> RuntimeOperation {
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let mut outputs = HashMap::new();
        outputs.insert(Label::value(), take_input(&mut inputs, Label::value())?);
        Ok(outputs)
    })
}

pub(crate) fn operation_sleep() -> RuntimeOperation {
    RuntimeOperation::new_fn_async(|mut inputs, _context| async move {
        let delay = validate_float_input(&mut inputs, "delay_secs")?;
        tokio::time::sleep(tokio::time::Duration::from_secs_f64(delay)).await;
        let mut outputs = HashMap::new();
        outputs.insert(Label::value(), take_input(&mut inputs, Label::value())?);
        Ok(outputs)
    })
}

pub(crate) fn operation_copy() -> RuntimeOperation {
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let value = take_input(&mut inputs, Label::value())?;
        let mut outputs = HashMap::new();
        outputs.insert(TryInto::try_into("value_0")?, value.clone());
        outputs.insert(TryInto::try_into("value_1")?, value);
        Ok(outputs)
    })
}

pub(crate) fn operation_discard() -> RuntimeOperation {
    RuntimeOperation::new_fn_simple(|_inputs, _context| Ok(HashMap::new()))
}

pub(crate) fn operation_equality() -> RuntimeOperation {
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let val0 = take_input(&mut inputs, "value_0")?;
        let val1 = take_input(&mut inputs, "value_1")?;
        let mut outputs = HashMap::new();
        outputs.insert(TryInto::try_into("result")?, Value::Bool(val0 == val1));
        Ok(outputs)
    })
}

pub(crate) fn operation_not_equality() -> RuntimeOperation {
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let val0 = take_input(&mut inputs, "value_0")?;
        let val1 = take_input(&mut inputs, "value_1")?;
        let mut outputs = HashMap::new();
        outputs.insert(TryInto::try_into("result")?, Value::Bool(val0 != val1));
        Ok(outputs)
    })
}

pub(crate) fn operation_switch() -> RuntimeOperation {
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let predicate = validate_bool_input(&mut inputs, "pred")?;

        let branch_true = take_input(&mut inputs, "if_true")?;
        let branch_false = take_input(&mut inputs, "if_false")?;

        let result = if predicate { branch_true } else { branch_false };

        let mut outputs = HashMap::new();
        outputs.insert(Label::value(), result);
        Ok(outputs)
    })
}

pub(crate) fn operation_make_pair() -> RuntimeOperation {
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let first = take_input(&mut inputs, "first")?;
        let second = take_input(&mut inputs, "second")?;

        let mut outputs = HashMap::new();
        outputs.insert(
            TryInto::try_into("pair")?,
            Value::Pair(Box::new((first, second))),
        );
        Ok(outputs)
    })
}

pub(crate) fn operation_unpack_pair() -> RuntimeOperation {
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let (first, second) = validate_input(&mut inputs, "pair", |x| match x {
            Value::Pair(pair) => Some((pair.0, pair.1)),
            _ => None,
        })?;

        let mut outputs = HashMap::new();
        outputs.insert(TryInto::try_into("first")?, first);
        outputs.insert(TryInto::try_into("second")?, second);
        Ok(outputs)
    })
}

pub(crate) fn operation_push() -> RuntimeOperation {
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let vec_l: Label = TryInto::try_into("vec")?;
        let mut vec = match take_input(&mut inputs, vec_l)? {
            Value::Vec(vec) => vec,
            _ => bail!("Push function expected vector input."),
        };

        let item = take_input(&mut inputs, "item")?;

        vec.push(item);
        let mut outputs = HashMap::new();
        outputs.insert(vec_l, Value::Vec(vec));
        Ok(outputs)
    })
}

pub(crate) fn operation_pop() -> RuntimeOperation {
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let vec_l: Label = TryInto::try_into("vec")?;
        let mut vec = match take_input(&mut inputs, vec_l)? {
            Value::Vec(vec) => {
                if vec.is_empty() {
                    Err(RuntimeError::EmptyVector)
                } else {
                    Ok(vec)
                }
            }
            _ => Err(RuntimeError::InvalidInput(vec_l)),
        }?;

        let item = vec.pop().unwrap();
        let mut outputs = HashMap::new();
        outputs.insert(TryInto::try_into("item")?, item);
        outputs.insert(vec_l, Value::Vec(vec));

        Ok(outputs)
    })
}

pub(crate) fn operation_loop() -> RuntimeOperation {
    RuntimeOperation::new_fn_async(|mut inputs, context| async move {
        let _ = &context;
        let body = validate_graph_input(&mut inputs, "body")?;
        let mut value = take_input(&mut inputs, Label::value())?;

        let node_trace = context
            .graph_trace
            .as_node_trace()
            .map_err(|_| anyhow!("loop function expected stack trace to correspond to a node"))?;
        for iteration in 1.. {
            // Run body
            let graph_trace = node_trace.clone().loop_iter(iteration);

            let body_output = RuntimeOperation::new_graph(body.clone())
                .run_simple(
                    context.runtime.clone(),
                    context.callback.clone(),
                    context.escape.clone(),
                    [(Label::value(), value)],
                    graph_trace,
                    context.checkpoint_client.clone(),
                )
                .into_task()
                .complete()
                .await
                .map_err(|err| {
                    // Formatting with debug gets us more/deeper context than anyhow! alone
                    let e = anyhow!(format!("{:?}", err.as_ref()));
                    e.context(format!("loop body (iteration {})", iteration))
                })?;

            let body_output = body_output.as_ref();
            if let Some(Value::Variant(label, b)) = body_output.get(&Label::value()) {
                value = *b.clone();
                if label == &Label::continue_() {
                    continue;
                } else if label == &Label::break_() {
                    break;
                }
            };
            // Both expected variants will have jumped elsewhere
            bail!(
                "loop node expected body to output a Variant (break | continue) on port 'value' (iteration {})",
                iteration
            )
        }
        Ok(HashMap::from([(Label::value(), value)]))
    })
}

pub(crate) fn operation_sequence() -> RuntimeOperation {
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let first = validate_graph_input(&mut inputs, "first")?;

        let second = validate_graph_input(&mut inputs, "second")?;

        let g3 = {
            let mut builder = GraphBuilder::new();
            let [input, output] = Graph::boundary();
            let inputs: Vec<Edge> = first.node_outputs(input).cloned().collect();
            let second_input_ports: HashSet<_> =
                second.node_outputs(input).map(|e| e.source.port).collect();
            // only wire up shared ports
            let middle: Vec<Edge> = first
                .node_inputs(output)
                .filter(|e| second_input_ports.contains(&e.target.port))
                .cloned()
                .collect();
            let outputs: Vec<Edge> = second.node_inputs(output).cloned().collect();

            let b1 = builder.add_node(Node::local_box(first))?;
            let b2 = builder.add_node(Node::local_box(second))?;

            for input_edge in inputs {
                builder.add_edge(
                    (input, input_edge.source.port),
                    (b1, input_edge.source.port),
                    input_edge.edge_type,
                )?;
            }
            for seq_edge in middle {
                builder.add_edge(
                    (b1, seq_edge.target.port),
                    (b2, seq_edge.target.port),
                    seq_edge.edge_type,
                )?;
            }

            for output_edge in outputs {
                builder.add_edge(
                    (b2, output_edge.target.port),
                    (output, output_edge.target.port),
                    output_edge.edge_type,
                )?;
            }
            builder.build()?
        };

        let mut outputs = HashMap::new();
        outputs.insert(TryInto::try_into("sequenced")?, Value::Graph(g3));

        Ok(outputs)
    })
}

pub(crate) fn operation_parallel() -> RuntimeOperation {
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let left = validate_graph_input(&mut inputs, "left")?;

        let right = validate_graph_input(&mut inputs, "right")?;

        let g3 = {
            let mut builder = GraphBuilder::new();
            let [input, output] = Graph::boundary();

            let inputs_left: Vec<Edge> = left.node_outputs(input).cloned().collect();
            let inputs_right: Vec<Edge> = right.node_outputs(input).cloned().collect();
            let outputs_left: Vec<Edge> = left.node_inputs(output).cloned().collect();
            let outputs_right: Vec<Edge> = right.node_inputs(output).cloned().collect();

            let b_left = builder.add_node(Node::local_box(left))?;
            let b_right = builder.add_node(Node::local_box(right))?;

            for left_input_edge in inputs_left {
                builder.add_edge(
                    (input, left_input_edge.source.port),
                    (b_left, left_input_edge.source.port),
                    left_input_edge.edge_type,
                )?;
            }

            for right_input_edge in inputs_right {
                builder.add_edge(
                    (input, right_input_edge.source.port),
                    (b_right, right_input_edge.source.port),
                    right_input_edge.edge_type,
                )?;
            }

            for left_output_edge in outputs_left {
                builder.add_edge(
                    (b_left, left_output_edge.target.port),
                    (output, left_output_edge.target.port),
                    left_output_edge.edge_type,
                )?;
            }

            for right_output_edge in outputs_right {
                builder.add_edge(
                    (b_right, right_output_edge.target.port),
                    (output, right_output_edge.target.port),
                    right_output_edge.edge_type,
                )?;
            }

            builder.build()?
        };

        let outputs = HashMap::from([(Label::value(), Value::Graph(g3))]);

        Ok(outputs)
    })
}

pub(crate) fn operation_make_struct() -> RuntimeOperation {
    RuntimeOperation::new_fn_simple(|inputs, _context| {
        let struc = Value::Struct(inputs);
        let mut outputs = HashMap::new();
        outputs.insert(TryInto::try_into("struct")?, struc);
        Ok(outputs)
    })
}

pub(crate) fn operation_unpack_struct() -> RuntimeOperation {
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let outputs = validate_input(&mut inputs, "struct", |x| match x {
            Value::Struct(fields) => Some(fields),
            _ => None,
        })?;
        Ok(outputs)
    })
}

pub(crate) fn operation_insert_key() -> RuntimeOperation {
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let map_l: Label = TryInto::try_into("map")?;
        let mut map = validate_input(&mut inputs, map_l, |x| match x {
            Value::Map(map) => Some(map),
            _ => None,
        })?;

        let key = take_input(&mut inputs, "key")?;
        let val = take_input(&mut inputs, "val")?;

        map.insert(key, val);
        let mut outputs = HashMap::new();
        outputs.insert(map_l, Value::Map(map));
        Ok(outputs)
    })
}

pub(crate) fn operation_remove_key() -> RuntimeOperation {
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let map_l: Label = TryInto::try_into("map")?;
        let mut map = validate_input(&mut inputs, map_l, |x| match x {
            Value::Map(map) => Some(map),
            _ => None,
        })?;

        let key = take_input(&mut inputs, "key")?;
        let val = map.remove(&key).ok_or(RuntimeError::KeyNotFound(key))?;
        let mut outputs = HashMap::new();
        outputs.insert(map_l, Value::Map(map));
        outputs.insert(TryInto::try_into("val")?, val);
        Ok(outputs)
    })
}

// numeric operations

pub(crate) fn binary_int_operation_with_error<F>(f: F) -> RuntimeOperation
where
    F: FnOnce(i64, i64) -> anyhow::Result<i64> + Sync + Send + 'static,
{
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let a = validate_int_input(&mut inputs, "a")?;
        let b = validate_int_input(&mut inputs, "b")?;

        let result = f(a, b)?;
        let mut outputs = HashMap::new();
        outputs.insert(Label::value(), Value::Int(result));
        Ok(outputs)
    })
}

pub(crate) fn binary_int_operation<F>(f: F) -> RuntimeOperation
where
    F: FnOnce(i64, i64) -> i64 + Sync + Send + 'static,
{
    binary_int_operation_with_error(|a, b| Ok(f(a, b)))
}

pub(crate) fn binary_flt_operation<F>(f: F) -> RuntimeOperation
where
    F: FnOnce(f64, f64) -> f64 + Sync + Send + 'static,
{
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let a = validate_float_input(&mut inputs, "a")?;
        let b = validate_float_input(&mut inputs, "b")?;

        let mut outputs = HashMap::new();
        outputs.insert(Label::value(), Value::Float(f(a, b)));
        Ok(outputs)
    })
}

pub(crate) fn binary_int_comparison<F>(f: F) -> RuntimeOperation
where
    F: FnOnce(i64, i64) -> bool + Sync + Send + 'static,
{
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let a = validate_int_input(&mut inputs, "a")?;
        let b = validate_int_input(&mut inputs, "b")?;

        let mut outputs = HashMap::new();
        outputs.insert(Label::value(), Value::Bool(f(a, b)));
        Ok(outputs)
    })
}

pub(crate) fn binary_flt_comparison<F>(f: F) -> RuntimeOperation
where
    F: FnOnce(f64, f64) -> bool + Sync + Send + 'static,
{
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let a = validate_float_input(&mut inputs, "a")?;
        let b = validate_float_input(&mut inputs, "b")?;

        let mut outputs = HashMap::new();
        outputs.insert(Label::value(), Value::Bool(f(a, b)));
        Ok(outputs)
    })
}

pub(crate) fn binary_bool_operation<F>(f: F) -> RuntimeOperation
where
    F: FnOnce(bool, bool) -> bool + Sync + Send + 'static,
{
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let a = validate_bool_input(&mut inputs, "a")?;
        let b = validate_bool_input(&mut inputs, "b")?;

        let mut outputs = HashMap::new();
        outputs.insert(Label::value(), Value::Bool(f(a, b)));
        Ok(outputs)
    })
}

pub(crate) fn operation_int_to_float() -> RuntimeOperation {
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let int = validate_int_input(&mut inputs, "int")?;

        let mut outputs = HashMap::new();
        outputs.insert(Label::value(), Value::Float(int as f64));
        Ok(outputs)
    })
}

pub(crate) fn operation_float_to_int() -> RuntimeOperation {
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let flt = validate_float_input(&mut inputs, "float")?;

        let mut outputs = HashMap::new();
        outputs.insert(Label::value(), Value::Int(flt as i64));
        Ok(outputs)
    })
}

pub(crate) fn operation_partial() -> RuntimeOperation {
    RuntimeOperation::new_fn_simple(|mut inputs, _context| {
        let thunk = validate_graph_input(&mut inputs, Label::thunk())?;

        let g3 = {
            let mut builder = GraphBuilder::new();
            let [input, output] = Graph::boundary();

            let input_edges: Vec<Edge> = thunk.node_outputs(input).cloned().collect();
            let output_edges: Vec<Edge> = thunk.node_inputs(output).cloned().collect();
            let b1 = builder.add_node(Node::local_box(thunk))?;

            // Note that if values are provided to 'partial' that do NOT
            // correspond to inputs to the thunk, this code will drop them here
            // (at closure-creation time), whereas the old code would have
            // passed them into the closure (which I think would have ignored them)
            for input_edge in input_edges {
                let port = input_edge.source.port;
                let source = match inputs.remove(&port) {
                    Some(value) => {
                        let new_const = builder.add_node(Node::Const(value))?;
                        (new_const, Label::value())
                    }
                    None => (input, port),
                };
                builder.add_edge(source, (b1, port), input_edge.edge_type)?;
            }
            //So: are there any ports in inputs.keys() that do not have
            // an edge in input_edges? If so, maybe raise an error?

            for out_edge in output_edges {
                builder.add_edge(
                    (b1, out_edge.target.port),
                    (output, out_edge.target.port),
                    out_edge.edge_type,
                )?;
            }

            builder.build()?
        };

        let mut outputs = HashMap::new();
        outputs.insert(Label::value(), Value::Graph(g3));

        Ok(outputs)
    })
}

pub(crate) fn operation_map() -> RuntimeOperation {
    RuntimeOperation::new_fn_async(|mut inputs, context| async move {
        let _ = &context;
        let thunk = validate_graph_input(&mut inputs, Label::thunk())?;

        let list = validate_input(&mut inputs, Label::value(), |x| match x {
            Value::Vec(a) => Some(a),
            _ => None,
        })?;

        let mut tasks = Vec::new();

        let node_trace = context
            .graph_trace
            .as_node_trace()
            .map_err(|_| anyhow!("map function expected stack trace to correspond to a node"))?;
        for (idx, x) in list.into_iter().enumerate() {
            let thunk_clone = thunk.clone();
            let runtime = context.runtime.clone();
            let callback = context.callback.clone();
            let escape = context.escape.clone();
            let checkpoint = context.checkpoint_client.clone();
            let graph_trace = node_trace.clone().list_elem(idx as u32);

            let span = tracing::Span::current();

            let t = tokio::spawn(
                async move {
                    let value: Value = RuntimeOperation::new_graph(thunk_clone)
                        .run_simple(
                            runtime,
                            callback,
                            escape,
                            [(Label::value(), x)],
                            graph_trace,
                            checkpoint,
                        )
                        .into_task()
                        .complete()
                        .await
                        .map_err(|err| {
                            let e = anyhow!(format!("{:?}", err.as_ref()));
                            e.context("map body".to_string())
                        })?
                        .get(&Label::value())
                        .ok_or_else(|| anyhow!("map thunk should output on value port"))?
                        .clone();
                    Ok::<Value, anyhow::Error>(value)
                }
                .instrument(span),
            );
            tasks.push(t);
        }
        // .collect();
        let x = future::join_all(tasks).await;
        let y: Result<Vec<_>, _> = x.into_iter().collect();
        let z: Result<Vec<Value>, _> = y?.into_iter().collect();

        let outputs = HashMap::from([(Label::value(), Value::Vec(z?))]);

        Ok(outputs)
    })
}

fn take_input<E>(
    inputs: &mut HashMap<Label, Value>,
    port: impl TryInto<Label, Error = E>,
) -> anyhow::Result<Value>
where
    E: Into<SymbolError>,
{
    let port = TryInto::try_into(port).map_err(|e| e.into())?;
    inputs
        .remove(&port)
        .ok_or_else(|| RuntimeError::MissingInput(port).into())
}

fn validate_input<E, T>(
    inputs: &mut HashMap<Label, Value>,
    port: impl TryInto<Label, Error = E>,
    validation: impl FnOnce(Value) -> Option<T>,
) -> anyhow::Result<T>
where
    E: Into<SymbolError>,
{
    let port = TryInto::try_into(port).map_err(|e| e.into())?;
    let input = inputs
        .remove(&port)
        .ok_or(RuntimeError::MissingInput(port))?;
    match validation(input) {
        Some(v) => Ok(v),
        None => Err(anyhow!(RuntimeError::InvalidInput(port))),
    }
}

fn validate_int_input<E>(
    inputs: &mut HashMap<Label, Value>,
    port: impl TryInto<Label, Error = E>,
) -> anyhow::Result<i64>
where
    E: Into<SymbolError>,
{
    validate_input(inputs, port, |x| match x {
        Value::Int(a) => Some(a),
        _ => None,
    })
}

fn validate_float_input<E>(
    inputs: &mut HashMap<Label, Value>,
    port: impl TryInto<Label, Error = E>,
) -> anyhow::Result<f64>
where
    E: Into<SymbolError>,
{
    validate_input(inputs, port, |x| match x {
        Value::Float(a) => Some(a),
        _ => None,
    })
}

fn validate_bool_input<E>(
    inputs: &mut HashMap<Label, Value>,
    port: impl TryInto<Label, Error = E>,
) -> anyhow::Result<bool>
where
    E: Into<SymbolError>,
{
    validate_input(inputs, port, |x| match x {
        Value::Bool(a) => Some(a),
        _ => None,
    })
}

fn validate_graph_input<E>(
    inputs: &mut HashMap<Label, Value>,
    port: impl TryInto<Label, Error = E>,
) -> anyhow::Result<Graph>
where
    E: Into<SymbolError>,
{
    validate_input(inputs, port, |x| match x {
        Value::Graph(a) => Some(a),
        _ => None,
    })
}

#[derive(Debug, Error)]
enum RuntimeError {
    #[error("Missing input on port {0}.")]
    MissingInput(Label),
    #[error("Invalid input on port {0}.")]
    InvalidInput(Label),
    #[error("Vector is empty.")]
    EmptyVector,
    #[error("Key not found in map.")]
    KeyNotFound(Value),
}