Merge pull request #2836 from CapSoftware/perf/gpui-screenshot-editor-copies

perf(desktop-gpui): convert screenshot editor frames in one pass and skip stale ones

Author: richiemcilroy

apps/desktop-gpui/src/screenshot_annotations.rs

@@ -2242,7 +2242,7 @@ impl ScreenshotEditorWindow {
     // -- Mask overlays ---------------------------------------------------------
 
     fn mask_requests(&self) -> Vec<MaskRequest> {
-        if self.frame_rgba.is_none() {
+        if self.frame.is_none() {
             return Vec::new();
         }
         let size = (
@@ -2312,14 +2312,16 @@ impl ScreenshotEditorWindow {
         }
 
         let previous = self.annotation_state.masks.clone();
-        let Some(rgba) = self.frame_rgba.clone() else {
+        let Some(frame) = self.frame.clone() else {
             return;
         };
         self.annotation_state.mask_pending = Some(requests.clone());
         self.annotation_state.mask_task = Some(cx.spawn_in(window, async move |this, cx| {
             let overlays = cx
                 .background_executor()
-                .spawn(async move { build_mask_overlays(&rgba, &previous, &requests) })
+                .spawn(async move {
+                    build_mask_overlays(frame.as_bytes(0).unwrap_or_default(), &previous, &requests)
+                })
                 .await;
             this.update_in(cx, |this, window, cx| {
                 let overlays = finish_mask_build(&this.annotation_state.masks, overlays);
@@ -3952,22 +3954,26 @@ fn checker_quarters(size: f32) -> Vec<AnyElement> {
 /// `renderMaskOverlays` + `blurRegion` (`Preview.tsx:537-687`): resample the
 /// frame's own pixels inside each mask and hand back one image per region.
 /// Runs on the background executor -- the regions are small, but a 50-block
-/// pixelate over a 4K screenshot is still not main-thread work.
+/// pixelate over a 4K screenshot is still not main-thread work. The filter
+/// treats every channel alike, so it runs on the atlas's BGRA as-is and the
+/// overlays come out already in atlas order.
 fn build_mask_overlays(
-    rgba: &[u8],
+    bgra: &[u8],
     previous: &[MaskEntry<MaskOverlay>],
     requests: &[MaskRequest],
 ) -> Vec<MaskEntry<MaskOverlay>> {
     update_masks(previous, requests, |request| {
-        let processed = request.filter(rgba)?;
+        let processed = request.filter(bgra)?;
         Some(MaskOverlay {
             rect: Rect {
                 x: request.region.0 as f64,
                 y: request.region.1 as f64,
                 width: processed.width() as f64,
                 height: processed.height() as f64,
             },
-            image: crate::library::rgba_to_render_image(processed),
+            image: Arc::new(RenderImage::new(smallvec::smallvec![image::Frame::new(
+                processed
+            )])),
         })
     })
 }
@@ -4523,6 +4529,68 @@ mod tests {
         assert_eq!(mask_region(&away, rect, (200, 200)), None);
     }
 
+    #[test]
+    fn mask_overlays_from_bgra_match_filtering_rgba_then_swapping() {
+        let (width, height) = (37u32, 23u32);
+        let stride = (width as usize * 4).next_multiple_of(256);
+        let padded: Vec<u8> = (0..stride * height as usize)
+            .map(|index| (index * 31 + index / 7) as u8)
+            .collect();
+        let rgba = cap_gpui_kernels::frame::unpad_rows(
+            &padded,
+            width as usize * 4,
+            stride,
+            height as usize,
+        )
+        .unwrap();
+        let bgra = cap_gpui_kernels::frame::unpad_rgba_to_bgra(
+            &padded,
+            width as usize,
+            stride,
+            height as usize,
+        )
+        .unwrap();
+        let frame = MaskFrame {
+            revision: 1,
+            size: (width, height),
+            content: [0., 0., f64::from(width), f64::from(height)].map(f64::to_bits),
+            padding: 0_f64.to_bits(),
+            crop: None,
+        };
+        let mut masks = Vec::new();
+        for (index, (mask_type, level, region)) in [
+            (MaskType::Blur, 4., (0., 0., 37., 23.)),
+            (MaskType::Blur, 9., (3., 1., 17., 11.)),
+            (MaskType::Blur, 16., (20., 5., 17., 18.)),
+            (MaskType::Blur, 33., (1., 2., 35., 21.)),
+            (MaskType::Pixelate, 2., (0., 0., 37., 23.)),
+            (MaskType::Pixelate, 7., (5., 3., 13., 9.)),
+            (MaskType::Pixelate, 20., (11., 7., 26., 16.)),
+        ]
+        .into_iter()
+        .enumerate()
+        {
+            let mut mask = shape(AnnotationType::Mask, region.0, region.1, region.2, region.3);
+            mask.id = index.to_string();
+            mask.mask_type = Some(mask_type);
+            mask.mask_level = Some(level);
+            masks.push(mask);
+        }
+        let requests: Vec<_> = masks
+            .iter()
+            .filter_map(|mask| mask_request(mask, frame))
+            .collect();
+        assert_eq!(requests.len(), masks.len());
+
+        let overlays = build_mask_overlays(&bgra, &[], &requests);
+        for (request, entry) in requests.iter().zip(&overlays) {
+            let expected = crate::library::rgba_to_render_image(request.filter(&rgba).unwrap());
+            let overlay = entry.value.as_ref().unwrap();
+            assert_eq!(overlay.image.size(0), expected.size(0));
+            assert_eq!(overlay.image.as_bytes(0), expected.as_bytes(0));
+        }
+    }
+
     /// The resample itself: one overlay per mask, at the region's own size,
     /// with the row slicing staying inside the frame buffer.
     #[test]

apps/desktop-gpui/src/screenshot_editor.rs

@@ -912,12 +912,10 @@ pub struct ScreenshotEditorWindow {
     /// renders.
     exporting: bool,
     export_status: ExportStatus,
-    /// The latest GPU frame, already BGRA in gpui's atlas format.
-    frame: Option<Arc<RenderImage>>,
-    /// The same frame as tight RGBA. `Preview.tsx` reads the preview canvas
-    /// back to blur/pixelate mask regions (`:537-687`); the mask overlay
-    /// resamples this instead of the atlas image.
-    pub(crate) frame_rgba: Option<Arc<Vec<u8>>>,
+    /// The latest GPU frame, already BGRA in gpui's atlas format. `Preview.tsx`
+    /// reads the preview canvas back to blur/pixelate mask regions
+    /// (`:537-687`); the mask overlay resamples these same pixels.
+    pub(crate) frame: Option<Arc<RenderImage>>,
     pub(crate) mask_frame_revision: u64,
     pub(crate) frame_size: (f32, f32),
     /// `originalImageSize` -- the source PNG's dimensions, which the crop
@@ -1003,7 +1001,6 @@ impl ScreenshotEditorWindow {
             exporting: false,
             export_status: ExportStatus::Idle,
             frame: None,
-            frame_rgba: None,
             mask_frame_revision: 0,
             frame_size: (1., 1.),
             image_size: None,
@@ -1126,15 +1123,13 @@ impl ScreenshotEditorWindow {
     pub fn frame_arrived(
         &mut self,
         image: Arc<RenderImage>,
-        rgba: Arc<Vec<u8>>,
         size: (u32, u32),
         window: &mut Window,
         cx: &mut Context<Self>,
     ) {
         let next = (size.0.max(1) as f32, size.1.max(1) as f32);
         self.rescale_annotations_for(next);
         self.frame_size = next;
-        self.frame_rgba = Some(rgba);
         self.mask_frame_revision = self
             .mask_frame_revision
             .checked_add(1)
@@ -4747,26 +4742,10 @@ fn render_input_key(project: &mut ProjectConfiguration) -> Option<Vec<u8>> {
     Some(key)
 }
 
-/// One rendered frame, in both the forms the window needs: tight RGBA (what a
-/// mask overlay resamples) and the BGRA copy gpui's atlas takes. The rows are
-/// copied once and the swap runs over a clone, so the padded GPU buffer is
-/// walked a single time.
-fn frame_buffers(frame: &RenderedFrame) -> Option<(Arc<Vec<u8>>, Arc<RenderImage>)> {
-    let rgba = cap_gpui_kernels::frame::unpad_rows(
-        &frame.data,
-        frame.width as usize * 4,
-        frame.padded_bytes_per_row as usize,
-        frame.height as usize,
-    )?;
-    let mut bgra = rgba.clone();
-    cap_gpui_kernels::frame::swap_red_blue(&mut bgra);
-    let buffer = image::RgbaImage::from_raw(frame.width, frame.height, bgra)?;
-    Some((
-        Arc::new(rgba),
-        Arc::new(RenderImage::new(smallvec::smallvec![image::Frame::new(
-            buffer
-        )])),
-    ))
+/// Only the newest queued frame is ever shown, so older ones are dropped
+/// before the copy rather than converted and replaced.
+fn latest_frame<T>(frame: T, frames: &flume::Receiver<T>) -> T {
+    frames.drain().last().unwrap_or(frame)
 }
 
 /// Read the bundle, start the still renderer on tokio, pump frames back --
@@ -4837,13 +4816,14 @@ pub fn load_screenshot_project(
             return;
         }
 
-        while let Ok((frame, revision)) = frame_rx.recv_async().await {
+        while let Ok(frame) = frame_rx.recv_async().await {
+            let (frame, revision) = latest_frame(frame, &frame_rx);
             let size = (frame.width, frame.height);
-            let buffers = cx
+            let image = cx
                 .background_executor()
-                .spawn(async move { frame_buffers(&frame) })
+                .spawn(async move { crate::editor_window::frame_image(&frame) })
                 .await;
-            let Some((rgba, image)) = buffers else {
+            let Some(image) = image else {
                 tracing::warn!("a rendered screenshot frame could not be converted for display");
                 continue;
             };
@@ -4855,7 +4835,7 @@ pub fn load_screenshot_project(
             );
             if handle
                 .update(cx, |view, window, cx| {
-                    view.frame_arrived(image, rgba, size, window, cx)
+                    view.frame_arrived(image, size, window, cx)
                 })
                 .is_err()
             {
@@ -4887,6 +4867,56 @@ mod tests {
         assert_eq!(CHECKER_TILE % (CHECKER_CELL * 2), 0);
     }
 
+    #[test]
+    fn the_displayed_frame_matches_unpadding_then_swapping_a_copy() {
+        for (width, height) in [(1u32, 1u32), (3, 5), (7, 2), (33, 9), (1001, 3)] {
+            let row_bytes = width * 4;
+            for stride in [row_bytes, row_bytes + 4, row_bytes.next_multiple_of(256)] {
+                let frame = RenderedFrame {
+                    data: Arc::new(
+                        (0..stride as usize * height as usize)
+                            .map(|index| (index * 7 + index / 251) as u8)
+                            .collect(),
+                    ),
+                    width,
+                    height,
+                    padded_bytes_per_row: stride,
+                    frame_number: 0,
+                    target_time_ns: 0,
+                };
+                let rgba = cap_gpui_kernels::frame::unpad_rows(
+                    &frame.data,
+                    row_bytes as usize,
+                    stride as usize,
+                    height as usize,
+                )
+                .unwrap();
+                let mut bgra = rgba.clone();
+                cap_gpui_kernels::frame::swap_red_blue(&mut bgra);
+
+                let image = crate::editor_window::frame_image(&frame).unwrap();
+                assert_eq!(
+                    image.size(0),
+                    gpui::size(
+                        gpui::DevicePixels::from(width),
+                        gpui::DevicePixels::from(height)
+                    )
+                );
+                assert_eq!(image.as_bytes(0).unwrap(), bgra.as_slice());
+            }
+        }
+    }
+
+    #[test]
+    fn only_the_newest_queued_frame_is_converted() {
+        let (frame_tx, frame_rx) = flume::bounded(2);
+        assert_eq!(latest_frame(1, &frame_rx), 1);
+        frame_tx.send(2).unwrap();
+        frame_tx.send(3).unwrap();
+        assert_eq!(latest_frame(1, &frame_rx), 3);
+        assert!(frame_rx.is_empty());
+    }
+
     /// `PaddingPopover`'s guard (`context.tsx:30-42`).
     #[test]
     fn an_invisible_background_is_one_with_nothing_to_show() {
@@ -5182,7 +5212,7 @@ mod tests {
         ));
         setup_rx.recv_async().await.unwrap().unwrap();
         let (first, _) = frame_rx.recv_async().await.unwrap();
-        assert!(frame_buffers(&first).is_some());
+        assert!(crate::editor_window::frame_image(&first).is_some());
 
         let renders = Arc::new(std::sync::atomic::AtomicUsize::new(0));
         let render_ms = Arc::new(std::sync::Mutex::new(Vec::new()));
@@ -5194,10 +5224,12 @@ mod tests {
                 (renders.clone(), render_ms.clone(), sent_at.clone());
             tokio::spawn(async move {
                 while let Ok((frame, revision)) = frame_rx.recv_async().await {
-                    let buffers = tokio::task::spawn_blocking(move || frame_buffers(&frame))
-                        .await
-                        .unwrap();
-                    assert!(buffers.is_some());
+                    let image = tokio::task::spawn_blocking(move || {
+                        crate::editor_window::frame_image(&frame)
+                    })
+                    .await
+                    .unwrap();
+                    assert!(image.is_some());
                     if let Some(sent) = sent_at.lock().unwrap().remove(&revision) {
                         render_ms
                             .lock()