async_nodes/

run.rs

//! Run-scoped state: the value store, the schedulable snapshot [`RunState`], the
//! node-facing [`RunContext`], and suspension packaging.
//!
//! The engine owns one trunk [`RunState`] per run. It carries the consumer state `S`
//! (`user`), the pin-value store (the dataflow profile's IO surface), and all activation
//! accounting — pending tokens, back-edge iteration counters, per-node execution counters,
//! the current frontier, and the superstep counter. Because *all* of that lives in one
//! struct, suspending mid-loop and resuming later is just "stop, keep the struct, continue"
//! — loop counters continue from where they were, they do not reset.

use std::any::Any;

use indexmap::{IndexMap, IndexSet};
use serde::{Deserialize, Serialize};

use super::{
    BspContext,
    error::BspError,
    graph::{EdgeId, InputPin, NodeId, OutputPin, PortIndex},
    value::{BspOutputMapExt, BspValue},
};

/// The engine-owned store of pin values (the dataflow profile's IO surface).
///
/// Values are type-erased (`Arc<dyn Any>`) and are **never serialized** — suspension
/// snapshots restore pin values through `waiting_inputs` re-injection only.
#[derive(Debug, Default)]
pub struct ValueStore {
    pub(crate) outputs: IndexMap<OutputPin, BspValue>,
    pub(crate) injected: IndexMap<InputPin, BspValue>,
}

impl ValueStore {
    /// Returns the value currently held by an output pin, if any.
    #[must_use]
    pub fn get_output(&self, pin: OutputPin) -> Option<&BspValue> {
        self.outputs.get(&pin)
    }

    /// Returns the value currently held by an output pin, downcast to `T`.
    #[must_use]
    pub fn get_output_typed<T: Any + Send + Sync + Clone>(&self, pin: OutputPin) -> Option<T> {
        self.outputs.get(&pin)?.downcast_ref::<T>().cloned()
    }

    /// The full map of output-pin values written so far.
    #[must_use]
    pub const fn outputs(&self) -> &IndexMap<OutputPin, BspValue> {
        &self.outputs
    }

    /// Returns a value injected into an input pin (e.g. by a resume injection), if any.
    #[must_use]
    pub fn get_injected(&self, pin: InputPin) -> Option<&BspValue> {
        self.injected.get(&pin)
    }

    /// The full map of injected input-pin values.
    #[must_use]
    pub const fn injected(&self) -> &IndexMap<InputPin, BspValue> {
        &self.injected
    }

    /// Writes a typed value into an output pin.
    ///
    /// Useful for restoring linear state profiles after out-of-band deserialization.
    pub fn set_output_typed<T: Any + Send + Sync + 'static>(&mut self, pin: OutputPin, value: T) {
        self.outputs.insert(pin, BspValue::new(value));
    }

    /// Injects a typed value into an input pin.
    pub fn set_injected_typed<T: Any + Send + Sync + 'static>(&mut self, pin: InputPin, value: T) {
        self.injected.insert(pin, BspValue::new(value));
    }
}

/// The complete schedulable state of one run: consumer state, value store, and all
/// activation accounting. Suspension, resumption, and step-limit results all package this.
#[derive(Debug)]
pub struct RunState<S> {
    /// The consumer-owned state. The engine moves it; it never inspects it.
    pub user: S,
    pub(crate) values: ValueStore,
    /// Pending (unconsumed) activation tokens per edge. Edge ids are declaration indices.
    pub(crate) pending_tokens: IndexMap<EdgeId, u64>,
    /// Per-edge iteration counters for edges carrying a `max_iterations` budget.
    pub(crate) iterations: IndexMap<EdgeId, u64>,
    /// Per-node execution counters (counts `execute` invocations, including retries).
    pub(crate) executions: IndexMap<NodeId, u64>,
    /// The node errors recorded so far (populated by `RouteToErrorEdge` recovery).
    pub errors: IndexMap<NodeId, BspError>,
    /// Nodes scheduled for the current superstep (used to re-enter after suspension).
    pub(crate) frontier: IndexSet<NodeId>,
    /// Number of supersteps executed so far (counts toward `max_steps`).
    pub step: u64,
}

impl<S> RunState<S> {
    /// Creates the initial run state for a fresh run.
    pub(crate) fn new(user: S) -> Self {
        Self {
            user,
            values: ValueStore::default(),
            pending_tokens: IndexMap::new(),
            iterations: IndexMap::new(),
            executions: IndexMap::new(),
            errors: IndexMap::new(),
            frontier: IndexSet::new(),
            step: 0,
        }
    }

    /// The engine-owned value store.
    #[must_use]
    pub const fn values(&self) -> &ValueStore {
        &self.values
    }

    /// Mutable access to the engine-owned value store (e.g. to restore pin state out of band).
    pub const fn values_mut(&mut self) -> &mut ValueStore {
        &mut self.values
    }

    /// Per-edge iteration counters, keyed by edge declaration index.
    #[must_use]
    pub const fn iterations(&self) -> &IndexMap<EdgeId, u64> {
        &self.iterations
    }

    /// Per-node execution counters (counts `execute` invocations, including retries).
    #[must_use]
    pub const fn executions(&self) -> &IndexMap<NodeId, u64> {
        &self.executions
    }

    /// How many times a node has executed so far.
    #[must_use]
    pub fn execution_count(&self, node_id: NodeId) -> u64 {
        self.executions.get(&node_id).copied().unwrap_or(0)
    }

    /// How many times an edge has fired against its iteration budget so far.
    #[must_use]
    pub fn iteration_count(&self, edge: EdgeId) -> u64 {
        self.iterations.get(&edge).copied().unwrap_or(0)
    }

    /// Retrieves a typed value associated with an output pin, if it exists.
    #[must_use]
    pub fn get_output_typed<T: Any + Send + Sync + Clone>(&self, pin: OutputPin) -> Option<T> {
        self.values.get_output_typed(pin)
    }
}

/// Everything a node may touch, nothing it may not.
///
/// The engine hands one of these to each `execute` call. Nodes mutate the consumer state
/// through [`RunContext::state`], read gathered pin inputs, and stage pin outputs; they
/// **never name their successors** — the engine decides what happens next from the edges.
///
/// # Retry & rollback contract
///
/// When a node carries a retry policy, a failed attempt is re-invoked with the *same*
/// gathered inputs and an incremented [`RunContext::execution_no`]. A node that stages
/// writes into `state` should therefore treat `execution_no > 0` as "undo my partial
/// writes from the previous attempt first" — that is the exactly-once-ish rollback hook.
pub struct RunContext<'a, S> {
    /// The consumer state (history, stores, metrics…). `&mut`, owned by the engine
    /// between executions — required for fork/merge and suspension snapshots.
    pub state: &'a mut S,
    /// The id of the node being executed.
    pub node_id: NodeId,
    /// Dependency-injection context (unchanged from the dataflow engine).
    pub ctx: &'a BspContext,
    /// This node's Nth `execute` invocation, 0-based (loop- and retry-aware).
    pub execution_no: u64,
    pub(crate) inputs: IndexMap<PortIndex, BspValue>,
    pub(crate) outputs: IndexMap<PortIndex, BspValue>,
    pub(crate) skip: bool,
}

impl<'a, S> RunContext<'a, S> {
    pub(crate) fn new(
        state: &'a mut S,
        node_id: NodeId,
        ctx: &'a BspContext,
        execution_no: u64,
        inputs: IndexMap<PortIndex, BspValue>,
    ) -> Self {
        Self {
            state,
            node_id,
            ctx,
            execution_no,
            inputs,
            outputs: IndexMap::new(),
            skip: false,
        }
    }

    /// Reads a gathered input value, downcast to `T`, without consuming it.
    #[must_use]
    pub fn input<T: Any + Send + Sync + Clone>(&self, port: PortIndex) -> Option<T> {
        self.inputs.get(&port)?.downcast_ref::<T>().cloned()
    }

    /// Reads a gathered input value without consuming it.
    #[must_use]
    pub fn input_value(&self, port: PortIndex) -> Option<&BspValue> {
        self.inputs.get(&port)
    }

    /// Takes a gathered input value out of the context (dataflow profile; used by the
    /// `#[node]` macros to unpack named ports).
    pub fn take_input(&mut self, port: PortIndex) -> Option<BspValue> {
        self.inputs.swap_remove(&port)
    }

    /// Stages a typed value for an output port.
    ///
    /// Values staged here merge with (and are overridden by) the port map returned in
    /// [`crate::node::NodeOutcome::Success`].
    pub fn set_output<T: Any + Send + Sync>(&mut self, port: PortIndex, value: T) {
        self.outputs.insert(port, BspValue::new(value));
    }

    /// Marks every declared output port to be written with the `BspSkip` marker
    /// (dataflow profile skip propagation), regardless of any staged outputs.
    pub fn skip_outputs(&mut self) {
        self.skip = true;
    }

    /// Drains the staged outputs and the skip flag after execution.
    pub(crate) fn take_staged(&mut self) -> (IndexMap<PortIndex, BspValue>, bool) {
        (std::mem::take(&mut self.outputs), self.skip)
    }
}

/// Represents the return status of a BSP execution run.
///
/// This is the single, unified result type of every entry point
/// ([`crate::engine::BspEngine::run`], `resume`, and their sequential counterparts).
/// A tripped `max_steps` budget is a first-class outcome carrying the live run state —
/// never an `Ok` with silently salvaged state, and never an error: hitting a circuit
/// breaker is an expected operational outcome for bounded loops.
#[derive(Debug)]
pub enum EngineResult<S = ()> {
    /// The graph ran to completion: the frontier emptied with no pending tokens left.
    Complete(Box<RunState<S>>),
    /// At least one node requested suspension; the packaged snapshot can be resumed.
    Suspended(Box<SuspendedExecution<S>>),
    /// The global `max_steps` superstep budget was exhausted before natural completion.
    StepLimitReached(Box<RunState<S>>),
}

/// Package containing suspension parameters to resume execution later.
#[derive(Debug)]
pub struct SuspendedExecution<S = ()> {
    /// The nodes waiting for external input, with their payloads and expected pins.
    pub suspensions: Vec<PendingSuspension>,
    pub(crate) state: RunState<S>,
}

/// Single record of a suspended node waiting for external inputs.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PendingSuspension {
    /// The node that requested suspension.
    pub node_id: NodeId,
    /// Explicit indices of inputs that must be provided to resume.
    pub waiting_inputs: Vec<PortIndex>,
    /// Arbitrary context parameters passed back to the host application.
    pub payload: serde_json::Value,
}

impl<S> SuspendedExecution<S> {
    /// Returns an immutable reference to the inner run state.
    #[must_use]
    pub const fn state(&self) -> &RunState<S> {
        &self.state
    }

    /// Returns a mutable reference to the inner run state.
    pub const fn state_mut(&mut self) -> &mut RunState<S> {
        &mut self.state
    }

    /// Consumes the package, returning the inner run state.
    #[must_use]
    pub fn into_state(self) -> RunState<S> {
        self.state
    }
}

/// The serializable subset of a suspended run.
///
/// **The value store is deliberately excluded** (see [`SuspendedExecution::to_json`]).
#[derive(Serialize, Deserialize)]
struct SuspensionSnapshot<S> {
    suspensions: Vec<PendingSuspension>,
    user: S,
    pending_tokens: IndexMap<EdgeId, u64>,
    iterations: IndexMap<EdgeId, u64>,
    executions: IndexMap<NodeId, u64>,
    frontier: Vec<NodeId>,
    step: u64,
}

impl<S> SuspendedExecution<S>
where
    S: Serialize + for<'de> Deserialize<'de>,
{
    /// Serializes the suspension snapshot to JSON.
    ///
    /// **The pin-value store is NOT serialized.** `BspValue` is type-erased
    /// (`Arc<dyn Any>`) and has no serialized form; pretending otherwise would silently
    /// corrupt resume fidelity. The snapshot covers the consumer state `S`, token
    /// accounting, iteration/execution counters, the frontier, and the suspension list.
    /// On resume, dataflow-profile pin values restore exclusively through `waiting_inputs`
    /// re-injection (plus any consumer-managed restore of `S`).
    ///
    /// Resuming a deserialized snapshot requires a `BspEngine` built over an identical
    /// graph: node and edge ids are positional.
    pub fn to_json(&self) -> Result<String, BspError> {
        let snapshot = SuspensionSnapshot {
            suspensions: self.suspensions.clone(),
            user: &self.state.user,
            pending_tokens: self.state.pending_tokens.clone(),
            iterations: self.state.iterations.clone(),
            executions: self.state.executions.clone(),
            frontier: self.state.frontier.iter().copied().collect(),
            step: self.state.step,
        };
        serde_json::to_string(&snapshot).map_err(|e| BspError::SuspensionSerde(e.to_string()))
    }

    /// Restores a suspension snapshot from JSON produced by [`SuspendedExecution::to_json`].
    ///
    /// The value store comes back empty: pin values must be re-supplied through resume
    /// injections keyed by each suspension's `waiting_inputs`.
    pub fn from_json(json: &str) -> Result<Self, BspError> {
        let snapshot: SuspensionSnapshot<S> =
            serde_json::from_str(json).map_err(|e| BspError::SuspensionSerde(e.to_string()))?;
        Ok(Self {
            suspensions: snapshot.suspensions,
            state: RunState {
                user: snapshot.user,
                values: ValueStore::default(),
                pending_tokens: snapshot.pending_tokens,
                iterations: snapshot.iterations,
                executions: snapshot.executions,
                errors: IndexMap::new(),
                frontier: snapshot.frontier.into_iter().collect(),
                step: snapshot.step,
            },
        })
    }
}

impl<S> BspOutputMapExt for RunState<S> {
    /// Typed access to the value store's output map, so a completed run reads exactly like
    /// the dataflow engine's output map did.
    fn get_typed<T: Any + Send + Sync + Clone>(&self, pin: OutputPin) -> Result<T, BspError> {
        self.values.outputs().get_typed(pin)
    }
}

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

    #[test]
    fn test_value_store_roundtrip() {
        let mut store = ValueStore::default();
        let pin = OutputPin {
            node_id: 1,
            port: 0,
        };
        store.set_output_typed(pin, 42_i64);
        assert_eq!(store.get_output_typed::<i64>(pin), Some(42));
        assert_eq!(store.get_output_typed::<String>(pin), None);

        let input = InputPin {
            node_id: 2,
            port: 1,
        };
        store.set_injected_typed(input, true);
        assert!(store.get_injected(input).is_some());
    }

    #[test]
    fn test_run_state_counters() {
        let mut state = RunState::new(());
        state.executions.insert(3, 7);
        state.iterations.insert(9, 2);
        assert_eq!(state.execution_count(3), 7);
        assert_eq!(state.execution_count(4), 0);
        assert_eq!(state.iteration_count(9), 2);
        assert_eq!(state.iteration_count(10), 0);
    }

    #[test]
    fn test_run_context_io() {
        let mut user = 0_u64;
        let ctx = BspContext::new();
        let mut inputs = IndexMap::new();
        inputs.insert(0, BspValue::new(21_i64));
        {
            let mut run = RunContext::new(&mut user, 5, &ctx, 3, inputs);
            assert_eq!(run.node_id, 5);
            assert_eq!(run.execution_no, 3);
            assert_eq!(run.input::<i64>(0), Some(21));
            assert_eq!(
                run.take_input(0)
                    .and_then(|v| v.downcast::<i64>().ok())
                    .map(|a| *a),
                Some(21)
            );
            assert!(run.input_value(0).is_none());
            run.set_output(1, "done".to_string());
            run.skip_outputs();
            let (staged, skip) = run.take_staged();
            assert!(skip);
            assert!(staged.contains_key(&1));
        }
    }

    #[test]
    fn test_suspension_snapshot_roundtrip_preserves_counters() {
        let mut state: RunState<String> = RunState::new("consumer-state".to_string());
        state.pending_tokens.insert(4, 2);
        state.iterations.insert(4, 5);
        state.executions.insert(1, 6);
        state.frontier.insert(1);
        state.step = 11;
        let suspended = SuspendedExecution {
            suspensions: vec![PendingSuspension {
                node_id: 1,
                waiting_inputs: vec![0],
                payload: serde_json::json!({"why": "test"}),
            }],
            state,
        };

        let json = suspended.to_json().unwrap();
        // The value store is never serialized.
        assert!(!json.contains("outputs"));

        let restored = SuspendedExecution::<String>::from_json(&json).unwrap();
        assert_eq!(restored.state.user, "consumer-state");
        assert_eq!(restored.state.iteration_count(4), 5);
        assert_eq!(restored.state.execution_count(1), 6);
        assert_eq!(restored.state.pending_tokens.get(&4), Some(&2));
        assert_eq!(restored.state.step, 11);
        assert!(restored.state.frontier.contains(&1));
        assert_eq!(restored.suspensions[0].waiting_inputs, vec![0]);
    }

    #[test]
    fn test_suspension_snapshot_rejects_garbage() {
        let err = SuspendedExecution::<String>::from_json("{not json").unwrap_err();
        assert!(matches!(err, BspError::SuspensionSerde(_)));
    }
}