Added offscreen rendering support for wgpu & tiny-skia exposed with the window::screenshot command.

6 files changed, 329 insertions, 11 deletions

diff --git a/wgpu/src/backend.rs b/wgpu/src/backend.rs
index b524c615..8f37f285 100644
--- a/wgpu/src/backend.rs
+++ b/wgpu/src/backend.rs
@@ -1,4 +1,3 @@
-use crate::core;
 use crate::core::{Color, Font, Point, Size};
 use crate::graphics::backend;
 use crate::graphics::color;
@@ -6,6 +5,7 @@ use crate::graphics::{Primitive, Transformation, Viewport};
 use crate::quad;
 use crate::text;
 use crate::triangle;
+use crate::{core, offscreen};
 use crate::{Layer, Settings};
 
 #[cfg(feature = "tracing")]
@@ -123,6 +123,65 @@ impl Backend {
         self.image_pipeline.end_frame();
     }
 
+    /// Performs an offscreen render pass. If the `format` selected by WGPU is not
+    /// `wgpu::TextureFormat::Rgba8UnormSrgb`, a conversion compute pipeline will run.
+    ///
+    /// Returns `None` if the `frame` is `Rgba8UnormSrgb`, else returns the newly
+    /// converted texture view in `Rgba8UnormSrgb`.
+    pub fn offscreen<T: AsRef<str>>(
+        &mut self,
+        device: &wgpu::Device,
+        queue: &wgpu::Queue,
+        encoder: &mut wgpu::CommandEncoder,
+        clear_color: Option<Color>,
+        frame: &wgpu::TextureView,
+        format: wgpu::TextureFormat,
+        primitives: &[Primitive],
+        viewport: &Viewport,
+        overlay_text: &[T],
+        texture_extent: wgpu::Extent3d,
+    ) -> Option<wgpu::Texture> {
+        #[cfg(feature = "tracing")]
+        let _ = info_span!("iced_wgpu::offscreen", "DRAW").entered();
+
+        self.present(
+            device,
+            queue,
+            encoder,
+            clear_color,
+            frame,
+            primitives,
+            viewport,
+            overlay_text,
+        );
+
+        if format != wgpu::TextureFormat::Rgba8UnormSrgb {
+            log::info!("Texture format is {format:?}; performing conversion to rgba8..");
+            let pipeline = offscreen::Pipeline::new(device);
+
+            let texture = device.create_texture(&wgpu::TextureDescriptor {
+                label: Some("iced_wgpu.offscreen.conversion.source_texture"),
+                size: texture_extent,
+                mip_level_count: 1,
+                sample_count: 1,
+                dimension: wgpu::TextureDimension::D2,
+                format: wgpu::TextureFormat::Rgba8Unorm,
+                usage: wgpu::TextureUsages::STORAGE_BINDING
+                    | wgpu::TextureUsages::COPY_SRC,
+                view_formats: &[],
+            });
+
+            let view =
+                texture.create_view(&wgpu::TextureViewDescriptor::default());
+
+            pipeline.convert(device, texture_extent, frame, &view, encoder);
+
+            return Some(texture);
+        }
+
+        None
+    }
+
     fn prepare_text(
         &mut self,
         device: &wgpu::Device,
diff --git a/wgpu/src/lib.rs b/wgpu/src/lib.rs
index 0a5726b5..827acb89 100644
--- a/wgpu/src/lib.rs
+++ b/wgpu/src/lib.rs
@@ -46,6 +46,7 @@ pub mod geometry;
 
 mod backend;
 mod buffer;
+mod offscreen;
 mod quad;
 mod text;
 mod triangle;
diff --git a/wgpu/src/offscreen.rs b/wgpu/src/offscreen.rs
new file mode 100644
index 00000000..29913d02
--- /dev/null
+++ b/wgpu/src/offscreen.rs
@@ -0,0 +1,102 @@
+use std::borrow::Cow;
+
+/// A simple compute pipeline to convert any texture to Rgba8UnormSrgb.
+#[derive(Debug)]
+pub struct Pipeline {
+    pipeline: wgpu::ComputePipeline,
+    layout: wgpu::BindGroupLayout,
+}
+
+impl Pipeline {
+    pub fn new(device: &wgpu::Device) -> Self {
+        let shader =
+            device.create_shader_module(wgpu::ShaderModuleDescriptor {
+                label: Some("iced_wgpu.offscreen.blit.shader"),
+                source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(include_str!(
+                    "shader/offscreen_blit.wgsl"
+                ))),
+            });
+
+        let bind_group_layout =
+            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
+                label: Some("iced_wgpu.offscreen.blit.bind_group_layout"),
+                entries: &[
+                    wgpu::BindGroupLayoutEntry {
+                        binding: 0,
+                        visibility: wgpu::ShaderStages::COMPUTE,
+                        ty: wgpu::BindingType::Texture {
+                            sample_type: wgpu::TextureSampleType::Float {
+                                filterable: false,
+                            },
+                            view_dimension: wgpu::TextureViewDimension::D2,
+                            multisampled: false,
+                        },
+                        count: None,
+                    },
+                    wgpu::BindGroupLayoutEntry {
+                        binding: 1,
+                        visibility: wgpu::ShaderStages::COMPUTE,
+                        ty: wgpu::BindingType::StorageTexture {
+                            access: wgpu::StorageTextureAccess::WriteOnly,
+                            format: wgpu::TextureFormat::Rgba8Unorm,
+                            view_dimension: wgpu::TextureViewDimension::D2,
+                        },
+                        count: None,
+                    },
+                ],
+            });
+
+        let pipeline_layout =
+            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
+                label: Some("iced_wgpu.offscreen.blit.pipeline_layout"),
+                bind_group_layouts: &[&bind_group_layout],
+                push_constant_ranges: &[],
+            });
+
+        let pipeline =
+            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
+                label: Some("iced_wgpu.offscreen.blit.pipeline"),
+                layout: Some(&pipeline_layout),
+                module: &shader,
+                entry_point: "main",
+            });
+
+        Self {
+            pipeline,
+            layout: bind_group_layout,
+        }
+    }
+
+    pub fn convert(
+        &self,
+        device: &wgpu::Device,
+        extent: wgpu::Extent3d,
+        frame: &wgpu::TextureView,
+        view: &wgpu::TextureView,
+        encoder: &mut wgpu::CommandEncoder,
+    ) {
+        let bind = device.create_bind_group(&wgpu::BindGroupDescriptor {
+            label: Some("iced_wgpu.offscreen.blit.bind_group"),
+            layout: &self.layout,
+            entries: &[
+                wgpu::BindGroupEntry {
+                    binding: 0,
+                    resource: wgpu::BindingResource::TextureView(frame),
+                },
+                wgpu::BindGroupEntry {
+                    binding: 1,
+                    resource: wgpu::BindingResource::TextureView(view),
+                },
+            ],
+        });
+
+        let mut compute_pass =
+            encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
+                label: Some("iced_wgpu.offscreen.blit.compute_pass"),
+            });
+
+        compute_pass.set_pipeline(&self.pipeline);
+        compute_pass.set_bind_group(0, &bind, &[]);
+        compute_pass.dispatch_workgroups(extent.width, extent.height, 1);
+    }
+}
diff --git a/wgpu/src/shader/offscreen_blit.wgsl b/wgpu/src/shader/offscreen_blit.wgsl
new file mode 100644
index 00000000..9c764c36
--- /dev/null
+++ b/wgpu/src/shader/offscreen_blit.wgsl
@@ -0,0 +1,22 @@
+@group(0) @binding(0) var u_texture: texture_2d<f32>;
+@group(0) @binding(1) var out_texture: texture_storage_2d<rgba8unorm, write>;
+
+fn srgb(color: f32) -> f32 {
+    if (color <= 0.0031308) {
+        return 12.92 * color;
+    } else {
+        return (1.055 * (pow(color, (1.0/2.4)))) - 0.055;
+    }
+}
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    // texture coord must be i32 due to a naga bug:
+    // https://github.com/gfx-rs/naga/issues/1997
+    let coords = vec2(i32(id.x), i32(id.y));
+
+    let src: vec4<f32> = textureLoad(u_texture, coords, 0);
+    let srgb_color: vec4<f32> = vec4(srgb(src.x), srgb(src.y), srgb(src.z), src.w);
+
+    textureStore(out_texture, coords, srgb_color);
+}
diff --git a/wgpu/src/triangle/msaa.rs b/wgpu/src/triangle/msaa.rs
index 4afbdb32..320b5b12 100644
--- a/wgpu/src/triangle/msaa.rs
+++ b/wgpu/src/triangle/msaa.rs
@@ -16,15 +16,8 @@ impl Blit {
         format: wgpu::TextureFormat,
         antialiasing: graphics::Antialiasing,
     ) -> Blit {
-        let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
-            address_mode_u: wgpu::AddressMode::ClampToEdge,
-            address_mode_v: wgpu::AddressMode::ClampToEdge,
-            address_mode_w: wgpu::AddressMode::ClampToEdge,
-            mag_filter: wgpu::FilterMode::Nearest,
-            min_filter: wgpu::FilterMode::Nearest,
-            mipmap_filter: wgpu::FilterMode::Nearest,
-            ..Default::default()
-        });
+        let sampler =
+            device.create_sampler(&wgpu::SamplerDescriptor::default());
 
         let constant_layout =
             device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
diff --git a/wgpu/src/window/compositor.rs b/wgpu/src/window/compositor.rs
index 2eaafde0..43c3dce5 100644
--- a/wgpu/src/window/compositor.rs
+++ b/wgpu/src/window/compositor.rs
@@ -1,5 +1,5 @@
 //! Connect a window with a renderer.
-use crate::core::Color;
+use crate::core::{Color, Size};
 use crate::graphics;
 use crate::graphics::color;
 use crate::graphics::compositor;
@@ -283,4 +283,145 @@ impl<Theme> graphics::Compositor for Compositor<Theme> {
             )
         })
     }
+
+    fn screenshot<T: AsRef<str>>(
+        &mut self,
+        renderer: &mut Self::Renderer,
+        _surface: &mut Self::Surface,
+        viewport: &Viewport,
+        background_color: Color,
+        overlay: &[T],
+    ) -> Vec<u8> {
+        renderer.with_primitives(|backend, primitives| {
+            screenshot(
+                self,
+                backend,
+                primitives,
+                viewport,
+                background_color,
+                overlay,
+            )
+        })
+    }
+}
+
+/// Renders the current surface to an offscreen buffer.
+///
+/// Returns RGBA bytes of the texture data.
+pub fn screenshot<Theme, T: AsRef<str>>(
+    compositor: &Compositor<Theme>,
+    backend: &mut Backend,
+    primitives: &[Primitive],
+    viewport: &Viewport,
+    background_color: Color,
+    overlay: &[T],
+) -> Vec<u8> {
+    let mut encoder = compositor.device.create_command_encoder(
+        &wgpu::CommandEncoderDescriptor {
+            label: Some("iced_wgpu.offscreen.encoder"),
+        },
+    );
+
+    let dimensions = BufferDimensions::new(viewport.physical_size());
+
+    let texture_extent = wgpu::Extent3d {
+        width: dimensions.width,
+        height: dimensions.height,
+        depth_or_array_layers: 1,
+    };
+
+    let texture = compositor.device.create_texture(&wgpu::TextureDescriptor {
+        label: Some("iced_wgpu.offscreen.source_texture"),
+        size: texture_extent,
+        mip_level_count: 1,
+        sample_count: 1,
+        dimension: wgpu::TextureDimension::D2,
+        format: compositor.format,
+        usage: wgpu::TextureUsages::RENDER_ATTACHMENT
+            | wgpu::TextureUsages::COPY_SRC
+            | wgpu::TextureUsages::TEXTURE_BINDING,
+        view_formats: &[],
+    });
+
+    let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
+
+    let rgba_texture = backend.offscreen(
+        &compositor.device,
+        &compositor.queue,
+        &mut encoder,
+        Some(background_color),
+        &view,
+        compositor.format,
+        primitives,
+        viewport,
+        overlay,
+        texture_extent,
+    );
+
+    let output_buffer =
+        compositor.device.create_buffer(&wgpu::BufferDescriptor {
+            label: Some("iced_wgpu.offscreen.output_texture_buffer"),
+            size: (dimensions.padded_bytes_per_row * dimensions.height as usize)
+                as u64,
+            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
+            mapped_at_creation: false,
+        });
+
+    encoder.copy_texture_to_buffer(
+        rgba_texture.unwrap_or(texture).as_image_copy(),
+        wgpu::ImageCopyBuffer {
+            buffer: &output_buffer,
+            layout: wgpu::ImageDataLayout {
+                offset: 0,
+                bytes_per_row: Some(dimensions.padded_bytes_per_row as u32),
+                rows_per_image: None,
+            },
+        },
+        texture_extent,
+    );
+
+    let index = compositor.queue.submit(Some(encoder.finish()));
+
+    let slice = output_buffer.slice(..);
+    slice.map_async(wgpu::MapMode::Read, |_| {});
+
+    let _ = compositor
+        .device
+        .poll(wgpu::Maintain::WaitForSubmissionIndex(index));
+
+    let mapped_buffer = slice.get_mapped_range();
+
+    mapped_buffer.chunks(dimensions.padded_bytes_per_row).fold(
+        vec![],
+        |mut acc, row| {
+            acc.extend(&row[..dimensions.unpadded_bytes_per_row]);
+            acc
+        },
+    )
+}
+
+#[derive(Clone, Copy, Debug)]
+struct BufferDimensions {
+    width: u32,
+    height: u32,
+    unpadded_bytes_per_row: usize,
+    padded_bytes_per_row: usize,
+}
+
+impl BufferDimensions {
+    fn new(size: Size<u32>) -> Self {
+        let unpadded_bytes_per_row = size.width as usize * 4; //slice of buffer per row; always RGBA
+        let alignment = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT as usize; //256
+        let padded_bytes_per_row_padding =
+            (alignment - unpadded_bytes_per_row % alignment) % alignment;
+        let padded_bytes_per_row =
+            unpadded_bytes_per_row + padded_bytes_per_row_padding;
+
+        Self {
+            width: size.width,
+            height: size.height,
+            unpadded_bytes_per_row,
+            padded_bytes_per_row,
+        }
+    }
 }