3 files changed, 49 insertions, 7 deletions

diff --git a/src/layout/mod.rs b/src/layout/mod.rs
index a49ec0e6..ae338ced 100644
--- a/src/layout/mod.rs
+++ b/src/layout/mod.rs
@@ -476,6 +476,13 @@ impl<W: LayoutElement> InteractiveMoveState<W> {
             _ => None,
         }
     }
+
+    fn moving_mut(&mut self) -> Option<&mut InteractiveMoveData<W>> {
+        match self {
+            InteractiveMoveState::Moving(move_) => Some(move_),
+            _ => None,
+        }
+    }
 }
 
 impl<W: LayoutElement> InteractiveMoveData<W> {
@@ -1608,6 +1615,28 @@ impl<W: LayoutElement> Layout<W> {
         moving_window.chain(mon_windows)
     }
 
+    pub fn windows_for_output_mut(&mut self, output: &Output) -> impl Iterator<Item = &mut W> + '_ {
+        let MonitorSet::Normal { monitors, .. } = &mut self.monitor_set else {
+            panic!()
+        };
+
+        let moving_window = self
+            .interactive_move
+            .as_mut()
+            .and_then(|x| x.moving_mut())
+            .filter(|move_| move_.output == *output)
+            .map(|move_| move_.tile.window_mut())
+            .into_iter();
+
+        let mon = monitors
+            .iter_mut()
+            .find(|mon| &mon.output == output)
+            .unwrap();
+        let mon_windows = mon.workspaces.iter_mut().flat_map(|ws| ws.windows_mut());
+
+        moving_window.chain(mon_windows)
+    }
+
     pub fn with_windows(&self, mut f: impl FnMut(&W, Option<&Output>, Option<WorkspaceId>)) {
         if let Some(InteractiveMoveState::Moving(move_)) = &self.interactive_move {
             f(move_.tile.window(), Some(&move_.output), None);
diff --git a/src/niri.rs b/src/niri.rs
index a7fc872f..96d52686 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -3895,7 +3895,7 @@ impl Niri {
         }
     }
 
-    pub fn send_frame_callbacks(&self, output: &Output) {
+    pub fn send_frame_callbacks(&mut self, output: &Output) {
         let _span = tracy_client::span!("Niri::send_frame_callbacks");
 
         let state = self.output_state.get(output).unwrap();
@@ -3936,8 +3936,8 @@ impl Niri {
 
         let frame_callback_time = get_monotonic_time();
 
-        for mapped in self.layout.windows_for_output(output) {
-            mapped.window.send_frame(
+        for mapped in self.layout.windows_for_output_mut(output) {
+            mapped.send_frame(
                 output,
                 frame_callback_time,
                 FRAME_CALLBACK_THROTTLE,
@@ -3985,7 +3985,7 @@ impl Niri {
         }
     }
 
-    pub fn send_frame_callbacks_on_fallback_timer(&self) {
+    pub fn send_frame_callbacks_on_fallback_timer(&mut self) {
         let _span = tracy_client::span!("Niri::send_frame_callbacks_on_fallback_timer");
 
         // Make up a bogus output; we don't care about it here anyway, just the throttling timer.
@@ -4002,8 +4002,8 @@ impl Niri {
 
         let frame_callback_time = get_monotonic_time();
 
-        self.layout.with_windows(|mapped, _, _| {
-            mapped.window.send_frame(
+        self.layout.with_windows_mut(|mapped, _| {
+            mapped.send_frame(
                 output,
                 frame_callback_time,
                 FRAME_CALLBACK_THROTTLE,
diff --git a/src/window/mapped.rs b/src/window/mapped.rs
index 2a1764f8..75bb5a06 100644
--- a/src/window/mapped.rs
+++ b/src/window/mapped.rs
@@ -13,7 +13,7 @@ use smithay::reexports::wayland_protocols::xdg::shell::server::xdg_toplevel;
 use smithay::reexports::wayland_server::protocol::wl_surface::WlSurface;
 use smithay::reexports::wayland_server::Resource as _;
 use smithay::utils::{Logical, Point, Rectangle, Scale, Serial, Size, Transform};
-use smithay::wayland::compositor::{remove_pre_commit_hook, with_states, HookId};
+use smithay::wayland::compositor::{remove_pre_commit_hook, with_states, HookId, SurfaceData};
 use smithay::wayland::seat::WaylandFocus;
 use smithay::wayland::shell::xdg::{SurfaceCachedState, ToplevelSurface};
 use wayland_backend::server::Credentials;
@@ -413,6 +413,19 @@ impl Mapped {
             WindowCastRenderElements::from(elem)
         })
     }
+
+    pub fn send_frame<T, F>(
+        &mut self,
+        output: &Output,
+        time: T,
+        throttle: Option<Duration>,
+        primary_scan_out_output: F,
+    ) where
+        T: Into<Duration>,
+        F: FnMut(&WlSurface, &SurfaceData) -> Option<Output> + Copy,
+    {
+        self.window.send_frame(output, time, throttle, primary_scan_out_output);
+    }
 }
 
 impl Drop for Mapped {