Implement Queue::write_buffer_with (#2777)

* implement Queue::write_buffer_with

* address comments

* update doc

* Fix copy span location

Co-authored-by: Connor Fitzgerald <[email protected]>

wgpu-core/src/hub.rs

@@ -5,7 +5,7 @@ use crate::{
     id,
     instance::{Adapter, HalSurface, Instance, Surface},
     pipeline::{ComputePipeline, RenderPipeline, ShaderModule},
-    resource::{Buffer, QuerySet, Sampler, Texture, TextureClearMode, TextureView},
+    resource::{Buffer, QuerySet, Sampler, StagingBuffer, Texture, TextureClearMode, TextureView},
     Epoch, Index,
 };
 
@@ -351,6 +351,7 @@ impl<A: HalApi> Access<Buffer<A>> for CommandBuffer<A> {}
 impl<A: HalApi> Access<Buffer<A>> for ComputePipeline<A> {}
 impl<A: HalApi> Access<Buffer<A>> for RenderPipeline<A> {}
 impl<A: HalApi> Access<Buffer<A>> for QuerySet<A> {}
+impl<A: HalApi> Access<StagingBuffer<A>> for Device<A> {}
 impl<A: HalApi> Access<Texture<A>> for Root {}
 impl<A: HalApi> Access<Texture<A>> for Device<A> {}
 impl<A: HalApi> Access<Texture<A>> for Buffer<A> {}
@@ -452,6 +453,7 @@ pub trait GlobalIdentityHandlerFactory:
     + IdentityHandlerFactory<id::ComputePipelineId>
     + IdentityHandlerFactory<id::QuerySetId>
     + IdentityHandlerFactory<id::BufferId>
+    + IdentityHandlerFactory<id::StagingBufferId>
     + IdentityHandlerFactory<id::TextureId>
     + IdentityHandlerFactory<id::TextureViewId>
     + IdentityHandlerFactory<id::SamplerId>
@@ -639,6 +641,7 @@ pub struct Hub<A: HalApi, F: GlobalIdentityHandlerFactory> {
     pub compute_pipelines: Registry<ComputePipeline<A>, id::ComputePipelineId, F>,
     pub query_sets: Registry<QuerySet<A>, id::QuerySetId, F>,
     pub buffers: Registry<Buffer<A>, id::BufferId, F>,
+    pub staging_buffers: Registry<StagingBuffer<A>, id::StagingBufferId, F>,
     pub textures: Registry<Texture<A>, id::TextureId, F>,
     pub texture_views: Registry<TextureView<A>, id::TextureViewId, F>,
     pub samplers: Registry<Sampler<A>, id::SamplerId, F>,
@@ -659,6 +662,7 @@ impl<A: HalApi, F: GlobalIdentityHandlerFactory> Hub<A, F> {
             compute_pipelines: Registry::new(A::VARIANT, factory),
             query_sets: Registry::new(A::VARIANT, factory),
             buffers: Registry::new(A::VARIANT, factory),
+            staging_buffers: Registry::new(A::VARIANT, factory),
             textures: Registry::new(A::VARIANT, factory),
             texture_views: Registry::new(A::VARIANT, factory),
             samplers: Registry::new(A::VARIANT, factory),

wgpu-core/src/id.rs

@@ -196,6 +196,7 @@ pub type DeviceId = Id<crate::device::Device<Dummy>>;
 pub type QueueId = DeviceId;
 // Resource
 pub type BufferId = Id<crate::resource::Buffer<Dummy>>;
+pub type StagingBufferId = Id<crate::resource::StagingBuffer<Dummy>>;
 pub type TextureViewId = Id<crate::resource::TextureView<Dummy>>;
 pub type TextureId = Id<crate::resource::Texture<Dummy>>;
 pub type SamplerId = Id<crate::resource::Sampler<Dummy>>;

wgpu-core/src/resource.rs

@@ -186,6 +186,24 @@ impl<A: hal::Api> Resource for Buffer<A> {
     }
 }
 
+pub struct StagingBuffer<A: hal::Api> {
+    pub(crate) raw: A::Buffer,
+    pub(crate) size: wgt::BufferAddress,
+    pub(crate) is_coherent: bool,
+}
+
+impl<A: hal::Api> Resource for StagingBuffer<A> {
+    const TYPE: &'static str = "StagingBuffer";
+
+    fn life_guard(&self) -> &LifeGuard {
+        unreachable!()
+    }
+
+    fn label(&self) -> &str {
+        "<StagingBuffer>"
+    }
+}
+
 pub type TextureDescriptor<'a> = wgt::TextureDescriptor<Label<'a>>;
 
 #[derive(Debug)]

wgpu/src/backend/direct.rs

@@ -2192,6 +2192,58 @@ impl crate::Context for Context {
         }
     }
 
+    fn queue_validate_write_buffer(
+        &self,
+        queue: &Self::QueueId,
+        buffer: &Self::BufferId,
+        offset: wgt::BufferAddress,
+        size: wgt::BufferSize,
+    ) {
+        let global = &self.0;
+        match wgc::gfx_select!(
+            *queue => global.queue_validate_write_buffer(*queue, buffer.id, offset, size.get())
+        ) {
+            Ok(()) => (),
+            Err(err) => self.handle_error_fatal(err, "Queue::write_buffer_with"),
+        }
+    }
+
+    fn queue_create_staging_buffer(
+        &self,
+        queue: &Self::QueueId,
+        size: wgt::BufferSize,
+    ) -> QueueWriteBuffer {
+        let global = &self.0;
+        match wgc::gfx_select!(
+            *queue => global.queue_create_staging_buffer(*queue, size, PhantomData)
+        ) {
+            Ok((buffer_id, ptr)) => QueueWriteBuffer {
+                buffer_id,
+                mapping: BufferMappedRange {
+                    ptr,
+                    size: size.get() as usize,
+                },
+            },
+            Err(err) => self.handle_error_fatal(err, "Queue::write_buffer_with"),
+        }
+    }
+
+    fn queue_write_staging_buffer(
+        &self,
+        queue: &Self::QueueId,
+        buffer: &Self::BufferId,
+        offset: wgt::BufferAddress,
+        staging_buffer: &QueueWriteBuffer,
+    ) {
+        let global = &self.0;
+        match wgc::gfx_select!(
+            *queue => global.queue_write_staging_buffer(*queue, buffer.id, offset, staging_buffer.buffer_id)
+        ) {
+            Ok(()) => (),
+            Err(err) => self.handle_error_fatal(err, "Queue::write_buffer_with"),
+        }
+    }
+
     fn queue_write_texture(
         &self,
         queue: &Self::QueueId,
@@ -2324,6 +2376,27 @@ fn default_error_handler(err: crate::Error) {
     panic!("wgpu error: {}\n", err);
 }
 
+#[derive(Debug)]
+pub struct QueueWriteBuffer {
+    buffer_id: wgc::id::StagingBufferId,
+    mapping: BufferMappedRange,
+}
+
+impl std::ops::Deref for QueueWriteBuffer {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        panic!("QueueWriteBuffer is write-only!");
+    }
+}
+
+impl std::ops::DerefMut for QueueWriteBuffer {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        use crate::BufferMappedRangeSlice;
+        self.mapping.slice_mut()
+    }
+}
+
 #[derive(Debug)]
 pub struct BufferMappedRange {
     ptr: *mut u8,

wgpu/src/backend/mod.rs

@@ -1,9 +1,9 @@
 #[cfg(all(target_arch = "wasm32", not(feature = "webgl")))]
 mod web;
 #[cfg(all(target_arch = "wasm32", not(feature = "webgl")))]
-pub(crate) use web::{BufferMappedRange, Context};
+pub(crate) use web::{BufferMappedRange, Context, QueueWriteBuffer};
 
 #[cfg(any(not(target_arch = "wasm32"), feature = "webgl"))]
 mod direct;
 #[cfg(any(not(target_arch = "wasm32"), feature = "webgl"))]
-pub(crate) use direct::{BufferMappedRange, Context};
+pub(crate) use direct::{BufferMappedRange, Context, QueueWriteBuffer};

wgpu/src/backend/web.rs

@@ -2197,6 +2197,34 @@ impl crate::Context for Context {
             );
     }
 
+    fn queue_validate_write_buffer(
+        &self,
+        _queue: &Self::QueueId,
+        _buffer: &Self::BufferId,
+        _offset: wgt::BufferAddress,
+        _size: wgt::BufferSize,
+    ) {
+        // TODO
+    }
+
+    fn queue_create_staging_buffer(
+        &self,
+        _queue: &Self::QueueId,
+        size: wgt::BufferSize,
+    ) -> QueueWriteBuffer {
+        QueueWriteBuffer(vec![0; size.get() as usize].into_boxed_slice())
+    }
+
+    fn queue_write_staging_buffer(
+        &self,
+        queue: &Self::QueueId,
+        buffer: &Self::BufferId,
+        offset: wgt::BufferAddress,
+        staging_buffer: &QueueWriteBuffer,
+    ) {
+        self.queue_write_buffer(queue, buffer, offset, staging_buffer)
+    }
+
     fn queue_write_texture(
         &self,
         queue: &Self::QueueId,
@@ -2262,6 +2290,23 @@ impl crate::Context for Context {
 
 pub(crate) type SurfaceOutputDetail = ();
 
+#[derive(Debug)]
+pub struct QueueWriteBuffer(Box<[u8]>);
+
+impl std::ops::Deref for QueueWriteBuffer {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        panic!("QueueWriteBuffer is write-only!");
+    }
+}
+
+impl std::ops::DerefMut for QueueWriteBuffer {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
 #[derive(Debug)]
 pub struct BufferMappedRange {
     actual_mapping: js_sys::Uint8Array,

wgpu/src/lib.rs

@@ -43,7 +43,7 @@ pub use wgt::{
     QUERY_SIZE, VERTEX_STRIDE_ALIGNMENT,
 };
 
-use backend::{BufferMappedRange, Context as C};
+use backend::{BufferMappedRange, Context as C, QueueWriteBuffer};
 
 /// Filter for error scopes.
 #[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
@@ -482,6 +482,25 @@ trait Context: Debug + Send + Sized + Sync {
         offset: BufferAddress,
         data: &[u8],
     );
+    fn queue_validate_write_buffer(
+        &self,
+        queue: &Self::QueueId,
+        buffer: &Self::BufferId,
+        offset: wgt::BufferAddress,
+        size: wgt::BufferSize,
+    );
+    fn queue_create_staging_buffer(
+        &self,
+        queue: &Self::QueueId,
+        size: BufferSize,
+    ) -> QueueWriteBuffer;
+    fn queue_write_staging_buffer(
+        &self,
+        queue: &Self::QueueId,
+        buffer: &Self::BufferId,
+        offset: BufferAddress,
+        staging_buffer: &QueueWriteBuffer,
+    );
     fn queue_write_texture(
         &self,
         queue: &Self::QueueId,
@@ -3376,6 +3395,40 @@ impl<'a> RenderBundleEncoder<'a> {
     }
 }
 
+/// A write-only view into a staging buffer
+pub struct QueueWriteBufferView<'a> {
+    queue: &'a Queue,
+    buffer: &'a Buffer,
+    offset: BufferAddress,
+    inner: QueueWriteBuffer,
+}
+
+impl<'a> std::ops::Deref for QueueWriteBufferView<'a> {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        panic!("QueueWriteBufferView is write-only!");
+    }
+}
+
+impl<'a> std::ops::DerefMut for QueueWriteBufferView<'a> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.inner
+    }
+}
+
+impl<'a> Drop for QueueWriteBufferView<'a> {
+    fn drop(&mut self) {
+        Context::queue_write_staging_buffer(
+            &*self.queue.context,
+            &self.queue.id,
+            &self.buffer.id,
+            self.offset,
+            &self.inner,
+        );
+    }
+}
+
 impl Queue {
     /// Schedule a data write into `buffer` starting at `offset`.
     ///
@@ -3388,6 +3441,32 @@ impl Queue {
         Context::queue_write_buffer(&*self.context, &self.id, &buffer.id, offset, data)
     }
 
+    /// Schedule a data write into `buffer` starting at `offset` via the returned [QueueWriteBufferView].
+    ///
+    /// The returned value can be dereferenced to a `&mut [u8]`; dereferencing it to a `&[u8]` panics!
+    ///
+    /// This method is intended to have low performance costs.
+    /// As such, the write is not immediately submitted, and instead enqueued
+    /// internally to happen at the start of the next `submit()` call.
+    ///
+    /// This method fails if `size` is greater than the size of `buffer` starting at `offset`.
+    #[must_use]
+    pub fn write_buffer_with<'a>(
+        &'a self,
+        buffer: &'a Buffer,
+        offset: BufferAddress,
+        size: BufferSize,
+    ) -> QueueWriteBufferView<'a> {
+        Context::queue_validate_write_buffer(&*self.context, &self.id, &buffer.id, offset, size);
+        let staging_buffer = Context::queue_create_staging_buffer(&*self.context, &self.id, size);
+        QueueWriteBufferView {
+            queue: self,
+            buffer,
+            offset,
+            inner: staging_buffer,
+        }
+    }
+
     /// Schedule a data write into `texture`.
     ///
     /// This method is intended to have low performance costs.