2 files changed, 12 insertions, 27 deletions

diff --git a/src/layout/mod.rs b/src/layout/mod.rs
index 9fbd1623..9376112f 100644
--- a/src/layout/mod.rs
+++ b/src/layout/mod.rs
@@ -1800,7 +1800,7 @@ impl<W: LayoutElement> Layout<W> {
         }
     }
 
-    pub fn update_render_elements(&mut self, output: &Output) {
+    pub fn update_render_elements(&mut self, output: Option<&Output>) {
         let _span = tracy_client::span!("Layout::update_render_elements");
 
         let MonitorSet::Normal {
@@ -1814,31 +1814,12 @@ impl<W: LayoutElement> Layout<W> {
         };
 
         for (idx, mon) in monitors.iter_mut().enumerate() {
-            if mon.output == *output {
+            if output.map_or(true, |output| mon.output == *output) {
                 mon.update_render_elements(idx == *active_monitor_idx);
-                return;
             }
         }
     }
 
-    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 a0f85da8..277fc45f 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -1278,7 +1278,7 @@ impl State {
             return;
         };
 
-        self.niri.layout.update_render_elements_all();
+        self.niri.update_render_elements(None);
 
         let Some(screenshots) = self
             .backend
@@ -2853,6 +2853,10 @@ impl Niri {
         }
     }
 
+    pub fn update_render_elements(&mut self, output: Option<&Output>) {
+        self.layout.update_render_elements(output);
+    }
+
     pub fn render<R: NiriRenderer>(
         &self,
         renderer: &mut R,
@@ -3066,7 +3070,7 @@ impl Niri {
             // Also keep redrawing during a screen transition.
             state.unfinished_animations_remain |= state.screen_transition.is_some();
 
-            self.layout.update_render_elements(output);
+            self.update_render_elements(Some(output));
 
             // Render.
             res = backend.render(self, output, target_presentation_time);
@@ -3891,7 +3895,7 @@ impl Niri {
             "screencopy output missing"
         );
 
-        self.layout.update_render_elements(output);
+        self.update_render_elements(Some(output));
 
         let elements = self.render(
             renderer,
@@ -4135,7 +4139,7 @@ impl Niri {
     ) -> anyhow::Result<()> {
         let _span = tracy_client::span!("Niri::screenshot");
 
-        self.layout.update_render_elements(output);
+        self.update_render_elements(Some(output));
 
         let size = output.current_mode().unwrap().size;
         let transform = output.current_transform();
@@ -4286,7 +4290,7 @@ impl Niri {
     ) -> anyhow::Result<()> {
         let _span = tracy_client::span!("Niri::screenshot_all_outputs");
 
-        self.layout.update_render_elements_all();
+        self.update_render_elements(None);
 
         let outputs: Vec<_> = self.global_space.outputs().cloned().collect();
 
@@ -4541,7 +4545,7 @@ 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();
+        self.update_render_elements(None);
 
         let textures: Vec<_> = self
             .output_state