From df36eac25bd7e0f29781e40b06ca34fc72cf3388 Mon Sep 17 00:00:00 2001
From: Ivan Molodetskikh <yalterz@gmail.com>
Date: Wed, 15 May 2024 20:08:37 +0400
Subject: Fix render elements looking off on screenshots

---
 src/layout/mod.rs | 18 ++++++++++++++++++
 src/niri.rs       | 18 ++++++++++++++++--
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/src/layout/mod.rs b/src/layout/mod.rs
index f4b88560..d091c024 100644
--- a/src/layout/mod.rs
+++ b/src/layout/mod.rs
@@ -1413,6 +1413,24 @@ impl<W: LayoutElement> Layout<W> {
         }
     }
 
+    pub fn update_render_elements_all(&mut self) {
+        let _span = tracy_client::span!("Layout::update_render_elements_all");
+
+        let MonitorSet::Normal {
+            monitors,
+            active_monitor_idx,
+            ..
+        } = &mut self.monitor_set
+        else {
+            error!("update_render_elements_all called with no monitors");
+            return;
+        };
+
+        for (idx, mon) in monitors.iter_mut().enumerate() {
+            mon.update_render_elements(idx == *active_monitor_idx);
+        }
+    }
+
     pub fn update_shaders(&mut self) {
         match &mut self.monitor_set {
             MonitorSet::Normal { monitors, .. } => {
diff --git a/src/niri.rs b/src/niri.rs
index 4e887570..a8d906c1 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -3268,6 +3268,8 @@ impl Niri {
         let output = screencopy.output().clone();
         ensure!(self.output_state.contains_key(&output), "output is missing");
 
+        self.layout.update_render_elements(&output);
+
         backend
             .with_primary_renderer(move |renderer| {
                 let elements = self
@@ -3359,6 +3361,8 @@ impl Niri {
             return;
         };
 
+        self.layout.update_render_elements_all();
+
         let screenshots = self
             .global_space
             .outputs()
@@ -3410,9 +3414,15 @@ impl Niri {
         self.queue_redraw_all();
     }
 
-    pub fn screenshot(&self, renderer: &mut GlesRenderer, output: &Output) -> anyhow::Result<()> {
+    pub fn screenshot(
+        &mut self,
+        renderer: &mut GlesRenderer,
+        output: &Output,
+    ) -> anyhow::Result<()> {
         let _span = tracy_client::span!("Niri::screenshot");
 
+        self.layout.update_render_elements(output);
+
         let size = output.current_mode().unwrap().size;
         let transform = output.current_transform();
         let size = transform.transform_size(size);
@@ -3555,13 +3565,15 @@ impl Niri {
 
     #[cfg(feature = "dbus")]
     pub fn screenshot_all_outputs(
-        &self,
+        &mut self,
         renderer: &mut GlesRenderer,
         include_pointer: bool,
         on_done: impl FnOnce(PathBuf) + Send + 'static,
     ) -> anyhow::Result<()> {
         let _span = tracy_client::span!("Niri::screenshot_all_outputs");
 
+        self.layout.update_render_elements_all();
+
         let outputs: Vec<_> = self.global_space.outputs().cloned().collect();
 
         // FIXME: support multiple outputs, needs fixing multi-scale handling and cropping.
@@ -3773,6 +3785,8 @@ impl Niri {
     pub fn do_screen_transition(&mut self, renderer: &mut GlesRenderer, delay_ms: Option<u16>) {
         let _span = tracy_client::span!("Niri::do_screen_transition");
 
+        self.layout.update_render_elements_all();
+
         let textures: Vec<_> = self
             .output_state
             .keys()
-- 
cgit