diff --git a/src/shell/mod.rs b/src/shell/mod.rs
index 1206b338..97e2e9de 100644
--- a/src/shell/mod.rs
+++ b/src/shell/mod.rs
@@ -4807,11 +4807,18 @@ impl Shell {
         {
             let mut from = set.sticky_layer.element_geometry(&mapped).unwrap();
             let mut was_maximized = false;
+            let mut restore_state = None;
             window = if let Some(stack) = mapped.stack_ref()
                 && stack.len() > 1
             {
-                let surface = stack.surfaces().find(|s| s == surface).unwrap();
-                stack.remove_window(&surface);
+                let idx = stack.surfaces().position(|s| &s == surface)?;
+                let surface = stack.remove_idx(idx)?;
+                restore_state = Some(FullscreenRestoreState::Stack {
+                    state: StackRestoreData {
+                        stack: mapped.key(),
+                        idx,
+                    },
+                });
                 surface
             } else {
                 // Must be set before `map_internal`/`unmap` below, as both may call
@@ -4844,7 +4851,7 @@ impl Shell {
             workspace.map_fullscreen(
                 &window,
                 &seat,
-                Some(FullscreenRestoreState::Sticky {
+                Some(restore_state.unwrap_or(FullscreenRestoreState::Sticky {
                     output: old_output,
                     state: FloatingRestoreData {
                         geometry: from,
@@ -4852,7 +4859,7 @@ impl Shell {
                         was_maximized,
                         was_snapped: None,
                     },
-                }),
+                })),
                 Some(from),
             );
         } else if let Some(workspace) = self.space_for_mut(&mapped) {