introduce FlushedStagingBuffer

Author: teoxoy

wgpu-core/src/device/queue.rs

@@ -17,8 +17,8 @@ use crate::{
     lock::RwLockWriteGuard,
     resource::{
         Buffer, BufferAccessError, BufferMapState, DestroyedBuffer, DestroyedResourceError,
-        DestroyedTexture, Labeled, ParentDevice, ResourceErrorIdent, StagingBuffer, Texture,
-        TextureInner, Trackable,
+        DestroyedTexture, FlushedStagingBuffer, Labeled, ParentDevice, ResourceErrorIdent,
+        StagingBuffer, Texture, TextureInner, Trackable,
     },
     resource_log,
     track::{self, Tracker, TrackerIndex},
@@ -136,7 +136,7 @@ pub struct WrappedSubmissionIndex {
 ///   submission, to be freed when it completes
 #[derive(Debug)]
 pub enum TempResource<A: HalApi> {
-    StagingBuffer(StagingBuffer<A>),
+    StagingBuffer(FlushedStagingBuffer<A>),
     DestroyedBuffer(DestroyedBuffer<A>),
     DestroyedTexture(DestroyedTexture<A>),
 }
@@ -256,7 +256,7 @@ impl<A: HalApi> PendingWrites<A> {
         self.temp_resources.push(resource);
     }
 
-    pub fn consume(&mut self, buffer: StagingBuffer<A>) {
+    pub fn consume(&mut self, buffer: FlushedStagingBuffer<A>) {
         self.temp_resources
             .push(TempResource::StagingBuffer(buffer));
     }
@@ -409,15 +409,15 @@ impl Global {
         let mut pending_writes = device.pending_writes.lock();
         let pending_writes = pending_writes.as_mut().unwrap();
 
-        unsafe {
+        let staging_buffer = unsafe {
             profiling::scope!("copy");
             ptr::copy_nonoverlapping(
                 data.as_ptr(),
                 staging_buffer_ptr.as_ptr(),
                 data_size.get() as usize,
             );
-            staging_buffer.flush();
-        }
+            staging_buffer.flush()
+        };
 
         let result = self.queue_write_staging_buffer_impl(
             &queue,
@@ -487,7 +487,7 @@ impl Global {
         // user. Platform validation requires that the staging buffer always
         // be freed, even if an error occurs. All paths from here must call
         // `device.pending_writes.consume`.
-        unsafe { staging_buffer.flush() };
+        let staging_buffer = unsafe { staging_buffer.flush() };
 
         let result = self.queue_write_staging_buffer_impl(
             &queue,
@@ -552,7 +552,7 @@ impl Global {
         queue: &Arc<Queue<A>>,
         device: &Arc<Device<A>>,
         pending_writes: &mut PendingWrites<A>,
-        staging_buffer: &StagingBuffer<A>,
+        staging_buffer: &FlushedStagingBuffer<A>,
         buffer_id: id::BufferId,
         buffer_offset: u64,
     ) -> Result<(), QueueWriteError> {
@@ -814,7 +814,7 @@ impl Global {
             }
         }
 
-        unsafe { staging_buffer.flush() };
+        let staging_buffer = unsafe { staging_buffer.flush() };
 
         let regions = (0..array_layer_count).map(|rel_array_layer| {
             let mut texture_base = dst_base.clone();

wgpu-core/src/resource.rs

@@ -672,7 +672,7 @@ impl<A: HalApi> Buffer<A> {
                 let mut pending_writes = device.pending_writes.lock();
                 let pending_writes = pending_writes.as_mut().unwrap();
 
-                unsafe { staging_buffer.flush() };
+                let staging_buffer = unsafe { staging_buffer.flush() };
 
                 self.use_at(device.active_submission_index.load(Ordering::Relaxed) + 1);
                 let region = wgt::BufferSize::new(self.size).map(|size| hal::BufferCopy {
@@ -853,7 +853,7 @@ impl<A: HalApi> Drop for DestroyedBuffer<A> {
 /// [`Device::pending_writes`]: crate::device::Device
 #[derive(Debug)]
 pub struct StagingBuffer<A: HalApi> {
-    raw: ManuallyDrop<A::Buffer>,
+    raw: A::Buffer,
     device: Arc<Device<A>>,
     pub(crate) size: wgt::BufferSize,
     is_coherent: bool,
@@ -873,11 +873,11 @@ impl<A: HalApi> StagingBuffer<A> {
             memory_flags: hal::MemoryFlags::TRANSIENT,
         };
 
-        let buffer = unsafe { device.raw().create_buffer(&stage_desc)? };
-        let mapping = unsafe { device.raw().map_buffer(&buffer, 0..size.get()) }?;
+        let raw = unsafe { device.raw().create_buffer(&stage_desc)? };
+        let mapping = unsafe { device.raw().map_buffer(&raw, 0..size.get()) }?;
 
         let staging_buffer = StagingBuffer {
-            raw: ManuallyDrop::new(buffer),
+            raw,
             device: device.clone(),
             size,
             is_coherent: mapping.is_coherent,
@@ -886,40 +886,52 @@ impl<A: HalApi> StagingBuffer<A> {
         Ok((staging_buffer, mapping.ptr))
     }
 
-    pub(crate) fn raw(&self) -> &A::Buffer {
-        &self.raw
-    }
-
-    pub(crate) unsafe fn flush(&self) {
+    pub(crate) fn flush(self) -> FlushedStagingBuffer<A> {
         use hal::Device;
         let device = self.device.raw();
         if !self.is_coherent {
-            unsafe { device.flush_mapped_ranges(self.raw(), iter::once(0..self.size.get())) };
+            unsafe { device.flush_mapped_ranges(&self.raw, iter::once(0..self.size.get())) };
+        }
+        unsafe { device.unmap_buffer(&self.raw) };
+
+        let StagingBuffer {
+            raw, device, size, ..
+        } = self;
+
+        FlushedStagingBuffer {
+            raw: ManuallyDrop::new(raw),
+            device,
+            size,
         }
-        unsafe { device.unmap_buffer(self.raw()) };
     }
 }
 
-impl<A: HalApi> Drop for StagingBuffer<A> {
+crate::impl_resource_type!(StagingBuffer);
+crate::impl_storage_item!(StagingBuffer);
+
+#[derive(Debug)]
+pub struct FlushedStagingBuffer<A: HalApi> {
+    raw: ManuallyDrop<A::Buffer>,
+    device: Arc<Device<A>>,
+    pub(crate) size: wgt::BufferSize,
+}
+
+impl<A: HalApi> FlushedStagingBuffer<A> {
+    pub(crate) fn raw(&self) -> &A::Buffer {
+        &self.raw
+    }
+}
+
+impl<A: HalApi> Drop for FlushedStagingBuffer<A> {
     fn drop(&mut self) {
         use hal::Device;
-        resource_log!("Destroy raw {}", self.error_ident());
+        resource_log!("Destroy raw StagingBuffer");
         // SAFETY: We are in the Drop impl and we don't use self.raw anymore after this point.
         let raw = unsafe { ManuallyDrop::take(&mut self.raw) };
         unsafe { self.device.raw().destroy_buffer(raw) };
     }
 }
 
-crate::impl_resource_type!(StagingBuffer);
-// TODO: add label
-impl<A: HalApi> Labeled for StagingBuffer<A> {
-    fn label(&self) -> &str {
-        ""
-    }
-}
-crate::impl_parent_device!(StagingBuffer);
-crate::impl_storage_item!(StagingBuffer);
-
 pub type TextureDescriptor<'a> = wgt::TextureDescriptor<Label<'a>, Vec<wgt::TextureFormat>>;
 
 #[derive(Debug)]