eframe: capture a screenshot using Frame::request_screenshot

Co-authored-by: Emil Ernerfeldt <emil.ernerfeldt@gmail.com>

crates/egui-wgpu/src/winit.rs

@@ -1,6 +1,6 @@
 use std::sync::Arc;
 
-use epaint::mutex::RwLock;
+use epaint::{self, mutex::RwLock};
 
 use tracing::error;
 
@@ -13,6 +13,65 @@ struct SurfaceState {
     height: u32,
 }
 
+/// A texture and a buffer for reading the rendered frame back to the cpu.
+/// The texture is required since [`wgpu::TextureUsages::COPY_DST`] is not an allowed
+/// flag for the surface texture on all platforms. This means that anytime we want to
+/// capture the frame, we first render it to this texture, and then we can copy it to
+/// both the surface texture and the buffer, from where we can pull it back to the cpu.
+struct CaptureState {
+    texture: wgpu::Texture,
+    buffer: wgpu::Buffer,
+    padding: BufferPadding,
+}
+
+impl CaptureState {
+    fn new(device: &Arc<wgpu::Device>, surface_texture: &wgpu::Texture) -> Self {
+        let texture = device.create_texture(&wgpu::TextureDescriptor {
+            label: Some("egui_screen_capture_texture"),
+            size: surface_texture.size(),
+            mip_level_count: surface_texture.mip_level_count(),
+            sample_count: surface_texture.sample_count(),
+            dimension: surface_texture.dimension(),
+            format: surface_texture.format(),
+            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
+            view_formats: &[],
+        });
+
+        let padding = BufferPadding::new(surface_texture.width());
+
+        let buffer = device.create_buffer(&wgpu::BufferDescriptor {
+            label: Some("egui_screen_capture_buffer"),
+            size: (padding.padded_bytes_per_row * texture.height()) as u64,
+            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
+            mapped_at_creation: false,
+        });
+
+        Self {
+            texture,
+            buffer,
+            padding,
+        }
+    }
+}
+
+struct BufferPadding {
+    unpadded_bytes_per_row: u32,
+    padded_bytes_per_row: u32,
+}
+
+impl BufferPadding {
+    fn new(width: u32) -> Self {
+        let bytes_per_pixel = std::mem::size_of::<u32>() as u32;
+        let unpadded_bytes_per_row = width * bytes_per_pixel;
+        let padded_bytes_per_row =
+            wgpu::util::align_to(unpadded_bytes_per_row, wgpu::COPY_BYTES_PER_ROW_ALIGNMENT);
+        Self {
+            unpadded_bytes_per_row,
+            padded_bytes_per_row,
+        }
+    }
+}
+
 /// Everything you need to paint egui with [`wgpu`] on [`winit`].
 ///
 /// Alternatively you can use [`crate::renderer`] directly.
@@ -22,6 +81,7 @@ pub struct Painter {
     support_transparent_backbuffer: bool,
     depth_format: Option<wgpu::TextureFormat>,
     depth_texture_view: Option<wgpu::TextureView>,
+    screen_capture_state: Option<CaptureState>,
 
     instance: wgpu::Instance,
     adapter: Option<wgpu::Adapter>,
@@ -59,6 +119,7 @@ impl Painter {
             support_transparent_backbuffer,
             depth_format: (depth_bits > 0).then_some(wgpu::TextureFormat::Depth32Float),
             depth_texture_view: None,
+            screen_capture_state: None,
 
             instance,
             adapter: None,
@@ -136,7 +197,7 @@ impl Painter {
         surface_state.surface.configure(
             &render_state.device,
             &wgpu::SurfaceConfiguration {
-                usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
+                usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_DST,
                 format: render_state.target_format,
                 width: surface_state.width,
                 height: surface_state.height,
@@ -274,23 +335,118 @@ impl Painter {
         }
     }
 
+    // CaptureState only needs to be updated when the size of the two textures don't match and we want to
+    // capture a frame
+    fn update_capture_state(
+        screen_capture_state: &mut Option<CaptureState>,
+        surface_texture: &wgpu::SurfaceTexture,
+        render_state: &RenderState,
+    ) {
+        let surface_texture = &surface_texture.texture;
+        match screen_capture_state {
+            Some(capture_state) => {
+                if capture_state.texture.size() != surface_texture.size() {
+                    *capture_state = CaptureState::new(&render_state.device, surface_texture);
+                }
+            }
+            None => {
+                *screen_capture_state =
+                    Some(CaptureState::new(&render_state.device, surface_texture));
+            }
+        }
+    }
+
+    // Handles copying from the CaptureState texture to the surface texture and the cpu
+    fn read_screen_rgba(
+        screen_capture_state: &CaptureState,
+        render_state: &RenderState,
+        output_frame: &wgpu::SurfaceTexture,
+    ) -> Option<epaint::ColorImage> {
+        let CaptureState {
+            texture: tex,
+            buffer,
+            padding,
+        } = screen_capture_state;
+
+        let device = &render_state.device;
+        let queue = &render_state.queue;
+
+        let tex_extent = tex.size();
+
+        let mut encoder = device.create_command_encoder(&Default::default());
+        encoder.copy_texture_to_buffer(
+            tex.as_image_copy(),
+            wgpu::ImageCopyBuffer {
+                buffer,
+                layout: wgpu::ImageDataLayout {
+                    offset: 0,
+                    bytes_per_row: Some(std::num::NonZeroU32::new(padding.padded_bytes_per_row)?),
+                    rows_per_image: None,
+                },
+            },
+            tex_extent,
+        );
+
+        encoder.copy_texture_to_texture(
+            tex.as_image_copy(),
+            output_frame.texture.as_image_copy(),
+            tex.size(),
+        );
+
+        let id = queue.submit(Some(encoder.finish()));
+        let buffer_slice = buffer.slice(..);
+        let (sender, receiver) = std::sync::mpsc::channel();
+        buffer_slice.map_async(wgpu::MapMode::Read, move |v| {
+            drop(sender.send(v));
+        });
+        device.poll(wgpu::Maintain::WaitForSubmissionIndex(id));
+        receiver.recv().ok()?.ok()?;
+
+        let to_rgba = match tex.format() {
+            wgpu::TextureFormat::Rgba8Unorm => [0, 1, 2, 3],
+            wgpu::TextureFormat::Bgra8Unorm => [2, 1, 0, 3],
+            _ => {
+                tracing::error!("Screen can't be captured unless the surface format is Rgba8Unorm or Bgra8Unorm. Current surface format is {:?}", tex.format());
+                return None;
+            }
+        };
+
+        let mut pixels = Vec::with_capacity((tex.width() * tex.height()) as usize);
+        for padded_row in buffer_slice
+            .get_mapped_range()
+            .chunks(padding.padded_bytes_per_row as usize)
+        {
+            let row = &padded_row[..padding.unpadded_bytes_per_row as usize];
+            for color in row.chunks(4) {
+                pixels.push(epaint::Color32::from_rgba_premultiplied(
+                    color[to_rgba[0]],
+                    color[to_rgba[1]],
+                    color[to_rgba[2]],
+                    color[to_rgba[3]],
+                ));
+            }
+        }
+        buffer.unmap();
+
+        Some(epaint::ColorImage {
+            size: [tex.width() as usize, tex.height() as usize],
+            pixels,
+        })
+    }
+
+    // Returns a vector with the frame's pixel data if it was requested.
     pub fn paint_and_update_textures(
         &mut self,
         pixels_per_point: f32,
         clear_color: [f32; 4],
         clipped_primitives: &[epaint::ClippedPrimitive],
         textures_delta: &epaint::textures::TexturesDelta,
-    ) {
+        capture: bool,
+    ) -> Option<epaint::ColorImage> {
         crate::profile_function!();
 
-        let render_state = match self.render_state.as_mut() {
-            Some(rs) => rs,
-            None => return,
-        };
-        let surface_state = match self.surface_state.as_ref() {
-            Some(rs) => rs,
-            None => return,
-        };
+        let render_state = self.render_state.as_mut()?;
+        let surface_state = self.surface_state.as_ref()?;
 
         let output_frame = {
             crate::profile_scope!("get_current_texture");
@@ -308,10 +464,10 @@ impl Painter {
                         render_state,
                         self.configuration.present_mode,
                     );
-                    return;
+                    return None;
                 }
                 SurfaceErrorAction::SkipFrame => {
-                    return;
+                    return None;
                 }
             },
         };
@@ -351,9 +507,21 @@ impl Painter {
 
         {
             let renderer = render_state.renderer.read();
-            let frame_view = output_frame
-                .texture
-                .create_view(&wgpu::TextureViewDescriptor::default());
+            let frame_view = if capture {
+                Self::update_capture_state(
+                    &mut self.screen_capture_state,
+                    &output_frame,
+                    render_state,
+                );
+                self.screen_capture_state
+                    .as_ref()?
+                    .texture
+                    .create_view(&wgpu::TextureViewDescriptor::default())
+            } else {
+                output_frame
+                    .texture
+                    .create_view(&wgpu::TextureViewDescriptor::default())
+            };
             let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                 color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                     view: &frame_view,
@@ -404,11 +572,18 @@ impl Painter {
                 .submit(user_cmd_bufs.into_iter().chain(std::iter::once(encoded)));
         };
 
+        let screenshot = if capture {
+            let screen_capture_state = self.screen_capture_state.as_ref()?;
+            Self::read_screen_rgba(screen_capture_state, render_state, &output_frame)
+        } else {
+            None
+        };
         // Redraw egui
         {
             crate::profile_scope!("present");
             output_frame.present();
         }
+        screenshot
     }
 
     #[allow(clippy::unused_self)]