weirflow 0.1.0

//! Tests from `tests.rs`.

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

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

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

    fn next_handle(&self) -> u64 {
        self.next.load(Ordering::Relaxed)
    }
}

impl private::Sealed for FakeResidentBackend {}

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

    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 graph tests do not dispatch through host bytes")
    }

    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> {
        for &(resource, bytes) in uploads {
            self.upload_resident(resource, bytes)?;
        }
        Ok(())
    }

    fn upload_resident_at_many(
        &self,
        uploads: &[(&Resource, usize, &[u8])],
    ) -> Result<(), BackendError> {
        for &(resource, offset, bytes) in uploads {
            assert_eq!(
                offset, 0,
                "fixed-point test backend only supports prefix resident 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(())
    }
}

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

impl FakeResidentPipeline {
    pub 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_fake_resident_pipeline"
    }

    fn dispatch(
        &self,
        _inputs: &[Vec<u8>],
        _config: &DispatchConfig,
    ) -> Result<Vec<Vec<u8>>, BackendError> {
        unreachable!("resident pipeline tests do not dispatch through host bytes")
    }

    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 frontier = match &inputs[5] {
            Resource::Borrowed(frontier) => frontier.clone(),
            Resource::Resident(_) => vec![1, 0, 0, 0],
        };
        Ok(vec![frontier])
    }
}

mod fixed_point_resident_graph_uploads_invariant_buffers_once;
mod fixed_point_resident_plan_pairs_pipeline_with_graph_resources;
mod fixed_point_resident_plan_reuses_frontier_resources_between_runs;
mod fixed_point_resident_plan_reuses_live_frontier_resources_between_runs;
mod fixed_point_resident_plan_reuses_points_to_frontier_resources_between_runs;
mod fixed_point_resident_plan_reuses_slice_frontier_resources_between_runs;
mod live_resident_plan_uses_resident_graph_resources;
mod points_to_resident_plan_uses_resident_graph_resources;
mod reaching_resident_plan_uses_resident_graph_resources;
mod resident_analysis_rejects_wrong_graph_family;
mod resident_graph_resource_builders_reuse_caller_resource_vector;
mod resident_resource_staging_production_path_returns_errors_instead_of_panicking;
mod slice_resident_plan_uses_resident_graph_resources;
mod zero_copy_resident_staging;