shell: Handle removed output state better

src/shell/layout/floating/mod.rs

@@ -18,6 +18,7 @@ use crate::{
     },
     state::State,
     utils::prelude::*,
+    wayland::protocols::toplevel_info::ToplevelInfoState,
 };
 
 mod grabs;
@@ -45,9 +46,30 @@ impl FloatingLayout {
         self.space.map_output(output, location)
     }
 
-    pub fn unmap_output(&mut self, output: &Output) {
+    pub fn unmap_output(
+        &mut self,
+        output: &Output,
+        toplevel_info: &mut ToplevelInfoState<State, CosmicSurface>,
+    ) {
+        let windows = self
+            .space
+            .elements_for_output(output)
+            .cloned()
+            .collect::<Vec<_>>();
+        for window in &windows {
+            for (toplevel, _) in window.windows() {
+                toplevel_info.toplevel_leave_output(&toplevel, output);
+            }
+        }
         self.space.unmap_output(output);
         self.refresh();
+        for window in &windows {
+            for output in self.space.outputs_for_element(&window) {
+                for (toplevel, _) in window.windows() {
+                    toplevel_info.toplevel_enter_output(&toplevel, &output);
+                }
+            }
+        }
     }
 
     pub fn map(

src/shell/layout/tiling/mod.rs

@@ -13,7 +13,9 @@ use crate::{
         CosmicSurface, OutputNotMapped,
     },
     utils::prelude::*,
-    wayland::handlers::xdg_shell::popup::get_popup_toplevel,
+    wayland::{
+        handlers::xdg_shell::popup::get_popup_toplevel, protocols::toplevel_info::ToplevelInfoState,
+    },
 };
 
 use id_tree::{InsertBehavior, MoveBehavior, Node, NodeId, NodeIdError, RemoveBehavior, Tree};
@@ -275,14 +277,37 @@ impl TilingLayout {
         }
     }
 
-    pub fn unmap_output(&mut self, output: &Output) {
+    pub fn unmap_output(
+        &mut self,
+        output: &Output,
+        toplevel_info: &mut ToplevelInfoState<State, CosmicSurface>,
+    ) {
         if let Some(src) = self.trees.remove(output) {
             // TODO: expects last remaining output
-            let Some((output, dst)) = self.trees.iter_mut().next() else { return; };
+            let Some((new_output, dst)) = self.trees.iter_mut().next() else { return; };
-            let orientation = match output.output.geometry().size {
+            let orientation = match new_output.output.geometry().size {
                 x if x.w >= x.h => Orientation::Vertical,
                 _ => Orientation::Horizontal,
             };
+            for node in src
+                .root_node_id()
+                .and_then(|root_id| src.traverse_pre_order(root_id).ok())
+                .into_iter()
+                .flatten()
+            {
+                if let Data::Mapped {
+                    mapped,
+                    last_geometry: _,
+                } = node.data()
+                {
+                    for (toplevel, _) in mapped.windows() {
+                        toplevel_info.toplevel_leave_output(&toplevel, output);
+                        toplevel_info.toplevel_enter_output(&toplevel, &new_output.output);
+                    }
+                    mapped.output_leave(output);
+                    mapped.output_enter(&new_output.output, mapped.bbox());
+                }
+            }
             TilingLayout::merge_trees(src, dst, orientation);
             self.refresh()
         }

src/shell/mod.rs

@@ -541,7 +541,7 @@ impl Shell {
 
                             // update mapping
                             workspace.map_output(new_output, (0, 0).into());
-                            workspace.unmap_output(output);
+                            workspace.unmap_output(output, &mut self.toplevel_info_state);
                             workspace.refresh();
 
                             new_set.workspaces.push(workspace);
@@ -559,7 +559,7 @@ impl Shell {
             WorkspaceMode::Global(set) => {
                 state.remove_group_output(&set.group, output);
                 for workspace in &mut set.workspaces {
-                    workspace.unmap_output(output);
+                    workspace.unmap_output(output, &mut self.toplevel_info_state);
                     workspace.refresh();
                 }
             }
@@ -1152,6 +1152,13 @@ impl Shell {
                 .shell
                 .toplevel_info_state
                 .toplevel_leave_workspace(&toplevel, &from_workspace.handle);
+            if from_output != to_output {
+                state
+                    .common
+                    .shell
+                    .toplevel_info_state
+                    .toplevel_leave_output(&toplevel, from_output);
+            }
         }
         let elements = from_workspace.mapped().cloned().collect::<Vec<_>>();
         std::mem::drop(from_workspace);
@@ -1180,6 +1187,13 @@ impl Shell {
                 .map(mapped.clone(), &seat, focus_stack.iter());
         }
         for (toplevel, _) in mapped.windows() {
+            if from_output != to_output {
+                state
+                    .common
+                    .shell
+                    .toplevel_info_state
+                    .toplevel_enter_output(&toplevel, to_output);
+            }
             state
                 .common
                 .shell

src/shell/workspace.rs

@@ -10,6 +10,7 @@ use crate::{
         handlers::screencopy::DropableSession,
         protocols::{
             screencopy::{BufferParams, Session as ScreencopySession},
+            toplevel_info::ToplevelInfoState,
             workspace::WorkspaceHandle,
         },
     },
@@ -95,12 +96,16 @@ impl Workspace {
         self.floating_layer.map_output(output, position);
     }
 
-    pub fn unmap_output(&mut self, output: &Output) {
+    pub fn unmap_output(
+        &mut self,
+        output: &Output,
+        toplevel_info: &mut ToplevelInfoState<State, CosmicSurface>,
+    ) {
         if let Some(dead_output_window) = self.fullscreen.remove(output) {
             self.unfullscreen_request(&dead_output_window);
         }
-        self.tiling_layer.unmap_output(output);
+        self.tiling_layer.unmap_output(output, toplevel_info);
-        self.floating_layer.unmap_output(output);
+        self.floating_layer.unmap_output(output, toplevel_info);
         self.refresh();
     }