gpu-alloc 0.4.5

use {
    crate::{
        align_down, align_up,
        error::AllocationError,
        heap::Heap,
        util::{arc_unwrap, is_arc_unique},
        MemoryBounds,
    },
    alloc::sync::Arc,
    core::{cmp::Ordering, ptr::NonNull},
    gpu_alloc_types::{AllocationFlags, DeviceMapError, MemoryDevice, MemoryPropertyFlags},
};

unsafe fn opt_ptr_add(ptr: Option<NonNull<u8>>, size: u64) -> Option<NonNull<u8>> {
    ptr.map(|ptr| {
        // Size is within memory region started at `ptr`.
        // size is within `chunk_size` that fits `isize`.
        NonNull::new_unchecked(ptr.as_ptr().offset(size as isize))
    })
}

#[derive(Debug)]
pub(super) struct FreeList<M> {
    array: Vec<FreeListRegion<M>>,
    total: u64,
    counter: u64,
}

impl<M> FreeList<M> {
    pub fn new() -> Self {
        FreeList {
            array: Vec::new(),
            total: 0,
            counter: 0,
        }
    }

    pub fn get_block_from_new_memory(
        &mut self,
        memory: Arc<M>,
        memory_size: u64,
        ptr: Option<NonNull<u8>>,
        align_mask: u64,
        size: u64,
    ) -> FreeListBlock<M> {
        debug_assert!(size <= memory_size);

        self.counter += 1;
        self.array.push(FreeListRegion {
            memory,
            ptr,
            chunk: self.counter,
            start: 0,
            end: memory_size,
        });
        self.total += memory_size;
        self.get_block_at(self.array.len() - 1, align_mask, size)
    }

    pub fn get_block(&mut self, align_mask: u64, size: u64) -> Option<FreeListBlock<M>> {
        match &self.array[..] {
            [] => None,
            [rest @ .., last] => {
                let (index, _) = std::iter::once((rest.len(), last))
                    .chain(rest.iter().enumerate())
                    .find(|(_, region)| match region.end.checked_sub(size) {
                        Some(start) => {
                            let aligned_start = align_down(start, align_mask);
                            aligned_start >= region.start
                        }
                        None => false,
                    })?;

                Some(self.get_block_at(index, align_mask, size))
            }
        }
    }

    fn get_block_at(&mut self, index: usize, align_mask: u64, size: u64) -> FreeListBlock<M> {
        let region = &mut self.array[index];

        let start = region.end - size;
        let aligned_start = align_down(start, align_mask);

        if aligned_start > region.start {
            let block = FreeListBlock {
                offset: aligned_start,
                size: region.end - aligned_start,
                chunk: region.chunk,
                ptr: unsafe { opt_ptr_add(region.ptr, aligned_start - region.start) },
                memory: region.memory.clone(),
            };

            region.end = aligned_start;
            self.total -= block.size;

            block
        } else {
            debug_assert_eq!(aligned_start, region.start);
            let region = self.array.remove(index);
            self.total -= region.end - region.start;
            region.into_block()
        }
    }

    pub fn insert_block(&mut self, block: FreeListBlock<M>) {
        let block_size = block.size;

        match self.array.binary_search_by(|b| b.cmp(&block)) {
            Ok(_) => {
                panic!("Overlapping block found in free list");
            }
            Err(index) if self.array.len() > index => match &mut self.array[..=index] {
                [] => unreachable!(),
                [next] => {
                    debug_assert!(!next.is_suffix_block(&block));

                    if next.is_prefix_block(&block) {
                        next.merge_prefix_block(block);
                    } else {
                        self.array.insert(0, FreeListRegion::from_block(block));
                    }
                }
                [.., prev, next] => {
                    debug_assert!(!prev.is_prefix_block(&block));
                    debug_assert!(!next.is_suffix_block(&block));

                    if next.is_prefix_block(&block) {
                        next.merge_prefix_block(block);

                        if prev.consecutive(&next) {
                            let next = self.array.remove(index);
                            let prev = &mut self.array[index - 1];
                            prev.merge(next);
                        }
                    } else if prev.is_suffix_block(&block) {
                        prev.merge_suffix_block(block);
                    } else {
                        self.array.insert(index, FreeListRegion::from_block(block));
                    }
                }
            },
            Err(_) => {
                self.array.push(FreeListRegion::from_block(block));
            }
        }
        self.total += block_size;
    }

    pub fn drain(&mut self, dealloc_threshold: u64) -> Option<impl Iterator<Item = M> + '_> {
        if self.total >= dealloc_threshold {
            // Time to deallocate

            // This code was adapted from `Vec::retain`
            let len = self.array.len();
            let mut del = 0;
            {
                let regions = &mut self.array[..];

                for i in 0..len {
                    if self.total >= dealloc_threshold && is_arc_unique(&mut regions[i].memory) {
                        del += 1;
                    } else if del > 0 {
                        regions.swap(i - del, i);
                    }
                }
            }

            if del > 0 {
                let total = &mut self.total;
                return Some(self.array.drain(len - del..).map(move |region| {
                    debug_assert_eq!(region.start, 0);
                    *total -= region.end;
                    unsafe { arc_unwrap(region.memory) }
                }));
            }
        }
        None
    }
}

#[derive(Debug)]
struct FreeListRegion<M> {
    memory: Arc<M>,
    ptr: Option<NonNull<u8>>,
    chunk: u64,
    start: u64,
    end: u64,
}

unsafe impl<M> Sync for FreeListRegion<M> where M: Sync {}
unsafe impl<M> Send for FreeListRegion<M> where M: Send {}

impl<M> FreeListRegion<M> {
    pub fn cmp(&self, block: &FreeListBlock<M>) -> Ordering {
        debug_assert_eq!(
            Arc::ptr_eq(&self.memory, &block.memory),
            self.chunk == block.chunk
        );

        if self.chunk == block.chunk {
            debug_assert_eq!(
                Ord::cmp(&self.start, &block.offset),
                Ord::cmp(&self.end, &(block.offset + block.size)),
                "Free region {{ start: {}, end: {} }} overlaps with block {{ offset: {}, size: {} }}",
                self.start,
                self.end,
                block.offset,
                block.size,
            );
        }

        Ord::cmp(&self.chunk, &block.chunk).then(Ord::cmp(&self.start, &block.offset))
    }

    fn from_block(block: FreeListBlock<M>) -> Self {
        FreeListRegion {
            memory: block.memory,
            chunk: block.chunk,
            ptr: block.ptr,
            start: block.offset,
            end: block.offset + block.size,
        }
    }

    fn into_block(self) -> FreeListBlock<M> {
        FreeListBlock {
            memory: self.memory,
            chunk: self.chunk,
            ptr: self.ptr,
            offset: self.start,
            size: self.end - self.start,
        }
    }

    fn consecutive(&self, other: &Self) -> bool {
        if self.chunk != other.chunk {
            return false;
        }

        debug_assert!(Arc::ptr_eq(&self.memory, &other.memory));

        debug_assert_eq!(
            Ord::cmp(&self.start, &other.start),
            Ord::cmp(&self.end, &other.end)
        );

        self.end == other.start
    }

    fn merge(&mut self, next: FreeListRegion<M>) {
        debug_assert!(self.consecutive(&next));
        self.end = next.end;
    }

    fn is_prefix_block(&self, block: &FreeListBlock<M>) -> bool {
        if self.chunk != block.chunk {
            return false;
        }

        debug_assert!(Arc::ptr_eq(&self.memory, &block.memory));

        debug_assert_eq!(
            Ord::cmp(&self.start, &block.offset),
            Ord::cmp(&self.end, &(block.offset + block.size))
        );

        self.start == (block.offset + block.size)
    }

    fn merge_prefix_block(&mut self, block: FreeListBlock<M>) {
        debug_assert!(self.is_prefix_block(&block));
        self.start = block.offset;
        self.ptr = block.ptr;
    }

    fn is_suffix_block(&self, block: &FreeListBlock<M>) -> bool {
        if self.chunk != block.chunk {
            return false;
        }

        debug_assert!(Arc::ptr_eq(&self.memory, &block.memory));

        debug_assert_eq!(
            Ord::cmp(&self.start, &block.offset),
            Ord::cmp(&self.end, &(block.offset + block.size))
        );

        self.end == block.offset
    }

    fn merge_suffix_block(&mut self, block: FreeListBlock<M>) {
        debug_assert!(self.is_suffix_block(&block));
        self.end += block.size;
    }
}

#[derive(Debug)]
pub struct FreeListBlock<M> {
    pub memory: Arc<M>,
    pub ptr: Option<NonNull<u8>>,
    pub chunk: u64,
    pub offset: u64,
    pub size: u64,
}

unsafe impl<M> Sync for FreeListBlock<M> where M: Sync {}
unsafe impl<M> Send for FreeListBlock<M> where M: Send {}

#[derive(Debug)]
pub(crate) struct FreeListAllocator<M> {
    freelist: FreeList<M>,
    dealloc_threshold: u64,
    chunk_size: u64,
    memory_type: u32,
    props: MemoryPropertyFlags,
    atom_mask: u64,
}

impl<M> FreeListAllocator<M>
where
    M: MemoryBounds + 'static,
{
    pub fn new(
        chunk_size: u64,
        dealloc_threshold: u64,
        memory_type: u32,
        props: MemoryPropertyFlags,
        atom_mask: u64,
    ) -> Self {
        debug_assert_eq!(align_down(chunk_size, atom_mask), chunk_size);

        let chunk_size = min(chunk_size, isize::max_value());

        FreeListAllocator {
            freelist: FreeList::new(),
            chunk_size,
            dealloc_threshold,
            memory_type,
            props,
            atom_mask,
        }
    }

    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self, device)))]
    pub unsafe fn alloc(
        &mut self,
        device: &impl MemoryDevice<M>,
        size: u64,
        align_mask: u64,
        flags: AllocationFlags,
        heap: &mut Heap,
        allocations_remains: &mut u32,
    ) -> Result<FreeListBlock<M>, AllocationError> {
        debug_assert!(
            self.chunk_size >= size,
            "GpuAllocator must not request allocations equal or greater to chunks size"
        );

        let size = align_up(size, self.atom_mask).expect(
            "Any value not greater than chunk size (which is aligned) has to fit for alignment",
        );

        let align_mask = align_mask | self.atom_mask;
        let host_visible = self.host_visible();

        if let Some(block) = self.freelist.get_block(align_mask, size) {
            return Ok(block);
        }

        // New allocation is required.
        if *allocations_remains == 0 {
            return Err(AllocationError::TooManyObjects);
        }
        let mut memory = device.allocate_memory(self.chunk_size, self.memory_type, flags)?;
        *allocations_remains -= 1;
        heap.alloc(self.chunk_size);

        // Map host visible allocations
        let ptr = if host_visible {
            match device.map_memory(&mut memory, 0, self.chunk_size) {
                Ok(ptr) => Some(ptr),
                Err(DeviceMapError::MapFailed) => {
                    #[cfg(feature = "tracing")]
                    tracing::error!("Failed to map host-visible memory in linear allocator");
                    device.deallocate_memory(memory);
                    *allocations_remains += 1;
                    heap.dealloc(self.chunk_size);

                    return Err(AllocationError::OutOfHostMemory);
                }
                Err(DeviceMapError::OutOfDeviceMemory) => {
                    return Err(AllocationError::OutOfDeviceMemory);
                }
                Err(DeviceMapError::OutOfHostMemory) => {
                    return Err(AllocationError::OutOfHostMemory);
                }
            }
        } else {
            None
        };

        let memory = Arc::new(memory);
        let block =
            self.freelist
                .get_block_from_new_memory(memory, self.chunk_size, ptr, align_mask, size);

        Ok(block)
    }

    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self, device)))]
    pub unsafe fn dealloc(
        &mut self,
        device: &impl MemoryDevice<M>,
        block: FreeListBlock<M>,
        heap: &mut Heap,
        allocations_remains: &mut u32,
    ) {
        debug_assert!(block.size < self.chunk_size);
        debug_assert_ne!(block.size, 0);
        self.freelist.insert_block(block);

        if let Some(memory) = self.freelist.drain(self.dealloc_threshold) {
            let chunk_size = self.chunk_size;
            memory.for_each(|m| {
                device.deallocate_memory(m);
                *allocations_remains += 1;
                heap.dealloc(chunk_size);
            });
        }
    }

    /// Deallocates leftover memory objects.
    /// Should be used before dropping.
    ///
    /// # Safety
    ///
    /// * `device` must be one with `DeviceProperties` that were provided to create this `GpuAllocator` instance
    /// * Same `device` instance must be used for all interactions with one `GpuAllocator` instance
    ///   and memory blocks allocated from it
    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self, device)))]
    pub unsafe fn cleanup(&mut self, device: &impl MemoryDevice<M>) {
        if let Some(memory) = self.freelist.drain(0) {
            memory.for_each(|m| device.deallocate_memory(m));
        }
    }

    fn host_visible(&self) -> bool {
        self.props.contains(MemoryPropertyFlags::HOST_VISIBLE)
    }
}

fn min<L, R>(l: L, r: R) -> L
where
    R: core::convert::TryInto<L>,
    L: Ord,
{
    match r.try_into() {
        Ok(r) => core::cmp::min(l, r),
        Err(_) => l,
    }
}