From 7fb2ec096e3b2074a0d6d0611e5c7181246fddb0 Mon Sep 17 00:00:00 2001
From: Nathan Adams <dinnerbone@dinnerbone.com>
Date: Fri, 23 Dec 2022 02:24:35 +0100
Subject: [PATCH] wgpu: Don't wait for the entire frame to be drawn if we're
 just capturing a segment of the render list

---
 render/wgpu/src/backend.rs |  6 +++---
 render/wgpu/src/target.rs  | 18 ++++++++++++------
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/render/wgpu/src/backend.rs b/render/wgpu/src/backend.rs
index ddf957fa2..dd123d5cb 100644
--- a/render/wgpu/src/backend.rs
+++ b/render/wgpu/src/backend.rs
@@ -114,7 +114,7 @@ impl WgpuRenderBackend<crate::target::TextureTarget> {
 
     pub fn capture_frame(&self, premultiplied_alpha: bool) -> Option<image::RgbaImage> {
         self.target
-            .capture(&self.descriptors.device, premultiplied_alpha)
+            .capture(&self.descriptors.device, premultiplied_alpha, None)
     }
 }
 
@@ -539,7 +539,7 @@ impl<T: RenderTarget + 'static> RenderBackend for WgpuRenderBackend<T> {
             commands,
             &mut self.frame_texture_pool,
         );
-        target.submit(
+        let index = target.submit(
             &self.descriptors.device,
             &self.descriptors.queue,
             command_buffers,
@@ -547,7 +547,7 @@ impl<T: RenderTarget + 'static> RenderBackend for WgpuRenderBackend<T> {
         );
 
         // Capture with premultiplied alpha, which is what we use for all textures
-        let image = target.capture(&self.descriptors.device, true);
+        let image = target.capture(&self.descriptors.device, true, Some(index));
 
         let image = image.map(|image| {
             Bitmap::new(
diff --git a/render/wgpu/src/target.rs b/render/wgpu/src/target.rs
index 6b8d5910f..db1b06654 100644
--- a/render/wgpu/src/target.rs
+++ b/render/wgpu/src/target.rs
@@ -29,7 +29,7 @@ pub trait RenderTarget: Debug + 'static {
         queue: &wgpu::Queue,
         command_buffers: I,
         frame: Self::Frame,
-    );
+    ) -> wgpu::SubmissionIndex;
 }
 
 #[derive(Debug)]
@@ -128,9 +128,10 @@ impl RenderTarget for SwapChainTarget {
         queue: &wgpu::Queue,
         command_buffers: I,
         frame: Self::Frame,
-    ) {
-        queue.submit(command_buffers);
+    ) -> wgpu::SubmissionIndex {
+        let index = queue.submit(command_buffers);
         frame.texture.present();
+        index
     }
 }
 
@@ -211,13 +212,18 @@ impl TextureTarget {
         &self,
         device: &wgpu::Device,
         premultiplied_alpha: bool,
+        index: Option<wgpu::SubmissionIndex>,
     ) -> Option<image::RgbaImage> {
         let (sender, receiver) = std::sync::mpsc::channel();
         let buffer_slice = self.buffer.slice(..);
         buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
             sender.send(result).unwrap();
         });
-        device.poll(wgpu::Maintain::Wait);
+        device.poll(
+            index
+                .map(wgpu::Maintain::WaitForSubmissionIndex)
+                .unwrap_or(wgpu::Maintain::Wait),
+        );
         let result = receiver.recv().unwrap();
         match result {
             Ok(()) => {
@@ -283,7 +289,7 @@ impl RenderTarget for TextureTarget {
         queue: &wgpu::Queue,
         command_buffers: I,
         _frame: Self::Frame,
-    ) {
+    ) -> wgpu::SubmissionIndex {
         let label = create_debug_label!("Render target transfer encoder");
         let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
             label: label.as_deref(),
@@ -305,6 +311,6 @@ impl RenderTarget for TextureTarget {
             },
             self.size,
         );
-        queue.submit(command_buffers.into_iter().chain(Some(encoder.finish())));
+        queue.submit(command_buffers.into_iter().chain(Some(encoder.finish())))
     }
 }