staging belt: allow creation with custom buffer usages

tests/tests/wgpu-validation/util.rs

@@ -1,6 +1,7 @@
 //! Tests of [`wgpu::util`].
 
 use nanorand::Rng;
+use wgpu::BufferUsages;
 
 /// Generate (deterministic) random staging belt operations to exercise its logic.
 #[test]
@@ -39,3 +40,44 @@ fn staging_belt_random_test() {
         belt.recall();
     }
 }
+
+#[test]
+fn staging_belt_panics_with_invalid_buffer_usages() {
+    fn test_if_panics(usage: BufferUsages) -> bool {
+        std::panic::catch_unwind(|| {
+            let (device, _queue) = wgpu::Device::noop(&wgpu::DeviceDescriptor::default());
+            let _belt = wgpu::util::StagingBelt::new_with_buffer_usages(device.clone(), 512, usage);
+        })
+        .is_err()
+    }
+
+    for mut usage in BufferUsages::all()
+        .difference(BufferUsages::COPY_SRC | BufferUsages::MAP_WRITE)
+        .iter()
+    {
+        assert!(test_if_panics(usage), "StagingBelt::new_with_buffer_usages should panic without MAPPABLE_PRIMARY_BUFFERS with usage={usage:?}");
+
+        usage.insert(BufferUsages::MAP_WRITE);
+        assert!(test_if_panics(usage), "StagingBelt::new_with_buffer_usages should panic without MAPPABLE_PRIMARY_BUFFERS with usage={usage:?}");
+    }
+}
+
+#[test]
+fn staging_belt_works_with_non_exclusive_buffer_usages() {
+    let (device, _queue) = wgpu::Device::noop(&wgpu::DeviceDescriptor::default());
+    let _belt = wgpu::util::StagingBelt::new_with_buffer_usages(
+        device.clone(),
+        512,
+        BufferUsages::COPY_SRC,
+    );
+    let _belt = wgpu::util::StagingBelt::new_with_buffer_usages(
+        device.clone(),
+        512,
+        BufferUsages::COPY_SRC | BufferUsages::MAP_WRITE,
+    );
+    let _belt = wgpu::util::StagingBelt::new_with_buffer_usages(
+        device.clone(),
+        512,
+        BufferUsages::MAP_WRITE,
+    );
+}

wgpu/src/util/belt.rs

@@ -5,6 +5,7 @@ use crate::{
 use alloc::vec::Vec;
 use core::fmt;
 use std::sync::mpsc;
+use wgt::Features;
 
 use crate::COPY_BUFFER_ALIGNMENT;
 
@@ -26,6 +27,11 @@ use crate::COPY_BUFFER_ALIGNMENT;
 pub struct StagingBelt {
     device: Device,
     chunk_size: BufferAddress,
+    /// User-specified [`BufferUsages`] with which the chunk buffers are created.
+    ///
+    /// [`new`](Self::new) guarantees that this always contains
+    /// [`MAP_WRITE`](BufferUsages::MAP_WRITE).
+    buffer_usages: BufferUsages,
     /// Chunks into which we are accumulating data to be transferred.
     active_chunks: Vec<Chunk>,
     /// Chunks that have scheduled transfers already; they are unmapped and some
@@ -51,11 +57,54 @@ impl StagingBelt {
     /// * 1-4 times less than the total amount of data uploaded per submission
     ///   (per [`StagingBelt::finish()`]); and
     /// * bigger is better, within these bounds.
+    ///
+    /// The buffers returned by this staging belt will have the buffer usages
+    /// [`MAP_READ | MAP_WRITE`](BufferUsages).
     pub fn new(device: Device, chunk_size: BufferAddress) -> Self {
+        Self::new_with_buffer_usages(device, chunk_size, BufferUsages::COPY_SRC)
+    }
+
+    /// Create a new staging belt.
+    ///
+    /// The `chunk_size` is the unit of internal buffer allocation; writes will be
+    /// sub-allocated within each chunk. Therefore, for optimal use of memory, the
+    /// chunk size should be:
+    ///
+    /// * larger than the largest single [`StagingBelt::write_buffer()`] operation;
+    /// * 1-4 times less than the total amount of data uploaded per submission
+    ///   (per [`StagingBelt::finish()`]); and
+    /// * bigger is better, within these bounds.
+    ///
+    /// `buffer_usages` specifies with which [`BufferUsages`] the staging buffers
+    /// will be created. [`MAP_WRITE`](BufferUsages::MAP_WRITE) will be added
+    /// automatically. The method will panic if the combination of usages is not
+    /// supported. Because [`MAP_WRITE`](BufferUsages::MAP_WRITE) is implied, only
+    /// [`COPY_SRC`](BufferUsages::COPY_SRC) can be used, except if
+    /// [`Features::MAPPABLE_PRIMARY_BUFFERS`] is enabled.
+    pub fn new_with_buffer_usages(
+        device: Device,
+        chunk_size: BufferAddress,
+        mut buffer_usages: BufferUsages,
+    ) -> Self {
         let (sender, receiver) = mpsc::channel();
+
+        // make sure anything other than MAP_WRITE | COPY_SRC is only allowed with MAPPABLE_PRIMARY_BUFFERS.
+        let extra_usages =
+            buffer_usages.difference(BufferUsages::MAP_WRITE | BufferUsages::COPY_SRC);
+        if !extra_usages.is_empty()
+            && !device
+                .features()
+                .contains(Features::MAPPABLE_PRIMARY_BUFFERS)
+        {
+            panic!("Only BufferUsages::COPY_SRC may be used when Features::MAPPABLE_PRIMARY_BUFFERS is not enabled. Specified buffer usages: {buffer_usages:?}");
+        }
+        // always set MAP_WRITE
+        buffer_usages.insert(BufferUsages::MAP_WRITE);
+
         StagingBelt {
             device,
             chunk_size,
+            buffer_usages,
             active_chunks: Vec::new(),
             closed_chunks: Vec::new(),
             free_chunks: Vec::new(),
@@ -117,7 +166,7 @@ impl StagingBelt {
     /// (The view must be dropped before [`StagingBelt::finish()`] is called.)
     ///
     /// You can then record your own GPU commands to perform with the slice,
-    /// such as copying it to a texture or executing a compute shader that reads it (whereas
+    /// such as copying it to a texture (whereas
     /// [`StagingBelt::write_buffer()`] can only write to other buffers).
     /// All commands involving this slice must be submitted after
     /// [`StagingBelt::finish()`] is called and before [`StagingBelt::recall()`] is called.
@@ -162,7 +211,7 @@ impl StagingBelt {
                     buffer: self.device.create_buffer(&BufferDescriptor {
                         label: Some("(wgpu internal) StagingBelt staging buffer"),
                         size: self.chunk_size.max(size.get()),
-                        usage: BufferUsages::MAP_WRITE | BufferUsages::COPY_SRC,
+                        usage: self.buffer_usages,
                         mapped_at_creation: true,
                     }),
                     offset: 0,
@@ -230,6 +279,7 @@ impl fmt::Debug for StagingBelt {
         let Self {
             device,
             chunk_size,
+            buffer_usages,
             active_chunks,
             closed_chunks,
             free_chunks,
@@ -239,6 +289,7 @@ impl fmt::Debug for StagingBelt {
         f.debug_struct("StagingBelt")
             .field("device", device)
             .field("chunk_size", chunk_size)
+            .field("buffer_usages", buffer_usages)
             .field("active_chunks", &active_chunks.len())
             .field("closed_chunks", &closed_chunks.len())
             .field("free_chunks", &free_chunks.len())