wgpu-hal 0.13.2

use crate::auxil::dxgi::result::HResult as _;
use bit_set::BitSet;
use parking_lot::Mutex;
use range_alloc::RangeAllocator;
use std::fmt;

const HEAP_SIZE_FIXED: usize = 64;

#[derive(Copy, Clone)]
pub(super) struct DualHandle {
    cpu: native::CpuDescriptor,
    pub gpu: native::GpuDescriptor,
    /// How large the block allocated to this handle is.
    count: u64,
}

impl fmt::Debug for DualHandle {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("DualHandle")
            .field("cpu", &self.cpu.ptr)
            .field("gpu", &self.gpu.ptr)
            .field("count", &self.count)
            .finish()
    }
}

type DescriptorIndex = u64;

pub(super) struct GeneralHeap {
    pub raw: native::DescriptorHeap,
    ty: native::DescriptorHeapType,
    handle_size: u64,
    total_handles: u64,
    start: DualHandle,
    ranges: Mutex<RangeAllocator<DescriptorIndex>>,
}

impl GeneralHeap {
    pub(super) fn new(
        device: native::Device,
        ty: native::DescriptorHeapType,
        total_handles: u64,
    ) -> Result<Self, crate::DeviceError> {
        let raw = {
            profiling::scope!("ID3D12Device::CreateDescriptorHeap");
            device
                .create_descriptor_heap(
                    total_handles as u32,
                    ty,
                    native::DescriptorHeapFlags::SHADER_VISIBLE,
                    0,
                )
                .into_device_result("Descriptor heap creation")?
        };

        Ok(Self {
            raw,
            ty,
            handle_size: device.get_descriptor_increment_size(ty) as u64,
            total_handles,
            start: DualHandle {
                cpu: raw.start_cpu_descriptor(),
                gpu: raw.start_gpu_descriptor(),
                count: 0,
            },
            ranges: Mutex::new(RangeAllocator::new(0..total_handles)),
        })
    }

    pub(super) fn at(&self, index: DescriptorIndex, count: u64) -> DualHandle {
        assert!(index < self.total_handles);
        DualHandle {
            cpu: self.cpu_descriptor_at(index),
            gpu: self.gpu_descriptor_at(index),
            count,
        }
    }

    fn cpu_descriptor_at(&self, index: u64) -> native::CpuDescriptor {
        native::CpuDescriptor {
            ptr: self.start.cpu.ptr + (self.handle_size * index) as usize,
        }
    }

    fn gpu_descriptor_at(&self, index: u64) -> native::GpuDescriptor {
        native::GpuDescriptor {
            ptr: self.start.gpu.ptr + self.handle_size * index,
        }
    }

    pub(super) fn allocate_slice(&self, count: u64) -> Result<DescriptorIndex, crate::DeviceError> {
        let range = self.ranges.lock().allocate_range(count).map_err(|err| {
            log::error!("Unable to allocate descriptors: {:?}", err);
            crate::DeviceError::OutOfMemory
        })?;
        Ok(range.start)
    }

    /// Free handles previously given out by this `DescriptorHeapSlice`.
    /// Do not use this with handles not given out by this `DescriptorHeapSlice`.
    pub(crate) fn free_slice(&self, handle: DualHandle) {
        let start = (handle.gpu.ptr - self.start.gpu.ptr) / self.handle_size;
        self.ranges.lock().free_range(start..start + handle.count);
    }
}

/// Fixed-size free-list allocator for CPU descriptors.
struct FixedSizeHeap {
    raw: native::DescriptorHeap,
    /// Bit flag representation of available handles in the heap.
    ///
    ///  0 - Occupied
    ///  1 - free
    availability: u64,
    handle_size: usize,
    start: native::CpuDescriptor,
}

impl FixedSizeHeap {
    fn new(device: native::Device, ty: native::DescriptorHeapType) -> Self {
        let (heap, _hr) = device.create_descriptor_heap(
            HEAP_SIZE_FIXED as _,
            ty,
            native::DescriptorHeapFlags::empty(),
            0,
        );

        Self {
            handle_size: device.get_descriptor_increment_size(ty) as _,
            availability: !0, // all free!
            start: heap.start_cpu_descriptor(),
            raw: heap,
        }
    }

    fn alloc_handle(&mut self) -> native::CpuDescriptor {
        // Find first free slot.
        let slot = self.availability.trailing_zeros() as usize;
        assert!(slot < HEAP_SIZE_FIXED);
        // Set the slot as occupied.
        self.availability ^= 1 << slot;

        native::CpuDescriptor {
            ptr: self.start.ptr + self.handle_size * slot,
        }
    }

    fn free_handle(&mut self, handle: native::CpuDescriptor) {
        let slot = (handle.ptr - self.start.ptr) / self.handle_size;
        assert!(slot < HEAP_SIZE_FIXED);
        assert_eq!(self.availability & (1 << slot), 0);
        self.availability ^= 1 << slot;
    }

    fn is_full(&self) -> bool {
        self.availability == 0
    }

    unsafe fn destroy(&self) {
        self.raw.destroy();
    }
}

#[derive(Clone, Copy)]
pub(super) struct Handle {
    pub raw: native::CpuDescriptor,
    heap_index: usize,
}

impl fmt::Debug for Handle {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Handle")
            .field("ptr", &self.raw.ptr)
            .field("heap_index", &self.heap_index)
            .finish()
    }
}

pub(super) struct CpuPool {
    device: native::Device,
    ty: native::DescriptorHeapType,
    heaps: Vec<FixedSizeHeap>,
    avaliable_heap_indices: BitSet,
}

impl CpuPool {
    pub(super) fn new(device: native::Device, ty: native::DescriptorHeapType) -> Self {
        Self {
            device,
            ty,
            heaps: Vec::new(),
            avaliable_heap_indices: BitSet::new(),
        }
    }

    pub(super) fn alloc_handle(&mut self) -> Handle {
        let heap_index = self
            .avaliable_heap_indices
            .iter()
            .next()
            .unwrap_or_else(|| {
                // Allocate a new heap
                let id = self.heaps.len();
                self.heaps.push(FixedSizeHeap::new(self.device, self.ty));
                self.avaliable_heap_indices.insert(id);
                id
            });

        let heap = &mut self.heaps[heap_index];
        let handle = Handle {
            raw: heap.alloc_handle(),
            heap_index,
        };
        if heap.is_full() {
            self.avaliable_heap_indices.remove(heap_index);
        }

        handle
    }

    pub(super) fn free_handle(&mut self, handle: Handle) {
        self.heaps[handle.heap_index].free_handle(handle.raw);
        self.avaliable_heap_indices.insert(handle.heap_index);
    }

    pub(super) unsafe fn destroy(&self) {
        for heap in &self.heaps {
            heap.destroy();
        }
    }
}

pub(super) struct CpuHeapInner {
    pub raw: native::DescriptorHeap,
    pub stage: Vec<native::CpuDescriptor>,
}

pub(super) struct CpuHeap {
    pub inner: Mutex<CpuHeapInner>,
    start: native::CpuDescriptor,
    handle_size: u32,
    total: u32,
}

unsafe impl Send for CpuHeap {}
unsafe impl Sync for CpuHeap {}

impl CpuHeap {
    pub(super) fn new(
        device: native::Device,
        ty: native::DescriptorHeapType,
        total: u32,
    ) -> Result<Self, crate::DeviceError> {
        let handle_size = device.get_descriptor_increment_size(ty);
        let raw = device
            .create_descriptor_heap(total, ty, native::DescriptorHeapFlags::empty(), 0)
            .into_device_result("CPU descriptor heap creation")?;

        Ok(Self {
            inner: Mutex::new(CpuHeapInner {
                raw,
                stage: Vec::new(),
            }),
            start: raw.start_cpu_descriptor(),
            handle_size,
            total,
        })
    }

    pub(super) fn at(&self, index: u32) -> native::CpuDescriptor {
        native::CpuDescriptor {
            ptr: self.start.ptr + (self.handle_size * index) as usize,
        }
    }

    pub(super) unsafe fn destroy(self) {
        self.inner.into_inner().raw.destroy();
    }
}

impl fmt::Debug for CpuHeap {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("CpuHeap")
            .field("start", &self.start.ptr)
            .field("handle_size", &self.handle_size)
            .field("total", &self.total)
            .finish()
    }
}

pub(super) unsafe fn upload(
    device: native::Device,
    src: &CpuHeapInner,
    dst: &GeneralHeap,
    dummy_copy_counts: &[u32],
) -> Result<DualHandle, crate::DeviceError> {
    let count = src.stage.len() as u32;
    let index = dst.allocate_slice(count as u64)?;
    device.CopyDescriptors(
        1,
        &dst.cpu_descriptor_at(index),
        &count,
        count,
        src.stage.as_ptr(),
        dummy_copy_counts.as_ptr(),
        dst.ty as u32,
    );
    Ok(dst.at(index, count as u64))
}