thuban_gpu 0.4.2

use std::cell::RefCell;

use thuban_error::Result;

use crate::buffer::Buffer;
use crate::device::Device;

struct RingSlot {
    buffer: Buffer,
    view: RefCell<Option<Box<wgpu::BufferViewMut>>>,
}

pub struct WriteRing {
    slots: Vec<RingSlot>,
    slot_bytes: u64,
}

impl WriteRing {
    pub fn new(device: &Device, slots: u32, slot_bytes: u64) -> Result<Self> {
        let slots = (0..slots)
            .map(|_| {
                Ok(RingSlot {
                    buffer: device.create_stage_buffer(slot_bytes)?,
                    view: RefCell::new(None),
                })
            })
            .collect::<Result<Vec<_>>>()?;
        Ok(Self { slots, slot_bytes })
    }

    pub fn buffer(&self, slot: u32) -> &Buffer {
        &self.slots[slot as usize].buffer
    }

    pub fn slot_of(&self, buffer: &Buffer) -> Option<u32> {
        self.slots
            .iter()
            .position(|s| s.buffer.same(buffer))
            .map(|i| i as u32)
    }

    pub fn map(&self, slot: u32) -> Result<()> {
        let slot = &self.slots[slot as usize];
        if slot.view.borrow().is_some() {
            return Ok(());
        }
        let slice = slot.buffer.buffer.slice(..);
        let (tx, rx) = std::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Write, move |res| {
            tx.send(res).ok();
        });
        slot.buffer
            .device
            .device
            .poll(wgpu::PollType::wait_indefinitely())
            .map_err(|e| thuban_error::Error::Gpu(format!("stage buffer map poll failed: {e}")))?;
        rx.recv()
            .map_err(|_| thuban_error::Error::Gpu("stage buffer map callback dropped".to_string()))?
            .map_err(|e| thuban_error::Error::Gpu(format!("stage buffer map failed: {e}")))?;
        *slot.view.borrow_mut() = Some(Box::new(
            slice
                .get_mapped_range_mut()
                .map_err(|e| thuban_error::Error::Gpu(format!("stage buffer view failed: {e}")))?,
        ));
        Ok(())
    }

    pub fn unmap(&self, slot: u32) {
        let slot = &self.slots[slot as usize];
        if slot.view.borrow_mut().take().is_none() {
            return;
        }
        slot.buffer.buffer.unmap();
    }

    pub fn write(&self, slot: u32, offset: u64, data: &[u8]) {
        let idx = slot as usize;
        let slot = &self.slots[idx];
        assert!(
            slot.view.borrow().is_some(),
            "slot {idx} must be mapped before writing"
        );
        assert!(
            offset + data.len() as u64 <= self.slot_bytes,
            "write exceeds the ring slot"
        );
        let start = offset as usize;
        slot.view
            .borrow_mut()
            .as_mut()
            .expect("slot is mapped")
            .slice(start..start + data.len())
            .copy_from_slice(data);
    }
}