weirflow 0.1.0

use super::*;
use std::{
    collections::HashSet,
    sync::{
        atomic::{AtomicU64, Ordering},
        Arc, Mutex,
    },
};
use vyre::backend::{private, BackendError, DispatchConfig, OutputBuffers, Resource};

struct FakeResidentBackend {
    next: AtomicU64,
    batch_uploads: AtomicU64,
    uploads: Mutex<Vec<(u64, usize)>>,
    freed: Mutex<Vec<u64>>,
    supported_ops: HashSet<vyre::ir::OpId>,
}

impl FakeResidentBackend {
    fn new() -> Self {
        Self {
            next: AtomicU64::new(1),
            batch_uploads: AtomicU64::new(0),
            uploads: Mutex::new(Vec::new()),
            freed: Mutex::new(Vec::new()),
            supported_ops: HashSet::new(),
        }
    }
}

impl private::Sealed for FakeResidentBackend {}

impl vyre::VyreBackend for FakeResidentBackend {
    fn id(&self) -> &'static str {
        "weir_test_fixed_point_resident_batch"
    }

    fn supported_ops(&self) -> &HashSet<vyre::ir::OpId> {
        &self.supported_ops
    }

    fn dispatch(
        &self,
        _program: &vyre::ir::Program,
        _inputs: &[Vec<u8>],
        _config: &DispatchConfig,
    ) -> Result<Vec<Vec<u8>>, BackendError> {
        unreachable!("resident fixed-point batch tests use persistent resources")
    }

    fn allocate_resident(&self, _byte_len: usize) -> Result<Resource, BackendError> {
        Ok(Resource::Resident(
            self.next.fetch_add(1, Ordering::Relaxed),
        ))
    }

    fn upload_resident(&self, resource: &Resource, bytes: &[u8]) -> Result<(), BackendError> {
        let Resource::Resident(handle) = resource else {
            panic!("fake backend only uploads resident resources");
        };
        self.uploads.lock().unwrap().push((*handle, bytes.len()));
        Ok(())
    }

    fn upload_resident_many(&self, uploads: &[(&Resource, &[u8])]) -> Result<(), BackendError> {
        self.batch_uploads.fetch_add(1, Ordering::Relaxed);
        for &(resource, bytes) in uploads {
            self.upload_resident(resource, bytes)?;
        }
        Ok(())
    }

    fn upload_resident_at_many(
        &self,
        uploads: &[(&Resource, usize, &[u8])],
    ) -> Result<(), BackendError> {
        self.batch_uploads.fetch_add(1, Ordering::Relaxed);
        for &(resource, offset, bytes) in uploads {
            assert_eq!(
                offset, 0,
                "resident fixed-point batch test backend only supports prefix uploads"
            );
            self.upload_resident(resource, bytes)?;
        }
        Ok(())
    }

    fn free_resident(&self, resource: Resource) -> Result<(), BackendError> {
        let Resource::Resident(handle) = resource else {
            panic!("fake backend only frees resident resources");
        };
        self.freed.lock().unwrap().push(handle);
        Ok(())
    }
}

struct FakeResidentPipeline {
    seen_resident_prefix: Mutex<Vec<usize>>,
}

impl FakeResidentPipeline {
    fn new() -> Self {
        Self {
            seen_resident_prefix: Mutex::new(Vec::new()),
        }
    }
}

impl private::Sealed for FakeResidentPipeline {}

impl vyre::CompiledPipeline for FakeResidentPipeline {
    fn id(&self) -> &str {
        "weir_test_fixed_point_resident_batch_pipeline"
    }

    fn dispatch(
        &self,
        _inputs: &[Vec<u8>],
        _config: &DispatchConfig,
    ) -> Result<Vec<Vec<u8>>, BackendError> {
        unreachable!("resident fixed-point batch tests use persistent resources")
    }

    fn dispatch_persistent_handles(
        &self,
        inputs: &[Resource],
        _config: &DispatchConfig,
    ) -> Result<OutputBuffers, BackendError> {
        let resident_prefix = inputs
            .iter()
            .take_while(|resource| matches!(resource, Resource::Resident(_)))
            .count();
        self.seen_resident_prefix
            .lock()
            .unwrap()
            .push(resident_prefix);
        Ok(vec![vec![1, 0, 0, 0]])
    }
}

struct FakeSequenceResidentPipeline {
    seen_resident_prefix: Mutex<Vec<usize>>,
    outputs: Mutex<std::collections::VecDeque<Vec<u8>>>,
}

impl FakeSequenceResidentPipeline {
    fn new(outputs: impl IntoIterator<Item = Vec<u8>>) -> Self {
        Self {
            seen_resident_prefix: Mutex::new(Vec::new()),
            outputs: Mutex::new(outputs.into_iter().collect()),
        }
    }
}

impl private::Sealed for FakeSequenceResidentPipeline {}

impl vyre::CompiledPipeline for FakeSequenceResidentPipeline {
    fn id(&self) -> &str {
        "weir_test_sized_resident_batch_pipeline"
    }

    fn dispatch(
        &self,
        _inputs: &[Vec<u8>],
        _config: &DispatchConfig,
    ) -> Result<Vec<Vec<u8>>, BackendError> {
        unreachable!("resident fixed-point batch tests use persistent resources")
    }

    fn dispatch_persistent_handles(
        &self,
        inputs: &[Resource],
        _config: &DispatchConfig,
    ) -> Result<OutputBuffers, BackendError> {
        let resident_prefix = inputs
            .iter()
            .take_while(|resource| matches!(resource, Resource::Resident(_)))
            .count();
        self.seen_resident_prefix
            .lock()
            .unwrap()
            .push(resident_prefix);
        let output = self
            .outputs
            .lock()
            .unwrap()
            .pop_front()
            .expect("fake sequence resident pipeline ran out of outputs");
        Ok(vec![output])
    }
}

fn graph_fixture() -> FixedPointForwardGraph {
    graph_fixture_for_kind(crate::fixed_point_graph::FixedPointAnalysisKind::Reaching)
}

fn graph_fixture_for_kind(
    kind: crate::fixed_point_graph::FixedPointAnalysisKind,
) -> FixedPointForwardGraph {
    FixedPointForwardGraph::new_for_kind(
        kind,
        "fixed_point_resident_batch_test",
        2,
        &[0, 1, 1],
        &[1],
        &[vyre_primitives::predicate::edge_kind::CONTROL],
    )
    .expect("valid graph must pack")
}

fn empty_graph_fixture(node_count: u32) -> FixedPointForwardGraph {
    empty_graph_fixture_for_kind(
        crate::fixed_point_graph::FixedPointAnalysisKind::Reaching,
        node_count,
    )
}

fn empty_graph_fixture_for_kind(
    kind: crate::fixed_point_graph::FixedPointAnalysisKind,
    node_count: u32,
) -> FixedPointForwardGraph {
    FixedPointForwardGraph::new_for_kind(
        kind,
        "fixed_point_resident_batch_empty_test",
        node_count,
        &vec![0; node_count as usize + 1],
        &[],
        &[],
    )
    .expect("valid empty graph must pack")
}

mod accounting_contracts;
mod cache_reuse_contracts;
mod result_buffer_contracts;
mod validation_contracts;