monolith

use crate::Error;

use crate::device::Shared;

use ash::vk;

use gpu_allocator::MemoryLocation;

use gpu_allocator::vulkan::{Allocation, AllocationCreateDesc, AllocationScheme};

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

use std::time::{Duration, Instant};

/// Pooled staging memory kept at most, past which dropped buffers are freed.

const POOL_BYTES: usize = 32 * 1024 * 1024;

/// How long a pooled buffer may go unused before it is freed.

const POOL_IDLE: Duration = Duration::from_secs(2);

struct Pooled {

    buffer: vk::Buffer,

    allocation: Allocation,

    size: usize,

    usage: vk::BufferUsageFlags,

    location: MemoryLocation,

    returned: Instant,

}

/// Staging buffers kept for reuse, see [`Staging`].

pub(crate) struct StagingPool {

    shared: Arc<Shared>,

    free: Mutex<Vec<Pooled>>,

    /// Granularity of flushes and invalidations of non-coherent memory.

    atom: u64,

}

impl StagingPool {

    pub fn new(shared: &Arc<Shared>) -> Arc<Self> {

        let limits = unsafe {

            shared

                .device

                .instance

                .get_physical_device_properties(shared.device.physical)

        }

        .limits;

        Arc::new(Self {

            shared: shared.clone(),

            free: Mutex::new(Vec::new()),

            atom: limits.non_coherent_atom_size.max(1),

        })

    }

    /// Frees pooled buffers that have sat unused for [`POOL_IDLE`].

    pub fn trim(&self) {

        let now = Instant::now();

        let expired: Vec<Pooled> = self

            .free

            .lock()

            .unwrap_or_else(|e| e.into_inner())

            .extract_if(.., |p| now - p.returned >= POOL_IDLE)

            .collect();

        self.destroy(expired.into_iter().map(|p| (p.buffer, p.allocation)));

    }

    fn destroy(&self, buffers: impl IntoIterator<Item = (vk::Buffer, Allocation)>) {

        let mut buffers = buffers.into_iter().peekable();

        if buffers.peek().is_none() {

            return;

        }

        let mut allocator = self

            .shared

            .allocator

            .lock()

            .unwrap_or_else(|e| e.into_inner());

        for (buffer, allocation) in buffers {

            unsafe { self.shared.device.device.destroy_buffer(buffer, None) };

            let _ = allocator.free(allocation);

        }

    }

}

impl Drop for StagingPool {

    fn drop(&mut self) {

        let free = std::mem::take(self.free.get_mut().unwrap_or_else(|e| e.into_inner()));

        self.destroy(free.into_iter().map(|p| (p.buffer, p.allocation)));

    }

}

/// A host-visible buffer for uploads and readbacks. Dropped ones go back to

/// a pool instead of being freed. Allocating and freeing large host-visible

/// memory takes a driver lock on NVIDIA that `vkQueueSubmit` also waits on,

/// so a staging buffer made per upload on the import worker stalled the main

/// thread's submits for 1 to 3 ms.

pub(crate) struct Staging {

    pub pool: Arc<StagingPool>,

    pub buffer: vk::Buffer,

    allocation: Option<Allocation>,

    size: usize,

    usage: vk::BufferUsageFlags,

    location: MemoryLocation,

}

impl Staging {

    pub fn new(

        pool: &Arc<StagingPool>,

        size: usize,

        usage: vk::BufferUsageFlags,

        location: MemoryLocation,

    ) -> Result<Self, Error> {

        let reused = {

            let mut free = pool.free.lock().unwrap_or_else(|e| e.into_inner());

            let fits = |p: &Pooled| {

                p.usage == usage && p.location == location && p.size >= size && p.size <= size * 2

            };

            free.iter().position(fits).map(|i| free.swap_remove(i))

        };

        if let Some(p) = reused {

            return Ok(Self {

                pool: pool.clone(),

                buffer: p.buffer,

                allocation: Some(p.allocation),

                size: p.size,

                usage,

                location,

            });

        }

        let d = &pool.shared.device.device;

        let buffer = unsafe {

            d.create_buffer(

                &vk::BufferCreateInfo::default()

                    .size(size as u64)

                    .usage(usage),

                None,

            )?

        };

        let requirements = unsafe { d.get_buffer_memory_requirements(buffer) };

        let allocation = pool

            .shared

            .allocator

            .lock()

            .unwrap_or_else(|e| e.into_inner())

            .allocate(&AllocationCreateDesc {

                name: "staging",

                requirements,

                location,

                linear: true,

                allocation_scheme: AllocationScheme::GpuAllocatorManaged,

            });

        let allocation = match allocation {

            Ok(a) => a,

            Err(e) => {

                unsafe { d.destroy_buffer(buffer, None) };

                return Err(e.into());

            }

        };

        unsafe { d.bind_buffer_memory(buffer, allocation.memory(), allocation.offset())? };

        Ok(Self {

            pool: pool.clone(),

            buffer,

            allocation: Some(allocation),

            size,

            usage,

            location,

        })

    }

    /// The mapped bytes, with whatever the GPU wrote to them made visible.

    pub fn bytes(&self) -> Result<&[u8], Error> {

        self.host_sync(false)?;

        self.allocation

            .as_ref()

            .and_then(|a| a.mapped_slice())

            .ok_or(Error::StagingUnmapped)

    }

    /// Makes what the host wrote through `bytes_mut` visible to the GPU.

    pub fn flush(&self) -> Result<(), Error> {

        self.host_sync(true)

    }

    /// Flushes or invalidates the mapping. Coherent memory needs neither.

    fn host_sync(&self, flush: bool) -> Result<(), Error> {

        let Some(allocation) = self.allocation.as_ref() else {

            return Ok(());

        };

        if allocation

            .memory_properties()

            .contains(vk::MemoryPropertyFlags::HOST_COHERENT)

        {

            return Ok(());

        }

        let atom = self.pool.atom;

        let range = [vk::MappedMemoryRange::default()

            .memory(unsafe { allocation.memory() })

            .offset(allocation.offset() / atom * atom)

            .size(vk::WHOLE_SIZE)];

        let d = &self.pool.shared.device.device;

        unsafe {

            if flush {

                d.flush_mapped_memory_ranges(&range)?;

            } else {

                d.invalidate_mapped_memory_ranges(&range)?;

            }

        }

        Ok(())

    }

    pub fn bytes_mut(&mut self) -> Result<&mut [u8], Error> {

        self.allocation

            .as_mut()

            .and_then(|a| a.mapped_slice_mut())

            .ok_or(Error::StagingUnmapped)

    }

}

impl Drop for Staging {

    fn drop(&mut self) {

        let Some(allocation) = self.allocation.take() else {

            return;

        };

        let mut free = self.pool.free.lock().unwrap_or_else(|e| e.into_inner());

        let pooled: usize = free.iter().map(|p| p.size).sum();

        if pooled + self.size <= POOL_BYTES {

            free.push(Pooled {

                buffer: self.buffer,

                allocation,

                size: self.size,

                usage: self.usage,

                location: self.location,

                returned: Instant::now(),

            });

            return;

        }

        drop(free);

        self.pool.destroy([(self.buffer, allocation)]);

    }

}