async-nodes 0.1.0

//! Core scheduler integration tests: chains, joins, step limits, bounded concurrency,
//! the event stream, and the sequential (!Clone) API.

use std::sync::{
    Arc, Mutex,
    atomic::{AtomicUsize, Ordering},
};

use async_nodes::{
    Activation, BspContext, BspEngine, BspError, BspGraph, BspInterceptor, BspNode, EngineResult,
    NodeOutcome, RunContext, async_trait::async_trait,
};
#[cfg(feature = "tokio")]
use async_nodes::{BspEvent, ChannelInterceptor};
use indexmap::IndexMap;

// =========================================================================
// TEST NODES
// =========================================================================

/// Adds a constant to a `u64` state whenever it fires.
struct AddNode {
    name: &'static str,
    amount: u64,
}

#[async_trait]
impl BspNode<u64> for AddNode {
    fn name(&self) -> &str {
        self.name
    }

    async fn execute(&self, run: &mut RunContext<'_, u64>) -> Result<NodeOutcome, BspError> {
        *run.state += self.amount;
        Ok(NodeOutcome::Success(IndexMap::new()))
    }
}

/// Does nothing; carries a configurable activation policy.
struct NoopNode {
    name: &'static str,
    activation: Activation,
}

#[async_trait]
impl<S: Send> BspNode<S> for NoopNode {
    fn name(&self) -> &str {
        self.name
    }

    fn activation(&self) -> Activation {
        self.activation
    }

    async fn execute(&self, _run: &mut RunContext<'_, S>) -> Result<NodeOutcome, BspError> {
        Ok(NodeOutcome::Success(IndexMap::new()))
    }
}

/// Records its `execution_no` into a Vec state each time it fires.
struct TraceNode {
    name: &'static str,
}

#[async_trait]
impl BspNode<Vec<u64>> for TraceNode {
    fn name(&self) -> &str {
        self.name
    }

    async fn execute(&self, run: &mut RunContext<'_, Vec<u64>>) -> Result<NodeOutcome, BspError> {
        run.state.push(run.execution_no);
        Ok(NodeOutcome::Success(IndexMap::new()))
    }
}

// =========================================================================
// LINEAR CHAINS & JOIN SEMANTICS
// =========================================================================

#[tokio::test]
async fn test_linear_chain_runs_in_declaration_order() {
    let mut graph = BspGraph::new();
    let a = graph.add_node(AddNode {
        name: "a",
        amount: 1,
    });
    let b = graph.add_node(AddNode {
        name: "b",
        amount: 10,
    });
    let c = graph.add_node(AddNode {
        name: "c",
        amount: 100,
    });
    graph.add_edge(a.id, b.id);
    graph.add_edge(b.id, c.id);

    let engine = BspEngine::new(graph);
    let outcome = engine.run(0_u64).await.unwrap();

    let EngineResult::Complete(state) = outcome else {
        panic!("expected completion, got {outcome:?}");
    };
    assert_eq!(state.user, 111);
    assert_eq!(state.step, 3);
    assert_eq!(state.execution_count(a.id), 1);
    assert_eq!(state.execution_count(b.id), 1);
    assert_eq!(state.execution_count(c.id), 1);
}

#[tokio::test]
async fn test_diamond_join_all_fires_once() {
    let mut graph = BspGraph::new();
    let src = graph.add_node(NoopNode {
        name: "src",
        activation: Activation::Any,
    });
    let a = graph.add_node(NoopNode {
        name: "a",
        activation: Activation::Any,
    });
    let b = graph.add_node(NoopNode {
        name: "b",
        activation: Activation::Any,
    });
    let join = graph.add_node(NoopNode {
        name: "join",
        activation: Activation::All,
    });
    graph.add_edge(src.id, a.id);
    graph.add_edge(src.id, b.id);
    graph.add_edge(a.id, join.id);
    graph.add_edge(b.id, join.id);

    let engine = BspEngine::new(graph);
    let EngineResult::Complete(state) = engine.run(()).await.unwrap() else {
        panic!("expected completion");
    };
    // The join waits for both branch tokens, then fires exactly once.
    assert_eq!(state.execution_count(join.id), 1);
    assert_eq!(state.execution_count(a.id), 1);
    assert_eq!(state.execution_count(b.id), 1);
}

#[tokio::test]
async fn test_diamond_join_any_fires_per_arriving_token() {
    let mut graph = BspGraph::new();
    let src = graph.add_node(NoopNode {
        name: "src",
        activation: Activation::Any,
    });
    let a = graph.add_node(NoopNode {
        name: "a",
        activation: Activation::Any,
    });
    let b = graph.add_node(NoopNode {
        name: "b",
        activation: Activation::Any,
    });
    let join = graph.add_node(NoopNode {
        name: "join",
        activation: Activation::Any,
    });
    graph.add_edge(src.id, a.id);
    graph.add_edge(src.id, b.id);
    graph.add_edge(a.id, join.id);
    graph.add_edge(b.id, join.id);

    let engine = BspEngine::new(graph);
    let EngineResult::Complete(state) = engine.run(()).await.unwrap() else {
        panic!("expected completion");
    };
    // `Any` fires once per pending token: both branch arrivals trigger it.
    assert_eq!(state.execution_count(join.id), 2);
}

// =========================================================================
// STEP LIMITS
// =========================================================================

#[tokio::test]
async fn test_step_limit_is_a_distinct_result() {
    // Unbounded self-loop: only the global budget can stop it.
    let mut graph = BspGraph::new();
    let looper = graph.add_node(NoopNode {
        name: "looper",
        activation: Activation::Any,
    });
    graph.add_edge(looper.id, looper.id);
    graph.add_entry(looper.id);

    let engine = BspEngine::new(graph).with_max_steps(7);
    let outcome = engine.run(()).await.unwrap();

    let EngineResult::StepLimitReached(state) = outcome else {
        panic!("expected StepLimitReached, got {outcome:?}");
    };
    assert_eq!(state.step, 7);
    assert_eq!(state.execution_count(looper.id), 7);

    // The unbounded cycle is a validation warning, not an error.
    let report = engine.validate_with_report();
    assert!(report.is_ok());
    assert!(
        report
            .warnings
            .iter()
            .any(|w| w.contains("no max_iterations budget"))
    );
}

// =========================================================================
// BOUNDED CONCURRENCY
// =========================================================================

async fn run_probe(max_concurrency: Option<usize>) -> usize {
    /// Shared in-flight instrumentation, injected through the DI context.
    struct Counters {
        in_flight: AtomicUsize,
        max_seen: AtomicUsize,
    }

    struct SharedCounters(Arc<Counters>);

    impl async_nodes::typemap_rev::TypeMapKey for SharedCounters {
        type Value = Self;
    }

    struct CounterProbeNode {
        name: &'static str,
    }

    #[async_trait]
    impl BspNode<()> for CounterProbeNode {
        fn name(&self) -> &str {
            self.name
        }

        async fn execute(&self, run: &mut RunContext<'_, ()>) -> Result<NodeOutcome, BspError> {
            let counters = run
                .ctx
                .get::<SharedCounters>()
                .expect("probe injected")
                .0
                .clone();
            let current = counters.in_flight.fetch_add(1, Ordering::SeqCst) + 1;
            counters.max_seen.fetch_max(current, Ordering::SeqCst);
            // Two yield points so co-scheduled tasks overlap measurably.
            tokio::task::yield_now().await;
            tokio::task::yield_now().await;
            counters.in_flight.fetch_sub(1, Ordering::SeqCst);
            Ok(NodeOutcome::Success(IndexMap::new()))
        }
    }

    let probe = Arc::new(Counters {
        in_flight: AtomicUsize::new(0),
        max_seen: AtomicUsize::new(0),
    });
    let ctx = BspContext::new();
    ctx.insert::<SharedCounters>(SharedCounters(probe.clone()));

    let mut graph = BspGraph::new();
    for name in ["p1", "p2", "p3", "p4"] {
        graph.add_node(CounterProbeNode { name });
    }
    let mut engine = BspEngine::new(graph).with_context(ctx);
    if let Some(limit) = max_concurrency {
        engine = engine.with_max_concurrency(limit);
    }
    let EngineResult::Complete(_) = engine.run(()).await.unwrap() else {
        panic!("expected completion");
    };
    probe.max_seen.load(Ordering::SeqCst)
}

#[tokio::test]
async fn test_max_concurrency_bounds_in_flight_tasks() {
    assert_eq!(run_probe(Some(1)).await, 1);
    assert_eq!(run_probe(Some(2)).await, 2);
}

#[tokio::test]
async fn test_unbounded_frontier_runs_fully_concurrent() {
    assert_eq!(run_probe(None).await, 4);
}

// =========================================================================
// EVENT STREAM
// =========================================================================

struct RecordingInterceptor {
    events: Mutex<Vec<String>>,
}

impl RecordingInterceptor {
    fn record(&self, event: String) {
        self.events.lock().unwrap().push(event);
    }
}

impl BspInterceptor for RecordingInterceptor {
    fn on_superstep_start(&self, step: u64, frontier: &[async_nodes::NodeId]) {
        self.record(format!("superstep:{step}:{frontier:?}"));
    }

    fn on_node_start(&self, _node_id: async_nodes::NodeId, name: &str) {
        self.record(format!("start:{name}"));
    }

    fn on_node_complete(
        &self,
        _node_id: async_nodes::NodeId,
        name: &str,
        _duration: std::time::Duration,
    ) {
        self.record(format!("done:{name}"));
    }

    fn on_edge_taken(&self, source: async_nodes::NodeId, target: async_nodes::NodeId, label: &str) {
        self.record(format!("edge:{source}->{target}:{label}"));
    }

    fn on_complete(&self, steps: u64) {
        self.record(format!("complete:{steps}"));
    }
}

#[tokio::test]
async fn test_event_stream_sequence() {
    let mut graph = BspGraph::new();
    let a = graph.add_node(AddNode {
        name: "a",
        amount: 1,
    });
    let b = graph.add_node(AddNode {
        name: "b",
        amount: 2,
    });
    graph.add_edge(a.id, b.id);

    let interceptor = Arc::new(RecordingInterceptor {
        events: Mutex::new(Vec::new()),
    });
    let engine = BspEngine::new(graph).with_interceptor(interceptor.clone());
    let outcome = engine.run(0_u64).await.unwrap();
    assert!(matches!(outcome, EngineResult::Complete(_)));

    let events = interceptor.events.lock().unwrap().clone();
    assert_eq!(
        events,
        vec![
            format!("superstep:0:[{}]", a.id),
            "start:a".to_string(),
            "done:a".to_string(),
            format!("edge:{}->{}:always", a.id, b.id),
            format!("superstep:1:[{}]", b.id),
            "start:b".to_string(),
            "done:b".to_string(),
            "complete:2".to_string(),
        ]
    );
}

#[cfg(feature = "tokio")]
#[tokio::test]
async fn test_channel_interceptor_broadcasts_run_events() {
    let channel = ChannelInterceptor::new(64);
    let mut rx = channel.subscribe();

    let mut graph = BspGraph::new();
    let a = graph.add_node(AddNode {
        name: "a",
        amount: 1,
    });
    let b = graph.add_node(AddNode {
        name: "b",
        amount: 2,
    });
    graph.add_edge(a.id, b.id);

    let engine = BspEngine::new(graph).with_interceptor(Arc::new(channel));
    let outcome = engine.run(0_u64).await.unwrap();
    assert!(matches!(outcome, EngineResult::Complete(_)));

    let mut events = Vec::new();
    while let Ok(event) = rx.try_recv() {
        events.push(event);
    }
    assert!(
        events
            .iter()
            .any(|e| matches!(e, BspEvent::SuperstepStart { step: 0, .. }))
    );
    assert!(events.iter().any(|e| matches!(
        e,
        BspEvent::EdgeTaken { source, target, .. } if *source == a.id && *target == b.id
    )));
    assert!(
        events
            .iter()
            .any(|e| matches!(e, BspEvent::Complete { steps: 2 }))
    );
}

// =========================================================================
// EXECUTION COUNTERS & THE SEQUENTIAL API
// =========================================================================

#[tokio::test]
async fn test_execution_no_counts_loop_iterations() {
    let mut graph = BspGraph::new();
    let entry = graph.add_node(NoopNode {
        name: "entry",
        activation: Activation::Any,
    });
    let tracer = graph.add_node(TraceNode { name: "tracer" });
    graph.add_edge(entry.id, tracer.id);
    graph.register_condition("again", |seen: &Vec<u64>, _| seen.len() < 3);
    graph.add_edge_with(
        tracer.id,
        tracer.id,
        async_nodes::Condition::Named("again".into()),
    );

    let engine = BspEngine::new(graph);
    let EngineResult::Complete(state) = engine.run(Vec::new()).await.unwrap() else {
        panic!("expected completion");
    };
    // Fired exactly three times, seeing 0-based execution numbers.
    assert_eq!(state.user, vec![0, 1, 2]);
    assert_eq!(state.execution_count(tracer.id), 3);
}

#[tokio::test]
async fn test_sequential_api_matches_concurrent_for_linear_chains() {
    let build = || {
        let mut graph = BspGraph::new();
        let a = graph.add_node(AddNode {
            name: "a",
            amount: 5,
        });
        let b = graph.add_node(AddNode {
            name: "b",
            amount: 7,
        });
        graph.add_edge(a.id, b.id);
        graph
    };

    let concurrent = BspEngine::new(build()).run(1_u64).await.unwrap();
    let sequential = BspEngine::new(build()).run_sequential(1_u64).await.unwrap();

    let EngineResult::Complete(a) = concurrent else {
        panic!("expected completion")
    };
    let EngineResult::Complete(b) = sequential else {
        panic!("expected completion")
    };
    assert_eq!(a.user, 13);
    assert_eq!(b.user, 13);
}

/// A state type that is deliberately not `Clone` (and still `Send`).
struct NoClone {
    count: u64,
    _marker: Box<dyn Fn() + Send + Sync>,
}

struct CountNode {
    name: &'static str,
}

#[async_trait]
impl BspNode<NoClone> for CountNode {
    fn name(&self) -> &str {
        self.name
    }

    async fn execute(&self, run: &mut RunContext<'_, NoClone>) -> Result<NodeOutcome, BspError> {
        run.state.count += 1;
        Ok(NodeOutcome::Success(IndexMap::new()))
    }
}

#[tokio::test]
async fn test_sequential_api_supports_non_clone_states() {
    let mut graph = BspGraph::new();
    let a = graph.add_node(CountNode { name: "a" });
    let b = graph.add_node(CountNode { name: "b" });
    graph.add_edge(a.id, b.id);

    let engine = BspEngine::new(graph);
    let EngineResult::Complete(state) = engine
        .run_sequential(NoClone {
            count: 0,
            _marker: Box::new(|| {}),
        })
        .await
        .unwrap()
    else {
        panic!("expected completion");
    };
    assert_eq!(state.user.count, 2);
}

// =========================================================================
// FORCED ENTRIES
// =========================================================================

#[tokio::test]
async fn test_run_with_entries_starts_at_the_forced_node() {
    // a -> b -> c; forcing b as the entry must skip a entirely.
    let mut graph = BspGraph::new();
    let a = graph.add_node(AddNode {
        name: "a",
        amount: 1,
    });
    let b = graph.add_node(AddNode {
        name: "b",
        amount: 10,
    });
    let c = graph.add_node(AddNode {
        name: "c",
        amount: 100,
    });
    graph.add_edge(a.id, b.id);
    graph.add_edge(b.id, c.id);

    let engine = BspEngine::new(graph);
    let EngineResult::Complete(state) = engine.run_with_entries(0, [b.id]).await.unwrap() else {
        panic!("expected completion");
    };
    assert_eq!(state.user, 110, "forced entry skips upstream nodes");
    assert_eq!(state.execution_count(a.id), 0);
    assert_eq!(state.execution_count(b.id), 1);
    assert_eq!(state.execution_count(c.id), 1);
}