Make screenshot UI render target-aware

3 files changed, 57 insertions, 40 deletions

diff --git a/resources/default-config.kdl b/resources/default-config.kdl
index 0bb31697..728baa53 100644
--- a/resources/default-config.kdl
+++ b/resources/default-config.kdl
@@ -402,14 +402,17 @@ animations {
     // To preview and set up this rule, check the preview-render option
     // in the debug section of the config.
     //
-    // WARNING: the window is NOT blocked out from the screenshot UI.
-    // If you open the screenshot UI while screencasting, blocked out windows
-    // WILL BE VISIBLE on the screencast.
+    // WARNING: the window is NOT blocked out from third-party screenshot tools.
+    // If you open some screenshot tool with preview while screencasting,
+    // blocked out windows WILL BE VISIBLE on the screencast.
+    //
+    // The built-in screenshot UI is not affected though, you can use it safely,
+    // and windows will remain blocked out even when screencasting it.
     block-out-from "screencast"
 
     // You can also block out the window out of all screen captures, including
-    // the screenshot UI. This way you avoid accidentally showing the window
-    // on a screencast when opening the screenshot UI.
+    // third-party screenshot tools. This way you avoid accidentally showing
+    // the window on a screencast when opening a third-party screenshot preview.
     block-out-from "screen-capture"
 
     // You can amend the window's minimum and maximum size in logical pixels.
@@ -459,7 +462,7 @@ window-rule {
     match app-id=r#"^org\.keepassxc\.KeePassXC$"#
     match app-id=r#"^org\.gnome\.World\.Secrets$"#
 
-    // Warning: will be visible when opening the screenshot UI.
+    // Warning: will be visible on third-party screenshot tools.
     block-out-from "screencast"
 
     // Use this instead to block out from all screen captures.
diff --git a/src/niri.rs b/src/niri.rs
index 88326119..02169a58 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -2307,7 +2307,7 @@ impl Niri {
         if self.screenshot_ui.is_open() {
             elements.extend(
                 self.screenshot_ui
-                    .render_output(output)
+                    .render_output(output, target)
                     .into_iter()
                     .map(OutputRenderElements::from),
             );
@@ -3104,31 +3104,37 @@ impl Niri {
                 let size = transform.transform_size(size);
 
                 let scale = Scale::from(output.current_scale().fractional_scale());
-                let elements = self.render::<GlesRenderer>(
-                    renderer,
-                    &output,
-                    true,
+                let targets = [
+                    RenderTarget::Output,
+                    RenderTarget::Screencast,
                     RenderTarget::ScreenCapture,
-                );
-                let elements = elements.iter().rev();
+                ];
+                let textures = targets.map(|target| {
+                    let elements = self.render::<GlesRenderer>(renderer, &output, true, target);
+                    let elements = elements.iter().rev();
 
-                let res = render_to_texture(
-                    renderer,
-                    size,
-                    scale,
-                    Transform::Normal,
-                    Fourcc::Abgr8888,
-                    elements,
-                );
-                let screenshot = match res {
-                    Ok((texture, _)) => texture,
-                    Err(err) => {
+                    let res = render_to_texture(
+                        renderer,
+                        size,
+                        scale,
+                        Transform::Normal,
+                        Fourcc::Abgr8888,
+                        elements,
+                    );
+
+                    if let Err(err) = &res {
                         warn!("error rendering output {}: {err:?}", output.name());
-                        return None;
                     }
-                };
 
-                Some((output, screenshot))
+                    res
+                });
+
+                if textures.iter().any(|res| res.is_err()) {
+                    return None;
+                }
+
+                let textures = textures.map(|res| res.unwrap().0);
+                Some((output, textures))
             })
             .collect();
 
diff --git a/src/ui/screenshot_ui.rs b/src/ui/screenshot_ui.rs
index 79245387..5892fde1 100644
--- a/src/ui/screenshot_ui.rs
+++ b/src/ui/screenshot_ui.rs
@@ -19,6 +19,7 @@ use smithay::utils::{Physical, Point, Rectangle, Size, Transform};
 
 use crate::niri_render_elements;
 use crate::render_helpers::primary_gpu_texture::PrimaryGpuTextureRenderElement;
+use crate::render_helpers::RenderTarget;
 
 const BORDER: i32 = 2;
 
@@ -42,8 +43,9 @@ pub struct OutputData {
     size: Size<i32, Physical>,
     scale: i32,
     transform: Transform,
-    texture: GlesTexture,
-    texture_buffer: TextureBuffer<GlesTexture>,
+    // Output, screencast, screen capture.
+    texture: [GlesTexture; 3],
+    texture_buffer: [TextureBuffer<GlesTexture>; 3],
     buffers: [SolidColorBuffer; 8],
     locations: [Point<i32, Physical>; 8],
 }
@@ -65,7 +67,8 @@ impl ScreenshotUi {
     pub fn open(
         &mut self,
         renderer: &GlesRenderer,
-        screenshots: HashMap<Output, GlesTexture>,
+        // Output, screencast, screen capture.
+        screenshots: HashMap<Output, [GlesTexture; 3]>,
         default_output: Output,
     ) -> bool {
         if screenshots.is_empty() {
@@ -110,13 +113,9 @@ impl ScreenshotUi {
                 let output_mode = output.current_mode().unwrap();
                 let size = transform.transform_size(output_mode.size);
                 let scale = output.current_scale().integer_scale();
-                let texture_buffer = TextureBuffer::from_texture(
-                    renderer,
-                    texture.clone(),
-                    scale,
-                    Transform::Normal,
-                    None,
-                );
+                let texture_buffer = texture.clone().map(|texture| {
+                    TextureBuffer::from_texture(renderer, texture, scale, Transform::Normal, None)
+                });
                 let buffers = [
                     SolidColorBuffer::new((0, 0), [1., 1., 1., 1.]),
                     SolidColorBuffer::new((0, 0), [1., 1., 1., 1.]),
@@ -243,7 +242,11 @@ impl ScreenshotUi {
         }
     }
 
-    pub fn render_output(&self, output: &Output) -> ArrayVec<ScreenshotUiRenderElement, 9> {
+    pub fn render_output(
+        &self,
+        output: &Output,
+        target: RenderTarget,
+    ) -> ArrayVec<ScreenshotUiRenderElement, 9> {
         let _span = tracy_client::span!("ScreenshotUi::render_output");
 
         let Self::Open { output_data, .. } = self else {
@@ -269,10 +272,15 @@ impl ScreenshotUi {
         }));
 
         // The screenshot itself goes last.
+        let index = match target {
+            RenderTarget::Output => 0,
+            RenderTarget::Screencast => 1,
+            RenderTarget::ScreenCapture => 2,
+        };
         elements.push(
             PrimaryGpuTextureRenderElement(TextureRenderElement::from_texture_buffer(
                 (0., 0.),
-                &output_data.texture_buffer,
+                &output_data.texture_buffer[index],
                 None,
                 None,
                 None,
@@ -307,7 +315,7 @@ impl ScreenshotUi {
             .to_buffer(1, Transform::Normal, &data.size.to_logical(1));
 
         let mapping = renderer
-            .copy_texture(&data.texture, buf_rect, Fourcc::Abgr8888)
+            .copy_texture(&data.texture[0], buf_rect, Fourcc::Abgr8888)
             .context("error copying texture")?;
         let copy = renderer
             .map_texture(&mapping)