From b1d2551e7ee529d7d0427173d8a0ec6f49f5b976 Mon Sep 17 00:00:00 2001
From: Andreas Reich <andreas@rerun.io>
Date: Mon, 16 Dec 2024 15:03:01 +0100
Subject: [PATCH] Make frame delay on screenshots consistently one frame on web
 as well (#5482)

Native is already delayed by a frame because it calls
`handle_viewport_output` -> `egui_winit::process_viewport_commands`
after drawing. On web however, we process input including viewport
commands separately from drawing.
This adds an arbitrary frame delay mechanism for web and then uses this
with 1 frame delay always
---
 crates/eframe/src/web/app_runner.rs | 28 ++++++++++++++++++++++------
 crates/egui/src/viewport.rs         |  2 +-
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/crates/eframe/src/web/app_runner.rs b/crates/eframe/src/web/app_runner.rs
index d8abd3d4bd7..13ad762874a 100644
--- a/crates/eframe/src/web/app_runner.rs
+++ b/crates/eframe/src/web/app_runner.rs
@@ -1,5 +1,4 @@
 use egui::{TexturesDelta, UserData, ViewportCommand};
-use std::mem;
 
 use crate::{epi, App};
 
@@ -17,8 +16,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>,
+    // If not empty, the painter should capture n frames from now.
+    // zero means capture the exact next frame.
+    screenshot_commands_with_frame_delay: Vec<(UserData, usize)>,
 
     // Output for the last run:
     textures_delta: TexturesDelta,
@@ -114,7 +114,7 @@ impl AppRunner {
             needs_repaint,
             last_save_time: now_sec(),
             text_agent,
-            screenshot_commands: vec![],
+            screenshot_commands_with_frame_delay: vec![],
             textures_delta: Default::default(),
             clipped_primitives: None,
         };
@@ -236,7 +236,8 @@ impl AppRunner {
             for command in viewport_output.commands {
                 match command {
                     ViewportCommand::Screenshot(user_data) => {
-                        self.screenshot_commands.push(user_data);
+                        self.screenshot_commands_with_frame_delay
+                            .push((user_data, 1));
                     }
                     _ => {
                         // TODO(emilk): handle some of the commands
@@ -259,12 +260,27 @@ impl AppRunner {
         let clipped_primitives = std::mem::take(&mut self.clipped_primitives);
 
         if let Some(clipped_primitives) = clipped_primitives {
+            let mut screenshot_commands = vec![];
+            self.screenshot_commands_with_frame_delay
+                .retain_mut(|(user_data, frame_delay)| {
+                    if *frame_delay == 0 {
+                        screenshot_commands.push(user_data.clone());
+                        false
+                    } else {
+                        *frame_delay -= 1;
+                        true
+                    }
+                });
+            if !self.screenshot_commands_with_frame_delay.is_empty() {
+                self.egui_ctx().request_repaint();
+            }
+
             if let Err(err) = self.painter.paint_and_update_textures(
                 self.app.clear_color(&self.egui_ctx.style().visuals),
                 &clipped_primitives,
                 self.egui_ctx.pixels_per_point(),
                 &textures_delta,
-                mem::take(&mut self.screenshot_commands),
+                screenshot_commands,
             ) {
                 log::error!("Failed to paint: {}", super::string_from_js_value(&err));
             }
diff --git a/crates/egui/src/viewport.rs b/crates/egui/src/viewport.rs
index adafeca43a8..91cd12e2b60 100644
--- a/crates/egui/src/viewport.rs
+++ b/crates/egui/src/viewport.rs
@@ -1056,7 +1056,7 @@ pub enum ViewportCommand {
     /// Enable mouse pass-through: mouse clicks pass through the window, used for non-interactable overlays.
     MousePassthrough(bool),
 
-    /// Take a screenshot.
+    /// Take a screenshot of the next frame after this.
     ///
     /// The results are returned in [`crate::Event::Screenshot`].
     Screenshot(crate::UserData),