Introduce ExternalStagingBuffer for staging buffer uses outside of wgpu-core

wgpu-core/src/device/queue.rs

@@ -17,8 +17,8 @@ use crate::{
     lock::RwLockWriteGuard,
     resource::{
         Buffer, BufferAccessError, BufferMapState, DestroyedBuffer, DestroyedResourceError,
-        DestroyedTexture, FlushedStagingBuffer, Labeled, ParentDevice, ResourceErrorIdent,
-        StagingBuffer, Texture, TextureInner, Trackable,
+        DestroyedTexture, ExternalStagingBuffer, FlushedStagingBuffer, Labeled, ParentDevice,
+        ResourceErrorIdent, StagingBuffer, Texture, TextureInner, Trackable,
     },
     resource_log,
     track::{self, Tracker, TrackerIndex},
@@ -438,8 +438,8 @@ impl Global {
         &self,
         queue_id: QueueId,
         buffer_size: wgt::BufferSize,
-        id_in: Option<id::StagingBufferId>,
-    ) -> Result<(id::StagingBufferId, NonNull<u8>), QueueWriteError> {
+        id_in: Option<id::ExternalStagingBufferId>,
+    ) -> Result<(id::ExternalStagingBufferId, NonNull<u8>), QueueWriteError> {
         profiling::scope!("Queue::create_staging_buffer");
         let hub = A::hub(self);
 
@@ -450,11 +450,10 @@ impl Global {
 
         let device = &queue.device;
 
-        let staging_buffer = StagingBuffer::new(device, buffer_size)?;
-        let ptr = unsafe { staging_buffer.ptr() };
+        let (staging_buffer, ptr) = ExternalStagingBuffer::new(device, buffer_size)?;
 
         let fid = hub.staging_buffers.prepare(id_in);
-        let id = fid.assign(Arc::new(staging_buffer));
+        let id = fid.assign(staging_buffer);
         resource_log!("Queue::create_staging_buffer {id:?}");
 
         Ok((id, ptr))
@@ -465,7 +464,7 @@ impl Global {
         queue_id: QueueId,
         buffer_id: id::BufferId,
         buffer_offset: wgt::BufferAddress,
-        staging_buffer_id: id::StagingBufferId,
+        staging_buffer_id: id::ExternalStagingBufferId,
     ) -> Result<(), QueueWriteError> {
         profiling::scope!("Queue::write_staging_buffer");
         let hub = A::hub(self);
@@ -483,6 +482,10 @@ impl Global {
             .and_then(Arc::into_inner)
             .ok_or_else(|| QueueWriteError::Transfer(TransferError::InvalidBufferId(buffer_id)))?;
 
+        // TODO: bubble up safety requirements, they currently hold but we should make functions
+        // unsafe and add safety comments
+        let staging_buffer = unsafe { staging_buffer.into_inner() };
+
         let mut pending_writes = device.pending_writes.lock();
 
         // At this point, we have taken ownership of the staging_buffer from the

wgpu-core/src/hub.rs

@@ -112,7 +112,7 @@ use crate::{
     instance::{Adapter, Surface},
     pipeline::{ComputePipeline, PipelineCache, RenderPipeline, ShaderModule},
     registry::{Registry, RegistryReport},
-    resource::{Buffer, QuerySet, Sampler, StagingBuffer, Texture, TextureView},
+    resource::{Buffer, ExternalStagingBuffer, QuerySet, Sampler, Texture, TextureView},
     storage::{Element, Storage},
 };
 use std::fmt::Debug;
@@ -184,7 +184,7 @@ pub struct Hub<A: HalApi> {
     pub(crate) pipeline_caches: Registry<PipelineCache<A>>,
     pub(crate) query_sets: Registry<QuerySet<A>>,
     pub(crate) buffers: Registry<Buffer<A>>,
-    pub(crate) staging_buffers: Registry<StagingBuffer<A>>,
+    pub(crate) staging_buffers: Registry<ExternalStagingBuffer<A>>,
     pub(crate) textures: Registry<Texture<A>>,
     pub(crate) texture_views: Registry<TextureView<A>>,
     pub(crate) samplers: Registry<Sampler<A>>,

wgpu-core/src/id.rs

@@ -291,7 +291,7 @@ ids! {
     pub type DeviceId Device;
     pub type QueueId Queue;
     pub type BufferId Buffer;
-    pub type StagingBufferId StagingBuffer;
+    pub type ExternalStagingBufferId ExternalStagingBuffer;
     pub type TextureViewId TextureView;
     pub type TextureId Texture;
     pub type SamplerId Sampler;

wgpu-core/src/resource.rs

@@ -790,6 +790,37 @@ impl<A: HalApi> Drop for DestroyedBuffer<A> {
     }
 }
 
+/// A temporary buffer, that can be created with [`queue_create_staging_buffer`]
+/// and used with [`queue_write_staging_buffer`].
+///
+/// [`queue_create_staging_buffer`]: Global::queue_create_staging_buffer
+/// [`queue_write_staging_buffer`]: Global::queue_write_staging_buffer
+#[derive(Debug)]
+pub struct ExternalStagingBuffer<A: HalApi> {
+    inner: StagingBuffer<A>,
+}
+
+impl<A: HalApi> ExternalStagingBuffer<A> {
+    pub(crate) fn new(
+        device: &Arc<Device<A>>,
+        size: wgt::BufferSize,
+    ) -> Result<(Arc<Self>, NonNull<u8>), DeviceError> {
+        let inner = StagingBuffer::new(device, size)?;
+        let ptr = unsafe { inner.ptr() };
+        let staging_buffer = Self { inner };
+        let staging_buffer = Arc::new(staging_buffer);
+        Ok((staging_buffer, ptr))
+    }
+
+    /// SAFETY: You must have stopped using the pointer returned by [`Self::new`].
+    pub(crate) unsafe fn into_inner(self) -> StagingBuffer<A> {
+        self.inner
+    }
+}
+
+crate::impl_resource_type!(ExternalStagingBuffer);
+crate::impl_storage_item!(ExternalStagingBuffer);
+
 #[cfg(send_sync)]
 unsafe impl<A: HalApi> Send for StagingBuffer<A> {}
 #[cfg(send_sync)]
@@ -801,17 +832,14 @@ unsafe impl<A: HalApi> Sync for StagingBuffer<A> {}
 /// is always created mapped, and the command that uses it destroys the buffer
 /// when it is done.
 ///
-/// [`StagingBuffer`]s can be created with [`queue_create_staging_buffer`] and
-/// used with [`queue_write_staging_buffer`]. They are also used internally by
-/// operations like [`queue_write_texture`] that need to upload data to the GPU,
-/// but that don't belong to any particular wgpu command buffer.
+/// They are used internally by operations like [`queue_write_texture`]
+/// that need to upload data to the GPU, but that don't belong to any
+/// particular wgpu command buffer.
 ///
 /// Used `StagingBuffer`s are accumulated in [`Device::pending_writes`], to be
 /// freed once their associated operation's queue submission has finished
 /// execution.
 ///
-/// [`queue_create_staging_buffer`]: Global::queue_create_staging_buffer
-/// [`queue_write_staging_buffer`]: Global::queue_write_staging_buffer
 /// [`queue_write_texture`]: Global::queue_write_texture
 /// [`Device::pending_writes`]: crate::device::Device
 #[derive(Debug)]
@@ -913,9 +941,6 @@ impl<A: HalApi> StagingBuffer<A> {
     }
 }
 
-crate::impl_resource_type!(StagingBuffer);
-crate::impl_storage_item!(StagingBuffer);
-
 #[derive(Debug)]
 pub struct FlushedStagingBuffer<A: HalApi> {
     raw: ManuallyDrop<A::Buffer>,

wgpu/src/backend/wgpu_core.rs

@@ -3435,7 +3435,7 @@ impl From<CreateShaderModuleError> for CompilationInfo {
 
 #[derive(Debug)]
 pub struct QueueWriteBuffer {
-    buffer_id: wgc::id::StagingBufferId,
+    buffer_id: wgc::id::ExternalStagingBufferId,
     mapping: BufferMappedRange,
 }