Move unmap snapshot from Mapped to Tile

7 files changed, 129 insertions, 114 deletions

diff --git a/niri-visual-tests/src/test_window.rs b/niri-visual-tests/src/test_window.rs
index 91da5bdb..644d98b4 100644
--- a/niri-visual-tests/src/test_window.rs
+++ b/niri-visual-tests/src/test_window.rs
@@ -227,10 +227,6 @@ impl LayoutElement for TestWindow {
         &EMPTY
     }
 
-    fn take_unmap_snapshot(&self) -> Option<LayoutElementRenderSnapshot> {
-        None
-    }
-
     fn animation_snapshot(&self) -> Option<&LayoutElementRenderSnapshot> {
         None
     }
diff --git a/src/handlers/compositor.rs b/src/handlers/compositor.rs
index 060980a5..37903a20 100644
--- a/src/handlers/compositor.rs
+++ b/src/handlers/compositor.rs
@@ -303,8 +303,9 @@ impl CompositorHandler for State {
         // Test client: alacritty with CSD.
         if let Some(root) = self.niri.root_surface.get(surface) {
             if let Some((mapped, _)) = self.niri.layout.find_window_and_output(root) {
+                let window = mapped.window.clone();
                 self.backend.with_primary_renderer(|renderer| {
-                    mapped.store_unmap_snapshot_if_empty(renderer);
+                    self.niri.layout.store_unmap_snapshot(renderer, &window);
                 });
             }
         }
diff --git a/src/handlers/xdg_shell.rs b/src/handlers/xdg_shell.rs
index 6f627850..9864ccdc 100644
--- a/src/handlers/xdg_shell.rs
+++ b/src/handlers/xdg_shell.rs
@@ -30,7 +30,6 @@ use smithay::{
 };
 
 use crate::layout::workspace::ColumnWidth;
-use crate::layout::LayoutElement as _;
 use crate::niri::{PopupGrabState, State};
 use crate::window::{InitialConfigureState, ResolvedWindowRules, Unmapped, WindowRef};
 
@@ -389,7 +388,7 @@ impl XdgShellHandler for State {
         let output = output.clone();
 
         self.backend.with_primary_renderer(|renderer| {
-            mapped.store_unmap_snapshot_if_empty(renderer);
+            self.niri.layout.store_unmap_snapshot(renderer, &window);
         });
         self.backend.with_primary_renderer(|renderer| {
             self.niri
@@ -881,22 +880,22 @@ pub fn add_mapped_toplevel_pre_commit_hook(toplevel: &ToplevelSurface) -> HookId
             false
         };
 
+        let window = mapped.window.clone();
         if got_unmapped {
             state.backend.with_primary_renderer(|renderer| {
-                mapped.store_unmap_snapshot_if_empty(renderer);
+                state.niri.layout.store_unmap_snapshot(renderer, &window);
             });
         } else {
-            // The toplevel remains mapped; clear any stored unmap snapshot.
-            let _ = mapped.take_unmap_snapshot();
-
             if animate {
                 state.backend.with_primary_renderer(|renderer| {
                     mapped.store_animation_snapshot(renderer);
                 });
 
-                let window = mapped.window.clone();
                 state.niri.layout.prepare_for_resize_animation(&window);
             }
+
+            // The toplevel remains mapped; clear any stored unmap snapshot.
+            state.niri.layout.clear_unmap_snapshot(&window);
         }
     })
 }
diff --git a/src/layout/mod.rs b/src/layout/mod.rs
index 4889e9a7..5ca48ad6 100644
--- a/src/layout/mod.rs
+++ b/src/layout/mod.rs
@@ -140,8 +140,6 @@ pub trait LayoutElement {
     /// Runs periodic clean-up tasks.
     fn refresh(&self);
 
-    fn take_unmap_snapshot(&self) -> Option<LayoutElementRenderSnapshot>;
-
     fn animation_snapshot(&self) -> Option<&LayoutElementRenderSnapshot>;
     fn take_animation_snapshot(&mut self) -> Option<LayoutElementRenderSnapshot>;
 }
@@ -1749,6 +1747,54 @@ impl<W: LayoutElement> Layout<W> {
         }
     }
 
+    pub fn store_unmap_snapshot(&mut self, renderer: &mut GlesRenderer, window: &W::Id) {
+        let _span = tracy_client::span!("Layout::store_unmap_snapshot");
+
+        match &mut self.monitor_set {
+            MonitorSet::Normal { monitors, .. } => {
+                for mon in monitors {
+                    for ws in &mut mon.workspaces {
+                        if ws.has_window(window) {
+                            ws.store_unmap_snapshot_if_empty(renderer, window);
+                            return;
+                        }
+                    }
+                }
+            }
+            MonitorSet::NoOutputs { workspaces, .. } => {
+                for ws in workspaces {
+                    if ws.has_window(window) {
+                        ws.store_unmap_snapshot_if_empty(renderer, window);
+                        return;
+                    }
+                }
+            }
+        }
+    }
+
+    pub fn clear_unmap_snapshot(&mut self, window: &W::Id) {
+        match &mut self.monitor_set {
+            MonitorSet::Normal { monitors, .. } => {
+                for mon in monitors {
+                    for ws in &mut mon.workspaces {
+                        if ws.has_window(window) {
+                            ws.clear_unmap_snapshot(window);
+                            return;
+                        }
+                    }
+                }
+            }
+            MonitorSet::NoOutputs { workspaces, .. } => {
+                for ws in workspaces {
+                    if ws.has_window(window) {
+                        ws.clear_unmap_snapshot(window);
+                        return;
+                    }
+                }
+            }
+        }
+    }
+
     pub fn start_close_animation_for_window(
         &mut self,
         renderer: &mut GlesRenderer,
@@ -1996,10 +2042,6 @@ mod tests {
             &EMPTY
         }
 
-        fn take_unmap_snapshot(&self) -> Option<LayoutElementRenderSnapshot> {
-            None
-        }
-
         fn animation_snapshot(&self) -> Option<&LayoutElementRenderSnapshot> {
             None
         }
diff --git a/src/layout/tile.rs b/src/layout/tile.rs
index 40198716..56ac6cd0 100644
--- a/src/layout/tile.rs
+++ b/src/layout/tile.rs
@@ -1,3 +1,4 @@
+use std::cell::RefCell;
 use std::cmp::max;
 use std::rc::Rc;
 use std::time::Duration;
@@ -17,11 +18,10 @@ use super::{
 use crate::animation::Animation;
 use crate::niri_render_elements;
 use crate::render_helpers::offscreen::OffscreenRenderElement;
-use crate::render_helpers::primary_gpu_texture::PrimaryGpuTextureRenderElement;
 use crate::render_helpers::renderer::NiriRenderer;
 use crate::render_helpers::resize::ResizeRenderElement;
 use crate::render_helpers::snapshot::RenderSnapshot;
-use crate::render_helpers::{render_to_encompassing_texture, RenderTarget, ToRenderElement};
+use crate::render_helpers::{render_to_encompassing_texture, RenderTarget};
 
 /// Toplevel window with decorations.
 #[derive(Debug)]
@@ -62,6 +62,9 @@ pub struct Tile<W: LayoutElement> {
     /// The animation of a tile visually moving vertically.
     move_y_animation: Option<MoveAnimation>,
 
+    /// Snapshot of the last render for use in the close animation.
+    unmap_snapshot: RefCell<Option<TileRenderSnapshot>>,
+
     /// Configurable properties of the layout.
     pub options: Rc<Options>,
 }
@@ -76,20 +79,8 @@ niri_render_elements! {
     }
 }
 
-niri_render_elements! {
-    TileSnapshotContentsRenderElement => {
-        Texture = PrimaryGpuTextureRenderElement,
-        SolidColor = SolidColorRenderElement,
-    }
-}
-
-niri_render_elements! {
-    TileSnapshotRenderElement => {
-        Contents = RescaleRenderElement<TileSnapshotContentsRenderElement>,
-        FocusRing = FocusRingRenderElement,
-        SolidColor = SolidColorRenderElement,
-    }
-}
+type TileRenderSnapshot =
+    RenderSnapshot<TileRenderElement<GlesRenderer>, TileRenderElement<GlesRenderer>>;
 
 #[derive(Debug)]
 struct ResizeAnimation {
@@ -121,6 +112,7 @@ impl<W: LayoutElement> Tile<W> {
             resize_animation: None,
             move_x_animation: None,
             move_y_animation: None,
+            unmap_snapshot: RefCell::new(None),
             options,
         }
     }
@@ -696,80 +688,58 @@ impl<W: LayoutElement> Tile<W> {
         }
     }
 
-    fn render_snapshot<E, C>(
+    pub fn store_unmap_snapshot_if_empty(
         &self,
         renderer: &mut GlesRenderer,
         scale: Scale<f64>,
         view_size: Size<i32, Logical>,
-        contents: Vec<C>,
-    ) -> Vec<TileSnapshotRenderElement>
-    where
-        E: Into<TileSnapshotContentsRenderElement>,
-        C: ToRenderElement<RenderElement = E>,
-    {
-        let alpha = if self.is_fullscreen {
-            1.
-        } else {
-            self.window.rules().opacity.unwrap_or(1.).clamp(0., 1.)
-        };
-
-        let window_size = self.window_size();
-        let animated_window_size = self.animated_window_size();
-        let animated_scale = animated_window_size.to_f64() / window_size.to_f64();
-
-        let mut rv = vec![];
-
-        for baked in contents {
-            let elem = baked.to_render_element(self.window_loc(), scale, alpha, Kind::Unspecified);
-            let elem: TileSnapshotContentsRenderElement = elem.into();
-
-            let origin = self.window_loc().to_physical_precise_round(scale);
-            let elem = RescaleRenderElement::from_element(elem, origin, animated_scale);
-            rv.push(elem.into());
-        }
-
-        if let Some(width) = self.effective_border_width() {
-            rv.extend(
-                self.border
-                    .render(renderer, Point::from((width, width)), scale, view_size)
-                    .map(Into::into),
-            );
-        }
-
-        if self.is_fullscreen {
-            let elem = SolidColorRenderElement::from_buffer(
-                &self.fullscreen_backdrop,
-                Point::from((0, 0)),
-                scale,
-                1.,
-                Kind::Unspecified,
-            );
-            rv.push(elem.into());
+    ) {
+        let mut snapshot = self.unmap_snapshot.borrow_mut();
+        if snapshot.is_some() {
+            return;
         }
 
-        rv
+        *snapshot = Some(self.render_snapshot(renderer, scale, view_size));
     }
 
-    pub fn take_snapshot_for_close_anim(
+    fn render_snapshot(
         &self,
         renderer: &mut GlesRenderer,
         scale: Scale<f64>,
         view_size: Size<i32, Logical>,
-    ) -> Option<RenderSnapshot<TileSnapshotRenderElement, TileSnapshotRenderElement>> {
-        let snapshot = self.window.take_unmap_snapshot()?;
-
-        Some(RenderSnapshot {
-            contents: self.render_snapshot(renderer, scale, view_size, snapshot.contents),
-            blocked_out_contents: self.render_snapshot(
-                renderer,
-                scale,
-                view_size,
-                snapshot.blocked_out_contents,
-            ),
-            block_out_from: snapshot.block_out_from,
+    ) -> TileRenderSnapshot {
+        let _span = tracy_client::span!("Tile::render_snapshot");
+
+        let contents = self.render_inner(
+            renderer,
+            Point::from((0, 0)),
+            scale,
+            view_size,
+            false,
+            RenderTarget::Output,
+        );
+
+        // A bit of a hack to render blocked out as for screencast, but I think it's fine here.
+        let blocked_out_contents = self.render_inner(
+            renderer,
+            Point::from((0, 0)),
+            scale,
+            view_size,
+            false,
+            RenderTarget::Screencast,
+        );
+
+        RenderSnapshot {
+            contents: contents.collect(),
+            blocked_out_contents: blocked_out_contents.collect(),
+            block_out_from: self.window.rules().block_out_from,
             size: self.animated_tile_size(),
             texture: Default::default(),
             blocked_out_texture: Default::default(),
-        })
+        }
+    }
+
+    pub fn take_unmap_snapshot(&self) -> Option<TileRenderSnapshot> {
+        self.unmap_snapshot.take()
     }
 }
diff --git a/src/layout/workspace.rs b/src/layout/workspace.rs
index fd5662b8..1aa1a02b 100644
--- a/src/layout/workspace.rs
+++ b/src/layout/workspace.rs
@@ -1176,6 +1176,32 @@ impl<W: LayoutElement> Workspace<W> {
         self.activate_column(column_idx);
     }
 
+    pub fn store_unmap_snapshot_if_empty(&mut self, renderer: &mut GlesRenderer, window: &W::Id) {
+        let (tile, _) = self
+            .tiles_in_render_order()
+            .find(|(tile, _)| tile.window().id() == window)
+            .unwrap();
+
+        // FIXME: workspaces should probably cache their last used scale so they can be correctly
+        // rendered even with no outputs connected.
+        let output_scale = self
+            .output
+            .as_ref()
+            .map(|o| Scale::from(o.current_scale().fractional_scale()))
+            .unwrap_or(Scale::from(1.));
+
+        tile.store_unmap_snapshot_if_empty(renderer, output_scale, self.view_size);
+    }
+
+    pub fn clear_unmap_snapshot(&mut self, window: &W::Id) {
+        let (tile, _) = self
+            .tiles_in_render_order()
+            .find(|(tile, _)| tile.window().id() == window)
+            .unwrap();
+
+        let _ = tile.take_unmap_snapshot();
+    }
+
     pub fn start_close_animation_for_window(
         &mut self,
         renderer: &mut GlesRenderer,
@@ -1194,9 +1220,7 @@ impl<W: LayoutElement> Workspace<W> {
             .map(|o| Scale::from(o.current_scale().fractional_scale()))
             .unwrap_or(Scale::from(1.));
 
-        let Some(snapshot) =
-            tile.take_snapshot_for_close_anim(renderer, output_scale, self.view_size)
-        else {
+        let Some(snapshot) = tile.take_unmap_snapshot() else {
             return;
         };
 
diff --git a/src/window/mapped.rs b/src/window/mapped.rs
index 9686a15f..90598431 100644
--- a/src/window/mapped.rs
+++ b/src/window/mapped.rs
@@ -47,9 +47,6 @@ pub struct Mapped {
     /// Buffer to draw instead of the window when it should be blocked out.
     block_out_buffer: RefCell<SolidColorBuffer>,
 
-    /// Snapshot of the last render for use in the close animation.
-    unmap_snapshot: RefCell<Option<LayoutElementRenderSnapshot>>,
-
     /// Whether the next configure should be animated, if the configured state changed.
     animate_next_configure: bool,
 
@@ -69,7 +66,6 @@ impl Mapped {
             need_to_recompute_rules: false,
             is_focused: false,
             block_out_buffer: RefCell::new(SolidColorBuffer::new((0, 0), [0., 0., 0., 1.])),
-            unmap_snapshot: RefCell::new(None),
             animate_next_configure: false,
             animate_serials: Vec::new(),
             animation_snapshot: None,
@@ -156,15 +152,6 @@ impl Mapped {
         }
     }
 
-    pub fn store_unmap_snapshot_if_empty(&self, renderer: &mut GlesRenderer) {
-        let mut snapshot = self.unmap_snapshot.borrow_mut();
-        if snapshot.is_some() {
-            return;
-        }
-
-        *snapshot = Some(self.render_snapshot(renderer));
-    }
-
     pub fn should_animate_commit(&mut self, commit_serial: Serial) -> bool {
         let mut should_animate = false;
         self.animate_serials.retain_mut(|serial| {
@@ -380,10 +367,6 @@ impl LayoutElement for Mapped {
         &self.rules
     }
 
-    fn take_unmap_snapshot(&self) -> Option<LayoutElementRenderSnapshot> {
-        self.unmap_snapshot.take()
-    }
-
     fn animation_snapshot(&self) -> Option<&LayoutElementRenderSnapshot> {
         self.animation_snapshot.as_ref()
     }