Add a way to override the element ID for primary output check

src/niri.rs

@@ -20,7 +20,8 @@ use smithay::backend::renderer::element::surface::{
 };
 use smithay::backend::renderer::element::utils::{select_dmabuf_feedback, RelocateRenderElement};
 use smithay::backend::renderer::element::{
-    default_primary_scanout_output_compare, AsRenderElements, Kind, RenderElementStates,
+    default_primary_scanout_output_compare, AsRenderElements, Id, Kind, PrimaryScanoutOutput,
+    RenderElementStates,
 };
 use smithay::backend::renderer::gles::GlesRenderer;
 use smithay::desktop::utils::{
@@ -274,6 +275,9 @@ struct SurfaceFrameThrottlingState {
     last_sent_at: RefCell<Option<(Output, u32)>>,
 }
 
+#[derive(Default)]
+pub struct WindowOffscreenId(pub RefCell<Option<Id>>);
+
 impl Default for SurfaceFrameThrottlingState {
     fn default() -> Self {
         Self {
@@ -2156,13 +2160,28 @@ impl Niri {
         // not in a unified way with the pointer surfaces, which makes the logic elsewhere simpler.
 
         for win in self.layout.windows_for_output(output) {
+            let offscreen_id = win
+                .user_data()
+                .get_or_insert(WindowOffscreenId::default)
+                .0
+                .borrow();
+            let offscreen_id = offscreen_id.as_ref();
+
             win.with_surfaces(|surface, states| {
-                update_surface_primary_scanout_output(
+                states
-                    surface,
+                    .data_map
+                    .insert_if_missing_threadsafe(Mutex::<PrimaryScanoutOutput>::default);
+                let surface_primary_scanout_output = states
+                    .data_map
+                    .get::<Mutex<PrimaryScanoutOutput>>()
+                    .unwrap();
+                surface_primary_scanout_output
+                    .lock()
+                    .unwrap()
+                    .update_from_render_element_states(
+                        offscreen_id.cloned().unwrap_or_else(|| surface.into()),
                         output,
-                    states,
                         render_element_states,
-                    // Windows are shown only on one output at a time.
                         |_, _, output, _| output,
                     );
             });