async-nodes 0.1.0

//! Fork/join branch-isolation tests: merger folds in frontier declaration order (never
//! completion order), and last-writer-wins with a validation warning when no merger is
//! named.

use std::sync::{Arc, Mutex};

use async_nodes::{
    Activation, BspEngine, BspError, BspGraph, BspNode, EngineResult, Merger, NodeOutcome,
    RunContext, async_trait::async_trait,
};
use indexmap::IndexMap;

type Bag = Vec<&'static str>;

/// Pushes a tag onto the bag; `yields` controls how long the branch takes so completion
/// order can be made to disagree with declaration order.
struct PushNode {
    name: &'static str,
    tag: &'static str,
    yields: usize,
}

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

    async fn execute(&self, run: &mut RunContext<'_, Bag>) -> Result<NodeOutcome, BspError> {
        for _ in 0..self.yields {
            tokio::task::yield_now().await;
        }
        run.state.push(self.tag);
        Ok(NodeOutcome::Success(IndexMap::new()))
    }
}

/// The join node; optionally names a merger.
struct JoinNode {
    merger: Option<&'static str>,
}

#[async_trait]
impl BspNode<Bag> for JoinNode {
    fn name(&self) -> &str {
        "join"
    }

    fn activation(&self) -> Activation {
        Activation::All
    }

    fn merger(&self) -> Option<&str> {
        self.merger
    }

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

/// Builds src -> {a, b} -> join. Branch `a` is slow (many yields), branch `b` is fast, so
/// task *completion* order is b-then-a while *declaration* order is a-then-b.
fn fork_join_graph(with_merger: bool) -> BspGraph<Bag> {
    let mut graph = BspGraph::new();
    let src = graph.add_node(PushNode {
        name: "src",
        tag: "start",
        yields: 0,
    });
    let a = graph.add_node(PushNode {
        name: "a",
        tag: "a",
        yields: 12,
    });
    let b = graph.add_node(PushNode {
        name: "b",
        tag: "b",
        yields: 0,
    });
    let join = graph.add_node(JoinNode {
        merger: with_merger.then_some("append-missing"),
    });
    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);
    graph.register_merger("append-missing", |trunk: &mut Bag, branch: Bag| {
        for item in branch {
            if !trunk.contains(&item) {
                trunk.push(item);
            }
        }
    });
    graph
}

#[tokio::test]
async fn test_merge_folds_in_declaration_order_deterministically() {
    for _ in 0..3 {
        let engine = BspEngine::new(fork_join_graph(true));
        let EngineResult::Complete(state) = engine.run(Vec::new()).await.unwrap() else {
            panic!("expected completion");
        };
        // Branch `b` completes first but `a` was declared first: the fold follows
        // declaration order, so the merge result is stable across runs.
        assert_eq!(state.user, vec!["start", "a", "b", "joined"]);
    }
}

#[tokio::test]
async fn test_join_without_merger_is_last_writer_wins_with_warning() {
    let engine = BspEngine::new(fork_join_graph(false));

    let report = engine.validate_with_report();
    assert!(
        report
            .warnings
            .iter()
            .any(|w| w.contains("last-writer-wins")),
        "expected a last-writer-wins warning, got {:?}",
        report.warnings
    );

    let EngineResult::Complete(state) = engine.run(Vec::new()).await.unwrap() else {
        panic!("expected completion");
    };
    // No merger: the last branch in declaration order wins the fold — deterministically,
    // regardless of `b` completing first.
    assert_eq!(state.user, vec!["start", "b", "joined"]);
}

// =========================================================================
// Merger::finish & BspEngine::with_branch_transform
// =========================================================================

/// An all-at-once merger: `merge` accumulates branch states into a side buffer, and
/// `finish` collapses the whole buffer into the trunk in one shot (the shape the
/// reference consumer's reducer barriers need). Records every branch it saw so tests
/// can inspect the fold inputs.
struct CollectMerger {
    seen: Mutex<Vec<Bag>>,
    finish_calls: Mutex<usize>,
    fail_on_finish: bool,
}

#[async_trait]
impl Merger<Bag> for CollectMerger {
    fn merge(&self, _trunk: &mut Bag, branch: Bag) {
        self.seen.lock().unwrap().push(branch);
    }

    async fn finish(&self, trunk: &mut Bag) -> Result<(), BspError> {
        *self.finish_calls.lock().unwrap() += 1;
        if self.fail_on_finish {
            return Err(BspError::ValidationError("finish blew up".to_string()));
        }
        let seen = std::mem::take(&mut *self.seen.lock().unwrap());
        // All-at-once collapse: keep the trunk's own items, then one item per branch.
        let mut collapsed: Bag = trunk.clone();
        for (i, _branch) in seen.iter().enumerate() {
            collapsed.push(match i {
                0 => "first",
                1 => "second",
                _ => "rest",
            });
        }
        *trunk = collapsed;
        Ok(())
    }
}

fn three_branch_graph<M: Merger<Bag> + 'static>(merger: M) -> BspGraph<Bag> {
    let mut graph = BspGraph::new();
    let src = graph.add_node(PushNode {
        name: "src",
        tag: "start",
        yields: 0,
    });
    let a = graph.add_node(PushNode {
        name: "a",
        tag: "a",
        yields: 12,
    });
    let b = graph.add_node(PushNode {
        name: "b",
        tag: "b",
        yields: 4,
    });
    let c = graph.add_node(PushNode {
        name: "c",
        tag: "c",
        yields: 0,
    });
    let join = graph.add_node(JoinNode {
        merger: Some("collect"),
    });
    graph.add_edge(src.id, a.id);
    graph.add_edge(src.id, b.id);
    graph.add_edge(src.id, c.id);
    graph.add_edge(a.id, join.id);
    graph.add_edge(b.id, join.id);
    graph.add_edge(c.id, join.id);
    graph.register_merger("collect", merger);
    graph
}

#[tokio::test]
async fn test_finish_collapses_all_at_once_and_runs_exactly_once() {
    let merger = CollectMerger {
        seen: Mutex::new(Vec::new()),
        finish_calls: Mutex::new(0),
        fail_on_finish: false,
    };
    let seen_handle: Arc<Mutex<Vec<Bag>>> = Arc::new(Mutex::new(Vec::new()));
    let calls_handle = Arc::new(Mutex::new(0_usize));
    // Wrap to observe from the outside after the engine consumed the graph.
    struct Observer {
        inner: CollectMerger,
        seen: Arc<Mutex<Vec<Bag>>>,
        calls: Arc<Mutex<usize>>,
    }
    #[async_trait]
    impl Merger<Bag> for Observer {
        fn merge(&self, trunk: &mut Bag, branch: Bag) {
            self.seen.lock().unwrap().push(branch.clone());
            self.inner.merge(trunk, branch);
        }
        async fn finish(&self, trunk: &mut Bag) -> Result<(), BspError> {
            *self.calls.lock().unwrap() += 1;
            self.inner.finish(trunk).await
        }
    }
    let observer = Observer {
        inner: merger,
        seen: seen_handle.clone(),
        calls: calls_handle.clone(),
    };

    let engine = BspEngine::new(three_branch_graph(observer));
    let EngineResult::Complete(state) = engine.run(Vec::new()).await.unwrap() else {
        panic!("expected completion");
    };

    // One barrier fold for the single 3-branch superstep: finish ran exactly once,
    // after all three branches folded.
    assert_eq!(*calls_handle.lock().unwrap(), 1);
    assert_eq!(seen_handle.lock().unwrap().len(), 3);
    // finish replaced the trunk with the collapsed view: the pre-join trunk items
    // (just "start", since merge accumulated rather than appended) + one per branch,
    // then the join node itself ran on the collapsed trunk.
    assert_eq!(
        state.user,
        vec!["start", "first", "second", "rest", "joined"]
    );
}

#[tokio::test]
async fn test_finish_error_aborts_the_run() {
    let engine = BspEngine::new(three_branch_graph(CollectMerger {
        seen: Mutex::new(Vec::new()),
        finish_calls: Mutex::new(0),
        fail_on_finish: true,
    }));
    let result = engine.run(Vec::new()).await;
    assert!(
        matches!(result, Err(BspError::ValidationError(ref m)) if m == "finish blew up"),
        "finish failure must abort the run loudly, got {result:?}"
    );
}

#[tokio::test]
async fn test_closure_merger_finish_defaults_to_noop() {
    let merger = |trunk: &mut Bag, mut branch: Bag| trunk.append(&mut branch);
    let mut trunk = vec!["a"];
    Merger::finish(&merger, &mut trunk).await.unwrap();
    assert_eq!(trunk, vec!["a"], "default finish must not touch the trunk");
}

#[tokio::test]
async fn test_branch_transform_applies_per_frontier_ordinal() {
    // Records what each branch clone looked like at fold time.
    let seen: Arc<Mutex<Vec<Bag>>> = Arc::new(Mutex::new(Vec::new()));
    let seen_clone = seen.clone();
    let graph = three_branch_graph(move |_trunk: &mut Bag, branch: Bag| {
        seen_clone.lock().unwrap().push(branch);
    });

    // Zero every sibling clone after the first (llm-net's fork metrics rule, ported).
    let engine = BspEngine::new(graph).with_branch_transform(Arc::new(
        |branch: &mut Bag, ordinal: usize| {
            if ordinal > 0 {
                branch.clear();
            }
        },
    ));
    let EngineResult::Complete(_) = engine.run(Vec::new()).await.unwrap() else {
        panic!("expected completion");
    };

    let seen = seen.lock().unwrap();
    assert_eq!(seen.len(), 3);
    // Branch 0 keeps the inherited trunk; siblings started from a cleared clone.
    assert!(
        seen[0].contains(&"start"),
        "first branch keeps the trunk clone"
    );
    for branch in &seen[1..] {
        assert!(
            !branch.contains(&"start"),
            "sibling branch must not inherit the trunk: {branch:?}"
        );
    }
}

#[tokio::test]
async fn test_branch_transform_default_is_identity() {
    let engine = BspEngine::new(fork_join_graph(true));
    let EngineResult::Complete(state) = engine.run(Vec::new()).await.unwrap() else {
        panic!("expected completion");
    };
    // Identical to the no-transform run: every branch was an exact clone.
    assert_eq!(state.user, vec!["start", "a", "b", "joined"]);
}