Add screenshot api (bevyengine#7163)

Fixes bevyengine#1207

# Objective

Right now, it's impossible to capture a screenshot of the entire window
without forking bevy. This is because
- The swapchain texture never has the COPY_SRC usage
- It can't be accessed without taking ownership of it
- Taking ownership of it breaks *a lot* of stuff

## Solution

- Introduce a dedicated api for taking a screenshot of a given bevy
window, and guarantee this screenshot will always match up with what
gets put on the screen.

---

## Changelog

- Added the `ScreenshotManager` resource with two functions,
`take_screenshot` and `save_screenshot_to_disk`

Author: TheRawMeatball

Diff for: Cargo.toml

@@ -1909,6 +1909,16 @@ description = "Illustrates how to customize the default window settings"
 category = "Window"
 wasm = true
 
+[[example]]
+name = "screenshot"
+path = "examples/window/screenshot.rs"
+
+[package.metadata.example.screenshot]
+name = "Screenshot"
+description = "Shows how to save screenshots to disk"
+category = "Window"
+wasm = false
+
 [[example]]
 name = "transparent_window"
 path = "examples/window/transparent_window.rs"

Diff for: crates/bevy_render/src/camera/camera.rs

@@ -423,7 +423,7 @@ impl NormalizedRenderTarget {
         match self {
             NormalizedRenderTarget::Window(window_ref) => windows
                 .get(&window_ref.entity())
-                .and_then(|window| window.swap_chain_texture.as_ref()),
+                .and_then(|window| window.swap_chain_texture_view.as_ref()),
             NormalizedRenderTarget::Image(image_handle) => {
                 images.get(image_handle).map(|image| &image.texture_view)
             }

Diff for: crates/bevy_render/src/camera/camera_driver_node.rs

@@ -52,7 +52,7 @@ impl Node for CameraDriverNode {
                 continue;
             }
 
-            let Some(swap_chain_texture) = &window.swap_chain_texture else {
+            let Some(swap_chain_texture) = &window.swap_chain_texture_view else {
                 continue;
             };
 

Diff for: crates/bevy_render/src/render_resource/texture.rs

@@ -51,31 +51,21 @@ define_atomic_id!(TextureViewId);
 render_resource_wrapper!(ErasedTextureView, wgpu::TextureView);
 render_resource_wrapper!(ErasedSurfaceTexture, wgpu::SurfaceTexture);
 
-/// This type combines wgpu's [`TextureView`](wgpu::TextureView) and
-/// [`SurfaceTexture`](wgpu::SurfaceTexture) into the same interface.
-#[derive(Clone, Debug)]
-pub enum TextureViewValue {
-    /// The value is an actual wgpu [`TextureView`](wgpu::TextureView).
-    TextureView(ErasedTextureView),
-
-    /// The value is a wgpu [`SurfaceTexture`](wgpu::SurfaceTexture), but dereferences to
-    /// a [`TextureView`](wgpu::TextureView).
-    SurfaceTexture {
-        // NOTE: The order of these fields is important because the view must be dropped before the
-        // frame is dropped
-        view: ErasedTextureView,
-        texture: ErasedSurfaceTexture,
-    },
-}
-
 /// Describes a [`Texture`] with its associated metadata required by a pipeline or [`BindGroup`](super::BindGroup).
-///
-/// May be converted from a [`TextureView`](wgpu::TextureView) or [`SurfaceTexture`](wgpu::SurfaceTexture)
-/// or dereferences to a wgpu [`TextureView`](wgpu::TextureView).
 #[derive(Clone, Debug)]
 pub struct TextureView {
     id: TextureViewId,
-    value: TextureViewValue,
+    value: ErasedTextureView,
+}
+
+pub struct SurfaceTexture {
+    value: ErasedSurfaceTexture,
+}
+
+impl SurfaceTexture {
+    pub fn try_unwrap(self) -> Option<wgpu::SurfaceTexture> {
+        self.value.try_unwrap()
+    }
 }
 
 impl TextureView {
@@ -84,34 +74,21 @@ impl TextureView {
     pub fn id(&self) -> TextureViewId {
         self.id
     }
-
-    /// Returns the [`SurfaceTexture`](wgpu::SurfaceTexture) of the texture view if it is of that type.
-    #[inline]
-    pub fn take_surface_texture(self) -> Option<wgpu::SurfaceTexture> {
-        match self.value {
-            TextureViewValue::TextureView(_) => None,
-            TextureViewValue::SurfaceTexture { texture, .. } => texture.try_unwrap(),
-        }
-    }
 }
 
 impl From<wgpu::TextureView> for TextureView {
     fn from(value: wgpu::TextureView) -> Self {
         TextureView {
             id: TextureViewId::new(),
-            value: TextureViewValue::TextureView(ErasedTextureView::new(value)),
+            value: ErasedTextureView::new(value),
         }
     }
 }
 
-impl From<wgpu::SurfaceTexture> for TextureView {
+impl From<wgpu::SurfaceTexture> for SurfaceTexture {
     fn from(value: wgpu::SurfaceTexture) -> Self {
-        let view = ErasedTextureView::new(value.texture.create_view(&Default::default()));
-        let texture = ErasedSurfaceTexture::new(value);
-
-        TextureView {
-            id: TextureViewId::new(),
-            value: TextureViewValue::SurfaceTexture { texture, view },
+        SurfaceTexture {
+            value: ErasedSurfaceTexture::new(value),
         }
     }
 }
@@ -121,10 +98,16 @@ impl Deref for TextureView {
 
     #[inline]
     fn deref(&self) -> &Self::Target {
-        match &self.value {
-            TextureViewValue::TextureView(value) => value,
-            TextureViewValue::SurfaceTexture { view, .. } => view,
-        }
+        &self.value
+    }
+}
+
+impl Deref for SurfaceTexture {
+    type Target = wgpu::SurfaceTexture;
+
+    #[inline]
+    fn deref(&self) -> &Self::Target {
+        &self.value
     }
 }
 

Diff for: crates/bevy_render/src/renderer/graph_runner.rs

@@ -57,9 +57,12 @@ impl RenderGraphRunner {
         render_device: RenderDevice,
         queue: &wgpu::Queue,
         world: &World,
+        finalizer: impl FnOnce(&mut wgpu::CommandEncoder),
     ) -> Result<(), RenderGraphRunnerError> {
         let mut render_context = RenderContext::new(render_device);
         Self::run_graph(graph, None, &mut render_context, world, &[], None)?;
+        finalizer(render_context.command_encoder());
+
         {
             #[cfg(feature = "trace")]
             let _span = info_span!("submit_graph_commands").entered();

Diff for: crates/bevy_render/src/renderer/mod.rs

@@ -35,6 +35,9 @@ pub fn render_system(world: &mut World) {
         render_device.clone(), // TODO: is this clone really necessary?
         &render_queue.0,
         world,
+        |encoder| {
+            crate::view::screenshot::submit_screenshot_commands(world, encoder);
+        },
     ) {
         error!("Error running render graph:");
         {
@@ -66,8 +69,8 @@ pub fn render_system(world: &mut World) {
 
         let mut windows = world.resource_mut::<ExtractedWindows>();
         for window in windows.values_mut() {
-            if let Some(texture_view) = window.swap_chain_texture.take() {
-                if let Some(surface_texture) = texture_view.take_surface_texture() {
+            if let Some(wrapped_texture) = window.swap_chain_texture.take() {
+                if let Some(surface_texture) = wrapped_texture.try_unwrap() {
                     surface_texture.present();
                 }
             }
@@ -81,6 +84,8 @@ pub fn render_system(world: &mut World) {
         );
     }
 
+    crate::view::screenshot::collect_screenshots(world);
+
     // update the time and send it to the app world
     let time_sender = world.resource::<TimeSender>();
     time_sender.0.try_send(Instant::now()).expect(

Diff for: crates/bevy_render/src/texture/image_texture_conversion.rs

@@ -174,6 +174,7 @@ impl Image {
     /// - `TextureFormat::R8Unorm`
     /// - `TextureFormat::Rg8Unorm`
     /// - `TextureFormat::Rgba8UnormSrgb`
+    /// - `TextureFormat::Bgra8UnormSrgb`
     ///
     /// To convert [`Image`] to a different format see: [`Image::convert`].
     pub fn try_into_dynamic(self) -> anyhow::Result<DynamicImage> {
@@ -196,6 +197,20 @@ impl Image {
                 self.data,
             )
             .map(DynamicImage::ImageRgba8),
+            // This format is commonly used as the format for the swapchain texture
+            // This conversion is added here to support screenshots
+            TextureFormat::Bgra8UnormSrgb => ImageBuffer::from_raw(
+                self.texture_descriptor.size.width,
+                self.texture_descriptor.size.height,
+                {
+                    let mut data = self.data;
+                    for bgra in data.chunks_exact_mut(4) {
+                        bgra.swap(0, 2);
+                    }
+                    data
+                },
+            )
+            .map(DynamicImage::ImageRgba8),
             // Throw and error if conversion isn't supported
             texture_format => {
                 return Err(anyhow!(