staging_belt.rs

use crossbeam_queue::SegQueue;
use futures::{
    future::FutureExt,
    stream::{FuturesUnordered, StreamExt},
};
use std::{future::Future, pin::Pin, sync::Arc};

type Fut = Pin<Box<dyn Future<Output = Result<wgpu::Buffer, wgpu::BufferAsyncError>>>>;

pub struct StagedBuffer {
    buffer: Option<wgpu::Buffer>,
    queue: Arc<SegQueue<Fut>>,
    map_mode: wgpu::MapMode,
}

pub struct StagingBelt {
    belt: Pin<Box<FuturesUnordered<Fut>>>,
    queue: Arc<SegQueue<Fut>>,
    buffer_size: wgpu::BufferAddress,
    map_mode: wgpu::MapMode,
    buffer_usage: wgpu::BufferUsage,
}

impl StagingBelt {
    pub fn new(
        buffer_size: wgpu::BufferAddress,
        buffer_usage: wgpu::BufferUsage,
        map_mode: wgpu::MapMode,
    ) -> Self {
        Self {
            belt: Box::pin(FuturesUnordered::new()),
            queue: Arc::new(SegQueue::new()),
            buffer_size,
            map_mode,
            buffer_usage,
        }
    }

    pub fn retrieve(&mut self, device: &wgpu::Device) -> StagedBuffer {
        while let Ok(fut) = self.queue.pop() {
            self.belt.push(fut);
        }

        let buffer = if let Some(Some(res)) = self.belt.as_mut().next().now_or_never() {
            res.expect("staging belt failed to map buffer")
        } else {
            device.create_buffer(&wgpu::BufferDescriptor {
                size: self.buffer_size,
                usage: self.buffer_usage | wgpu::BufferUsage::MAP_READ,
                mapped_at_creation: true,
                label: None,
            })
        };

        StagedBuffer {
            buffer: Some(buffer),
            queue: Arc::clone(&self.queue),
            map_mode: self.map_mode,
        }
    }
}

impl std::ops::Deref for StagedBuffer {
    type Target = wgpu::Buffer;

    fn deref(&self) -> &wgpu::Buffer {
        self.buffer.as_ref().unwrap()
    }
}

impl Drop for StagedBuffer {
    fn drop(&mut self) {
        let buffer = self.buffer.take().unwrap();

        buffer.unmap();

        let slice = buffer.slice(..);

        let fut = slice
            .map_async(self.map_mode)
            .map(|res| res.map(|_| buffer));

        self.queue.push(Box::pin(fut));
    }
}