tenferro-runtime 0.9.0

use std::collections::BTreeMap;
use std::num::NonZeroU64;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};

use super::super::execution::ScheduledEventDomains;
use super::super::schedule::{
    EventCompletion, EventDependency, EventDomainId, EventSlotId, ExecutionLocation,
    ScheduledCollective, ScheduledGraph, ScheduledNode, ScheduledNodeKind, ScheduledOperation,
    ScheduledTransfer, UnsupportedScheduledNodeError,
};
use super::super::{
    EngineId, EventDomainDriver, EventDomainRun, EventToken, FrozenTransferRegistry,
    ImmediateEventDomainDriver, ProviderDeviceIdentity, ProviderId, RegistrationIdentity,
    ResolvedTransferRoute, RuntimeEpoch, RuntimeId, StorageClass, TransferProvider,
    TransferRequest,
};
use crate::exec::{ExecInstruction, ExecOp, ExecProgram};
use crate::DType;

fn qualified_domain(
    runtime_value: u64,
    epoch_value: u64,
    issuer_value: u64,
    ordinal_value: u64,
) -> EventDomainId {
    EventDomainId::runtime_created_for_test(
        RuntimeId::from_nonzero(NonZeroU64::new(runtime_value).expect("runtime id")),
        RuntimeEpoch::from_nonzero(NonZeroU64::new(epoch_value).expect("runtime epoch")),
        RegistrationIdentity::new(
            NonZeroU64::new(issuer_value).expect("registration issuer"),
            NonZeroU64::new(ordinal_value).expect("registration ordinal"),
        ),
    )
}

fn location(engine: &str, domain: EventDomainId, storage: &str) -> ExecutionLocation {
    ExecutionLocation::new(
        EngineId::new(engine).expect("engine id"),
        ProviderDeviceIdentity::new(
            ProviderId::new("tenferro.test.schedule").expect("provider id"),
            format!("target:{engine}"),
        )
        .expect("provider target"),
        domain,
        StorageClass::new(storage).expect("storage class"),
    )
}

#[derive(Debug)]
struct TestTransferProvider;

impl TransferProvider for TestTransferProvider {
    fn transfer_blocking(
        &self,
        _request: TransferRequest<'_>,
    ) -> crate::Result<tenferro_tensor::Tensor> {
        Err(crate::Error::Internal("schedule test transfer".into()))
    }
}

fn test_identity(target: &str) -> ProviderDeviceIdentity {
    ProviderDeviceIdentity::new(
        ProviderId::new("tenferro.test.schedule").expect("provider id"),
        target,
    )
    .expect("provider target")
}

fn registry(routes: &[(&ExecutionLocation, &ExecutionLocation)]) -> FrozenTransferRegistry {
    let provider: Arc<dyn TransferProvider> = Arc::new(TestTransferProvider);
    FrozenTransferRegistry::new(
        routes
            .iter()
            .map(|(source, destination)| {
                (
                    ResolvedTransferRoute::new(
                        source.resolved_endpoint().clone(),
                        destination.resolved_endpoint().clone(),
                    ),
                    Arc::clone(&provider),
                )
            })
            .collect::<BTreeMap<_, _>>(),
    )
}

#[derive(Debug)]
struct BeginRunProbe {
    begin_runs: Arc<AtomicUsize>,
}

impl EventDomainDriver for BeginRunProbe {
    fn begin_run(&self, domain: EventDomainId) -> crate::Result<Box<dyn EventDomainRun>> {
        self.begin_runs.fetch_add(1, Ordering::SeqCst);
        ImmediateEventDomainDriver::new().begin_run(domain)
    }
}

fn instruction(
    semantic_operation_index: usize,
    input_slot: usize,
    output_slot: usize,
    last_use: bool,
) -> ExecInstruction {
    ExecInstruction {
        op: ExecOp::Negate,
        semantic_operation_index: Some(semantic_operation_index),
        input_slots: vec![input_slot],
        output_slots: vec![output_slot],
        dtype: DType::F64,
        output_shapes: Default::default(),
        output_extents: Default::default(),
        last_use: vec![last_use],
        input_views: Vec::new(),
    }
}

#[test]
fn schedule_emits_location_transfer_and_retains_source_for_split_use() {
    let source = location(
        "tenferro-test.engine.source",
        qualified_domain(1, 1, 1, 1),
        "tenferro-test.storage.shared",
    );
    let destination = location(
        "tenferro-test.engine.destination",
        qualified_domain(1, 1, 1, 2),
        "tenferro-test.storage.shared",
    );
    let program = ExecProgram {
        instructions: vec![
            instruction(0, 0, 1, true),
            instruction(1, 1, 2, false),
            instruction(2, 1, 3, true),
        ],
        input_slots: vec![0],
        output_slots: vec![2, 3],
        n_slots: 4,
        shape_guards: Vec::new(),
    };

    let graph = ScheduledGraph::from_exec_program(
        &program,
        source.clone(),
        std::slice::from_ref(&source),
        &[source.clone(), destination.clone(), source.clone()],
        &registry(&[(&source, &destination)]),
    )
    .expect("schedule");

    assert_eq!(graph.nodes_for_test().len(), 4);
    assert!(matches!(
        &graph.nodes_for_test()[0],
        ScheduledNode::Operation(operation)
            if operation.instruction_index() == 0 && operation.location() == &source
    ));
    assert!(matches!(
        &graph.nodes_for_test()[1],
        ScheduledNode::Transfer(transfer)
            if transfer.value_slot() == 1
                && transfer.source_location() == &source
                && transfer.destination_location() == &destination
    ));
    assert!(matches!(
        &graph.nodes_for_test()[2],
        ScheduledNode::Operation(operation)
            if operation.instruction_index() == 1
                && operation.location() == &destination
                && operation.dependencies()
                    == [EventDependency::new(
                        destination.event_domain_id(),
                        EventSlotId::new(1),
                        0,
                    )]
    ));
    assert!(matches!(
        &graph.nodes_for_test()[3],
        ScheduledNode::Operation(operation)
            if operation.instruction_index() == 2 && operation.location() == &source
    ));
}

#[test]
fn schedule_uses_a_copy_with_a_direct_route_to_the_destination() {
    let first = location(
        "tenferro-test.engine.first",
        qualified_domain(1, 1, 1, 1),
        "tenferro-test.storage.first",
    );
    let reachable = location(
        "tenferro-test.engine.reachable",
        qualified_domain(1, 1, 1, 2),
        "tenferro-test.storage.reachable",
    );
    let destination = location(
        "tenferro-test.engine.destination",
        qualified_domain(1, 1, 1, 3),
        "tenferro-test.storage.destination",
    );
    let program = ExecProgram {
        instructions: vec![
            instruction(0, 0, 1, true),
            instruction(1, 1, 2, false),
            instruction(2, 1, 3, true),
        ],
        input_slots: vec![0],
        output_slots: vec![2, 3],
        n_slots: 4,
        shape_guards: Vec::new(),
    };

    let graph = ScheduledGraph::from_exec_program(
        &program,
        first.clone(),
        std::slice::from_ref(&first),
        &[first.clone(), reachable.clone(), destination.clone()],
        &registry(&[(&first, &reachable), (&reachable, &destination)]),
    )
    .expect("schedule");
    let transfers = graph.transfers_for_test().collect::<Vec<_>>();

    assert_eq!(transfers.len(), 2);
    assert_eq!(transfers[1].source_location(), &reachable);
    assert_eq!(transfers[1].destination_location(), &destination);
}

#[test]
fn operation_dependencies_follow_input_order_and_are_deduplicated() {
    let execution = location(
        "tenferro-test.engine.execution",
        qualified_domain(1, 1, 1, 1),
        "tenferro-test.storage.execution",
    );
    let program = ExecProgram {
        instructions: vec![
            instruction(0, 0, 1, false),
            instruction(1, 0, 2, true),
            ExecInstruction {
                op: ExecOp::Negate,
                semantic_operation_index: Some(2),
                input_slots: vec![2, 1, 2],
                output_slots: vec![3],
                dtype: DType::F64,
                output_shapes: Default::default(),
                output_extents: Default::default(),
                last_use: vec![true, true, true],
                input_views: Vec::new(),
            },
        ],
        input_slots: vec![0],
        output_slots: vec![3],
        n_slots: 4,
        shape_guards: Vec::new(),
    };

    let graph = ScheduledGraph::from_exec_program(
        &program,
        execution.clone(),
        std::slice::from_ref(&execution),
        &[execution.clone(), execution.clone(), execution.clone()],
        &FrozenTransferRegistry::new(BTreeMap::new()),
    )
    .expect("schedule");
    let ScheduledNode::Operation(operation) = &graph.nodes_for_test()[2] else {
        panic!("third node should be an operation");
    };

    assert_eq!(
        operation.dependencies(),
        &[
            EventDependency::new(qualified_domain(1, 1, 1, 1), EventSlotId::new(1), 0,),
            EventDependency::new(qualified_domain(1, 1, 1, 1), EventSlotId::new(0), 0,),
        ]
    );
}

#[test]
fn schedule_validation_rejects_dependency_without_prior_completion() {
    let source = qualified_domain(1, 1, 1, 1);
    let destination = qualified_domain(1, 1, 1, 2);
    let source_location = location(
        "tenferro-test.engine.source",
        source,
        "tenferro-test.storage.source",
    );
    let destination_location = location(
        "tenferro-test.engine.destination",
        destination,
        "tenferro-test.storage.destination",
    );
    let graph = ScheduledGraph::for_test(vec![ScheduledNode::Transfer(ScheduledTransfer::new(
        0,
        source_location,
        destination_location,
        [EventDependency::new(source, EventSlotId::new(99), 0)],
        super::super::schedule::EventCompletion::new(destination, EventSlotId::new(0), 0),
    ))]);

    assert!(graph.validate().is_err());
}

#[test]
fn schedule_validation_rejects_duplicate_completion_identity() {
    let domain = qualified_domain(1, 1, 1, 1);
    let graph = ScheduledGraph::for_test(vec![
        ScheduledNode::Operation(ScheduledOperation::for_test(
            domain,
            test_identity("duplicate-first"),
        )),
        ScheduledNode::Operation(ScheduledOperation::for_test(
            domain,
            test_identity("duplicate-second"),
        )),
    ]);

    assert!(matches!(
        graph.validate(),
        Err(super::super::schedule::ScheduleValidationError::DuplicateCompletion { index: 1 })
    ));
}

#[test]
fn retained_bytes_include_operation_dependencies() {
    let execution = location(
        "tenferro-test.engine.execution",
        qualified_domain(1, 1, 1, 1),
        "tenferro-test.storage.execution",
    );
    let program = ExecProgram {
        instructions: vec![instruction(0, 0, 1, true), instruction(1, 1, 2, true)],
        input_slots: vec![0],
        output_slots: vec![2],
        n_slots: 3,
        shape_guards: Vec::new(),
    };
    let graph = ScheduledGraph::from_exec_program(
        &program,
        execution.clone(),
        std::slice::from_ref(&execution),
        &[execution.clone(), execution.clone()],
        &FrozenTransferRegistry::new(BTreeMap::new()),
    )
    .expect("schedule");
    let expected = std::mem::size_of::<ScheduledGraph>()
        + 2 * std::mem::size_of::<ScheduledNode>()
        + 4 * std::mem::size_of::<usize>()
        + std::mem::size_of::<EventDependency>()
        + std::mem::size_of::<usize>()
        + std::mem::size_of::<usize>()
        + 3 * std::mem::size_of::<usize>()
        + 3 * std::mem::size_of::<ExecutionLocation>();

    assert_eq!(graph.retained_bytes(), Some(expected));
}

#[test]
fn transfer_node_bridges_distinct_event_domains() {
    let source = qualified_domain(1, 1, 1, 1);
    let destination = qualified_domain(1, 1, 1, 2);
    let transfer = ScheduledTransfer::for_test(
        source,
        destination,
        test_identity("transfer-source"),
        test_identity("transfer-destination"),
    );

    assert_ne!(
        transfer.source_event_domain(),
        transfer.destination_event_domain()
    );
    assert_eq!(transfer.dependencies()[0].domain(), source);
    assert_eq!(transfer.completion().domain(), destination);
}

#[test]
fn collective_node_is_rejected_by_production_preflight_before_begin_run() {
    let begin_runs = Arc::new(AtomicUsize::new(0));
    let driver: Arc<dyn EventDomainDriver> = Arc::new(BeginRunProbe {
        begin_runs: Arc::clone(&begin_runs),
    });
    let graph = ScheduledGraph::for_test(vec![ScheduledNode::Collective(
        ScheduledCollective::unsupported_for_test(),
    )]);

    assert!(graph.contains_collective());
    assert!(graph.validate().is_ok());
    let engine_id = EngineId::new("tenferro-test.schedule-engine").expect("test engine id");
    let error = ScheduledEventDomains::new_for_test(&graph, &[(engine_id, driver)])
        .expect_err("unsupported collective must fail during production preflight");
    let crate::Error::RuntimeStateSource { source, .. } = error else {
        panic!("unsupported collective must be wrapped as RuntimeStateSource");
    };
    let unsupported = source
        .downcast_ref::<UnsupportedScheduledNodeError>()
        .expect("typed unsupported scheduled node source");
    assert_eq!(unsupported.node_index, 0);
    assert_eq!(unsupported.node_kind, ScheduledNodeKind::Collective);
    assert_eq!(begin_runs.load(Ordering::SeqCst), 0);
}

#[test]
fn mock_transfer_does_not_reuse_source_event_domain_as_destination_completion() {
    let source = qualified_domain(1, 1, 1, 7);
    let destination = qualified_domain(1, 1, 1, 8);
    let graph = ScheduledGraph::for_test(vec![
        ScheduledNode::Operation(ScheduledOperation::for_test(
            source,
            test_identity("mock-source"),
        )),
        ScheduledNode::Transfer(ScheduledTransfer::for_test(
            source,
            destination,
            test_identity("mock-source"),
            test_identity("mock-destination"),
        )),
        ScheduledNode::Operation(ScheduledOperation::new(
            1,
            location(
                "tenferro-test.engine.destination",
                destination,
                "tenferro-test.storage.destination",
            ),
            [],
            [],
            [EventDependency::new(destination, EventSlotId::new(0), 0)],
            super::super::schedule::EventCompletion::new(destination, EventSlotId::new(1), 0),
        )),
    ]);

    let transfer = graph
        .transfers_for_test()
        .next()
        .expect("scheduled transfer");

    assert_eq!(transfer.dependencies()[0].domain(), source);
    assert_eq!(transfer.completion().domain(), destination);
    assert_ne!(
        graph.nodes_for_test()[0].completion().domain(),
        transfer.completion().domain()
    );
    assert!(graph.validate().is_ok());
}

/// Driver whose run records every single-domain coalescing request.
#[derive(Debug)]
struct CoalescingProbe {
    enabled: Arc<AtomicUsize>,
    reports: bool,
    enqueued: Arc<Mutex<Vec<usize>>>,
}

impl CoalescingProbe {
    fn reporting(enabled: &Arc<AtomicUsize>, reports: bool) -> Self {
        Self {
            enabled: Arc::clone(enabled),
            reports,
            enqueued: Arc::new(Mutex::new(Vec::new())),
        }
    }
}

impl EventDomainDriver for CoalescingProbe {
    fn begin_run(&self, domain: EventDomainId) -> crate::Result<Box<dyn EventDomainRun>> {
        Ok(Box::new(CoalescingProbeRun {
            domain,
            enabled: Arc::clone(&self.enabled),
            reports: self.reports,
            enqueued: Arc::clone(&self.enqueued),
            accepted: false,
            inner: ImmediateEventDomainDriver::new().begin_run(domain)?,
        }))
    }
}

#[derive(Debug)]
struct CoalescingProbeRun {
    domain: EventDomainId,
    enabled: Arc<AtomicUsize>,
    /// Whether this run accepts submission order in place of per-node
    /// completions; the probe records the offer either way.
    reports: bool,
    /// Dependency counts of every enqueue, in submission order.
    enqueued: Arc<Mutex<Vec<usize>>>,
    accepted: bool,
    inner: Box<dyn EventDomainRun>,
}

impl EventDomainRun for CoalescingProbeRun {
    fn domain(&self) -> EventDomainId {
        self.domain
    }

    fn enable_single_domain_coalescing(&mut self) -> bool {
        self.enabled.fetch_add(1, Ordering::SeqCst);
        self.accepted = true;
        self.accepted && self.reports
    }

    fn enqueue(
        &mut self,
        dependencies: &[Arc<dyn EventToken>],
        launch: &mut dyn FnMut() -> crate::Result<()>,
    ) -> crate::Result<Arc<dyn EventToken>> {
        self.enqueued
            .lock()
            .expect("probe enqueue log")
            .push(dependencies.len());
        self.inner.enqueue(dependencies, launch)
    }

    fn drain(&mut self) -> crate::Result<()> {
        self.inner.drain()
    }
}

fn coalescing_engine_id() -> EngineId {
    EngineId::new("tenferro-test.coalescing-engine").expect("test engine id")
}

fn probe_operation(
    instruction_index: usize,
    domain: EventDomainId,
    dependencies: impl Into<Box<[EventDependency]>>,
    slot: u32,
) -> ScheduledNode {
    ScheduledNode::Operation(ScheduledOperation::new(
        instruction_index,
        ExecutionLocation::new(
            coalescing_engine_id(),
            test_identity("coalescing-probe"),
            domain,
            StorageClass::new("tenferro-test.coalescing-storage").expect("test storage class"),
        ),
        [],
        [],
        dependencies,
        EventCompletion::new(domain, EventSlotId::new(slot), u64::from(slot)),
    ))
}

#[test]
fn single_domain_schedule_requests_coalescing_once() {
    let enabled = Arc::new(AtomicUsize::new(0));
    let driver: Arc<dyn EventDomainDriver> = Arc::new(CoalescingProbe::reporting(&enabled, false));
    let domain = qualified_domain(1, 1, 1, 1);
    let graph = ScheduledGraph::for_test(vec![
        probe_operation(0, domain, [], 0),
        probe_operation(
            1,
            domain,
            [EventDependency::new(domain, EventSlotId::new(0), 0)],
            1,
        ),
    ]);
    assert!(graph.validate().is_ok());

    let _domains = ScheduledEventDomains::new_for_test(&graph, &[(coalescing_engine_id(), driver)])
        .expect("single-domain runs");
    assert_eq!(
        enabled.load(Ordering::SeqCst),
        1,
        "one run must be asked to accept the coalescing mode"
    );
}

#[test]
fn multi_domain_schedule_never_requests_coalescing() {
    let enabled = Arc::new(AtomicUsize::new(0));
    let driver: Arc<dyn EventDomainDriver> = Arc::new(CoalescingProbe::reporting(&enabled, false));
    let first_domain = qualified_domain(1, 1, 1, 1);
    let second_domain = qualified_domain(1, 1, 1, 2);
    let graph = ScheduledGraph::for_test(vec![
        probe_operation(0, first_domain, [], 0),
        probe_operation(1, second_domain, [], 1),
    ]);
    assert!(graph.validate().is_ok());

    let _domains = ScheduledEventDomains::new_for_test(&graph, &[(coalescing_engine_id(), driver)])
        .expect("multi-domain runs");
    assert_eq!(
        enabled.load(Ordering::SeqCst),
        0,
        "a multi-domain schedule must keep the per-node completion path"
    );
}

/// A run that accepts the coalescing offer receives no dependency tokens and
/// records no completion, so the scheduler never resolves a dependency for it.
#[test]
fn accepting_run_submits_dependencies_free_commands() {
    let enabled = Arc::new(AtomicUsize::new(0));
    let probe = Arc::new(CoalescingProbe::reporting(&enabled, true));
    let driver: Arc<dyn EventDomainDriver> = probe.clone();
    let domain = qualified_domain(1, 1, 1, 1);
    let graph = ScheduledGraph::for_test(vec![
        probe_operation(0, domain, [], 0),
        probe_operation(
            1,
            domain,
            [EventDependency::new(domain, EventSlotId::new(0), 0)],
            1,
        ),
    ]);
    assert!(graph.validate().is_ok());

    let mut domains =
        ScheduledEventDomains::new_for_test(&graph, &[(coalescing_engine_id(), driver)])
            .expect("single-domain runs");
    let mut launches = 0;
    for node in graph.nodes_for_test() {
        domains
            .enqueue(0, node, &mut || {
                launches += 1;
                Ok(())
            })
            .expect("coalesced enqueue");
    }
    assert_eq!(launches, 2, "every node still launches");
    assert!(
        domains.completion_count_for_test() == 0,
        "a reporting run must not make the scheduler record completions"
    );
    assert_eq!(
        probe.enqueued.lock().expect("probe enqueue log").as_slice(),
        &[0, 0],
        "a coalesced command carries no dependency tokens"
    );
}

/// A run that declines the offer keeps the per-node completion contract:
/// dependencies are resolved and passed to the run.
#[test]
fn declining_run_keeps_per_node_completions() {
    let enabled = Arc::new(AtomicUsize::new(0));
    let probe = Arc::new(CoalescingProbe::reporting(&enabled, false));
    let driver: Arc<dyn EventDomainDriver> = probe.clone();
    let domain = qualified_domain(1, 1, 1, 1);
    let graph = ScheduledGraph::for_test(vec![
        probe_operation(0, domain, [], 0),
        probe_operation(
            1,
            domain,
            [EventDependency::new(domain, EventSlotId::new(0), 0)],
            1,
        ),
    ]);
    assert!(graph.validate().is_ok());

    let mut domains =
        ScheduledEventDomains::new_for_test(&graph, &[(coalescing_engine_id(), driver)])
            .expect("single-domain runs");
    for node in graph.nodes_for_test() {
        domains
            .enqueue(0, node, &mut || Ok(()))
            .expect("enqueue with completions");
    }
    assert_eq!(
        domains.completion_count_for_test(),
        2,
        "a non-reporting run keeps one completion per node"
    );
    assert_eq!(
        probe.enqueued.lock().expect("probe enqueue log").as_slice(),
        &[0, 1],
        "the second command resolves its dependency"
    );
}