staging belt: allow creation with custom buffer usages

tests/tests/wgpu-validation/util.rs

@@ -39,3 +39,66 @@ fn staging_belt_random_test() {
         belt.recall();
     }
 }
+
+#[test]
+fn staging_belt_panics_with_invalid_buffer_usages() {
+    #[track_caller]
+    fn test_if_panics(usage: wgpu::BufferUsages) {
+        if let Err(panic) = 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);
+        }) {
+            // according to [1] the panic payload is either a `&str` or `String`
+            // [1]: https://doc.rust-lang.org/std/macro.panic.html
+
+            let message = if let Some(message) = panic.downcast_ref::<&str>() {
+                *message
+            } else if let Some(message) = panic.downcast_ref::<String>() {
+                message.as_str()
+            } else {
+                // don't know what this panic is, but it's not ours
+                std::panic::resume_unwind(panic);
+            };
+
+            let expected_message = format!("Only BufferUsages::COPY_SRC may be used when Features::MAPPABLE_PRIMARY_BUFFERS is not enabled. Specified buffer usages: {usage:?}");
+            if expected_message == message {
+                // panicked with the correct message
+            } else {
+                // This is not our panic (or the panic message was changed)
+                std::panic::resume_unwind(panic);
+            }
+        } else {
+            panic!("StagingBelt::new_with_buffer_usages should panic without MAPPABLE_PRIMARY_BUFFERS with usage={usage:?}");
+        }
+    }
+
+    for mut usage in wgpu::BufferUsages::all()
+        .difference(wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::MAP_WRITE)
+        .iter()
+    {
+        test_if_panics(usage);
+
+        usage.insert(wgpu::BufferUsages::MAP_WRITE);
+        test_if_panics(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,
+        wgpu::BufferUsages::COPY_SRC,
+    );
+    let _belt = wgpu::util::StagingBelt::new_with_buffer_usages(
+        device.clone(),
+        512,
+        wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::MAP_WRITE,
+    );
+    let _belt = wgpu::util::StagingBelt::new_with_buffer_usages(
+        device.clone(),
+        512,
+        wgpu::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,55 @@ 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.
+    #[track_caller]
+    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 +167,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 +212,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 +280,7 @@ impl fmt::Debug for StagingBelt {
         let Self {
             device,
             chunk_size,
+            buffer_usages,
             active_chunks,
             closed_chunks,
             free_chunks,
@@ -239,6 +290,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())