Add screenshot support for eframe web

crates/eframe/src/web/app_runner.rs

@@ -1,4 +1,5 @@
-use egui::TexturesDelta;
+use egui::{TexturesDelta, UserData, ViewportCommand};
+use std::mem;
 
 use crate::{epi, App};
 
@@ -16,6 +17,9 @@ pub struct AppRunner {
     last_save_time: f64,
     pub(crate) text_agent: TextAgent,
 
+    // If not empty, the painter should capture the next frame
+    screenshot_commands: Vec<UserData>,
+
     // Output for the last run:
     textures_delta: TexturesDelta,
     clipped_primitives: Option<Vec<egui::ClippedPrimitive>>,
@@ -36,7 +40,8 @@ impl AppRunner {
         app_creator: epi::AppCreator<'static>,
         text_agent: TextAgent,
     ) -> Result<Self, String> {
-        let painter = super::ActiveWebPainter::new(canvas, &web_options).await?;
+        let egui_ctx = egui::Context::default();
+        let painter = super::ActiveWebPainter::new(egui_ctx.clone(), canvas, &web_options).await?;
 
         let info = epi::IntegrationInfo {
             web_info: epi::WebInfo {
@@ -47,7 +52,6 @@ impl AppRunner {
         };
         let storage = LocalStorage::default();
 
-        let egui_ctx = egui::Context::default();
         egui_ctx.set_os(egui::os::OperatingSystem::from_user_agent(
             &super::user_agent().unwrap_or_default(),
         ));
@@ -110,6 +114,7 @@ impl AppRunner {
             needs_repaint,
             last_save_time: now_sec(),
             text_agent,
+            screenshot_commands: vec![],
             textures_delta: Default::default(),
             clipped_primitives: None,
         };
@@ -205,6 +210,8 @@ impl AppRunner {
     pub fn logic(&mut self) {
         // We sometimes miss blur/focus events due to the text agent, so let's just poll each frame:
         self.update_focus();
+        // We might have received a screenshot
+        self.painter.handle_screenshots(&mut self.input.raw.events);
 
         let canvas_size = super::canvas_size_in_points(self.canvas(), self.egui_ctx());
         let mut raw_input = self.input.new_frame(canvas_size);
@@ -225,12 +232,19 @@ impl AppRunner {
         if viewport_output.len() > 1 {
             log::warn!("Multiple viewports not yet supported on the web");
         }
-        for viewport_output in viewport_output.values() {
-            for command in &viewport_output.commands {
-                // TODO(emilk): handle some of the commands
-                log::warn!(
-                    "Unhandled egui viewport command: {command:?} - not implemented in web backend"
-                );
+        for (_viewport_id, viewport_output) in viewport_output {
+            for command in viewport_output.commands {
+                match command {
+                    ViewportCommand::Screenshot(user_data) => {
+                        self.screenshot_commands.push(user_data);
+                    }
+                    _ => {
+                        // TODO(emilk): handle some of the commands
+                        log::warn!(
+                            "Unhandled egui viewport command: {command:?} - not implemented in web backend"
+                        );
+                    }
+                }
             }
         }
 
@@ -250,6 +264,7 @@ impl AppRunner {
                 &clipped_primitives,
                 self.egui_ctx.pixels_per_point(),
                 &textures_delta,
+                mem::take(&mut self.screenshot_commands),
             ) {
                 log::error!("Failed to paint: {}", super::string_from_js_value(&err));
             }

crates/eframe/src/web/web_painter.rs

@@ -1,3 +1,4 @@
+use egui::{Event, UserData};
 use wasm_bindgen::JsValue;
 
 /// Renderer for a browser canvas.
@@ -16,14 +17,19 @@ pub(crate) trait WebPainter {
     fn max_texture_side(&self) -> usize;
 
     /// Update all internal textures and paint gui.
+    /// When `capture` isn't empty, the rendered screen should be captured.
+    /// Once the screenshot is ready, the screenshot should be returned via [`Self::handle_screenshots`].
     fn paint_and_update_textures(
         &mut self,
         clear_color: [f32; 4],
         clipped_primitives: &[egui::ClippedPrimitive],
         pixels_per_point: f32,
         textures_delta: &egui::TexturesDelta,
+        capture: Vec<UserData>,
     ) -> Result<(), JsValue>;
 
+    fn handle_screenshots(&mut self, events: &mut Vec<Event>);
+
     /// Destroy all resources.
     fn destroy(&mut self);
 }

crates/eframe/src/web/web_painter_glow.rs

@@ -1,24 +1,30 @@
+use egui::{Event, UserData, ViewportId};
+use egui_glow::glow;
+use std::sync::Arc;
 use wasm_bindgen::JsCast;
 use wasm_bindgen::JsValue;
 use web_sys::HtmlCanvasElement;
 
-use egui_glow::glow;
-
 use crate::{WebGlContextOption, WebOptions};
 
 use super::web_painter::WebPainter;
 
 pub(crate) struct WebPainterGlow {
     canvas: HtmlCanvasElement,
     painter: egui_glow::Painter,
+    screenshots: Vec<(egui::ColorImage, Vec<UserData>)>,
 }
 
 impl WebPainterGlow {
     pub fn gl(&self) -> &std::sync::Arc<glow::Context> {
         self.painter.gl()
     }
 
-    pub async fn new(canvas: HtmlCanvasElement, options: &WebOptions) -> Result<Self, String> {
+    pub async fn new(
+        _ctx: egui::Context,
+        canvas: HtmlCanvasElement,
+        options: &WebOptions,
+    ) -> Result<Self, String> {
         let (gl, shader_prefix) =
             init_glow_context_from_canvas(&canvas, options.webgl_context_option)?;
         #[allow(clippy::arc_with_non_send_sync)]
@@ -27,7 +33,11 @@ impl WebPainterGlow {
         let painter = egui_glow::Painter::new(gl, shader_prefix, None, options.dithering)
             .map_err(|err| format!("Error starting glow painter: {err}"))?;
 
-        Ok(Self { canvas, painter })
+        Ok(Self {
+            canvas,
+            painter,
+            screenshots: Vec::new(),
+        })
     }
 }
 
@@ -46,6 +56,7 @@ impl WebPainter for WebPainterGlow {
         clipped_primitives: &[egui::ClippedPrimitive],
         pixels_per_point: f32,
         textures_delta: &egui::TexturesDelta,
+        capture: Vec<UserData>,
     ) -> Result<(), JsValue> {
         let canvas_dimension = [self.canvas.width(), self.canvas.height()];
 
@@ -57,6 +68,11 @@ impl WebPainter for WebPainterGlow {
         self.painter
             .paint_primitives(canvas_dimension, pixels_per_point, clipped_primitives);
 
+        if !capture.is_empty() {
+            let image = self.painter.read_screen_rgba(canvas_dimension);
+            self.screenshots.push((image, capture));
+        }
+
         for &id in &textures_delta.free {
             self.painter.free_texture(id);
         }
@@ -67,6 +83,19 @@ impl WebPainter for WebPainterGlow {
     fn destroy(&mut self) {
         self.painter.destroy();
     }
+
+    fn handle_screenshots(&mut self, events: &mut Vec<Event>) {
+        for (image, data) in self.screenshots.drain(..) {
+            let image = Arc::new(image);
+            for data in data {
+                events.push(Event::Screenshot {
+                    viewport_id: ViewportId::default(),
+                    image: image.clone(),
+                    user_data: data,
+                });
+            }
+        }
+    }
 }
 
 /// Returns glow context and shader prefix.

crates/eframe/src/web/web_painter_wgpu.rs

@@ -1,14 +1,14 @@
-use std::sync::Arc;
+use std::sync::{mpsc, Arc};
 
+use super::web_painter::WebPainter;
+use crate::epaint::ColorImage;
+use crate::WebOptions;
+use egui::{Event, UserData, ViewportId};
+use egui_wgpu::capture::CaptureState;
+use egui_wgpu::{RenderState, SurfaceErrorAction, WgpuSetup};
 use wasm_bindgen::JsValue;
 use web_sys::HtmlCanvasElement;
 
-use egui_wgpu::{RenderState, SurfaceErrorAction, WgpuSetup};
-
-use crate::WebOptions;
-
-use super::web_painter::WebPainter;
-
 pub(crate) struct WebPainterWgpu {
     canvas: HtmlCanvasElement,
     surface: wgpu::Surface<'static>,
@@ -17,6 +17,10 @@ pub(crate) struct WebPainterWgpu {
     on_surface_error: Arc<dyn Fn(wgpu::SurfaceError) -> SurfaceErrorAction>,
     depth_format: Option<wgpu::TextureFormat>,
     depth_texture_view: Option<wgpu::TextureView>,
+    screen_capture_state: Option<CaptureState>,
+    capture_rx: mpsc::Receiver<(Vec<UserData>, ColorImage)>,
+    capture_tx: mpsc::Sender<(Vec<UserData>, ColorImage)>,
+    ctx: egui::Context,
 }
 
 impl WebPainterWgpu {
@@ -54,6 +58,7 @@ impl WebPainterWgpu {
 
     #[allow(unused)] // only used if `wgpu` is the only active feature.
     pub async fn new(
+        ctx: egui::Context,
         canvas: web_sys::HtmlCanvasElement,
         options: &WebOptions,
     ) -> Result<Self, String> {
@@ -119,17 +124,21 @@ impl WebPainterWgpu {
         .await
         .map_err(|err| err.to_string())?;
 
+        let default_configuration = surface
+            .get_default_config(&render_state.adapter, 0, 0) // Width/height is set later.
+            .ok_or("The surface isn't supported by this adapter")?;
+
         let surface_configuration = wgpu::SurfaceConfiguration {
             format: render_state.target_format,
             present_mode: options.wgpu_options.present_mode,
             view_formats: vec![render_state.target_format],
-            ..surface
-                .get_default_config(&render_state.adapter, 0, 0) // Width/height is set later.
-                .ok_or("The surface isn't supported by this adapter")?
+            ..default_configuration
         };
 
         log::debug!("wgpu painter initialized.");
 
+        let (capture_tx, capture_rx) = mpsc::channel();
+
         Ok(Self {
             canvas,
             render_state: Some(render_state),
@@ -138,6 +147,10 @@ impl WebPainterWgpu {
             depth_format,
             depth_texture_view: None,
             on_surface_error: options.wgpu_options.on_surface_error.clone(),
+            screen_capture_state: None,
+            capture_rx,
+            capture_tx,
+            ctx,
         })
     }
 }
@@ -159,7 +172,10 @@ impl WebPainter for WebPainterWgpu {
         clipped_primitives: &[egui::ClippedPrimitive],
         pixels_per_point: f32,
         textures_delta: &egui::TexturesDelta,
+        capture_data: Vec<UserData>,
     ) -> Result<(), JsValue> {
+        let capture = !capture_data.is_empty();
+
         let size_in_pixels = [self.canvas.width(), self.canvas.height()];
 
         let Some(render_state) = &self.render_state else {
@@ -220,7 +236,7 @@ impl WebPainter for WebPainterWgpu {
                 );
             }
 
-            let frame = match self.surface.get_current_texture() {
+            let output_frame = match self.surface.get_current_texture() {
                 Ok(frame) => frame,
                 Err(err) => match (*self.on_surface_error)(err) {
                     SurfaceErrorAction::RecreateSurface => {
@@ -236,12 +252,23 @@ impl WebPainter for WebPainterWgpu {
 
             {
                 let renderer = render_state.renderer.read();
-                let frame_view = frame
-                    .texture
-                    .create_view(&wgpu::TextureViewDescriptor::default());
+
+                let target_texture = if capture {
+                    let capture_state = self.screen_capture_state.get_or_insert_with(|| {
+                        CaptureState::new(&render_state.device, &output_frame.texture)
+                    });
+                    capture_state.update(&render_state.device, &output_frame.texture);
+
+                    &capture_state.texture
+                } else {
+                    &output_frame.texture
+                };
+                let target_view =
+                    target_texture.create_view(&wgpu::TextureViewDescriptor::default());
+
                 let render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                     color_attachments: &[Some(wgpu::RenderPassColorAttachment {
-                        view: &frame_view,
+                        view: &target_view,
                         resolve_target: None,
                         ops: wgpu::Operations {
                             load: wgpu::LoadOp::Clear(wgpu::Color {
@@ -280,7 +307,7 @@ impl WebPainter for WebPainterWgpu {
                 );
             }
 
-            Some(frame)
+            Some(output_frame)
         };
 
         {
@@ -296,12 +323,37 @@ impl WebPainter for WebPainterWgpu {
             .submit(user_cmd_bufs.into_iter().chain([encoder.finish()]));
 
         if let Some(frame) = frame {
+            if capture {
+                if let Some(capture_state) = &mut self.screen_capture_state {
+                    capture_state.read_screen_rgba(
+                        self.ctx.clone(),
+                        render_state,
+                        &frame,
+                        capture_data,
+                        self.capture_tx.clone(),
+                    );
+                }
+            };
+
             frame.present();
         }
 
         Ok(())
     }
 
+    fn handle_screenshots(&mut self, events: &mut Vec<Event>) {
+        for (user_data, screenshot) in self.capture_rx.try_iter() {
+            let screenshot = Arc::new(screenshot);
+            for data in user_data {
+                events.push(Event::Screenshot {
+                    viewport_id: ViewportId::default(),
+                    user_data: data,
+                    image: screenshot.clone(),
+                });
+            }
+        }
+    }
+
     fn destroy(&mut self) {
         self.render_state = None;
     }