shell: restore focus

src/input/mod.rs

@@ -268,7 +268,9 @@ impl Common {
                                                 let current_output = active_output(seat, &self);
                                                 let workspace =
                                                     self.shell.active_space_mut(&current_output);
-                                                if let Some(window) = workspace.focus_stack.last() {
+                                                if let Some(window) =
+                                                    workspace.focus_stack(seat).last()
+                                                {
                                                     #[allow(irrefutable_let_patterns)]
                                                     if let Kind::Xdg(xdg) = &window.toplevel() {
                                                         xdg.send_close();
@@ -536,7 +538,7 @@ impl Common {
                                         }
                                     };
 
-                                    self.shell.set_focus(under.as_ref(), seat);
+                                    self.set_focus(under.as_ref(), seat, None);
                                 }
                                 wl_pointer::ButtonState::Pressed
                             }

src/main.rs

@@ -57,6 +57,7 @@ fn main() -> Result<()> {
         display.borrow_mut().flush_clients(state);
         // trigger routines
         state.common.shell.refresh();
+        state.common.refresh_focus();
     })?;
 
     let _log = state.destroy();

src/shell/handler.rs

@@ -4,8 +4,8 @@ use crate::{config::Config, input::active_output, state::State, utils::SurfaceDr
 use smithay::{
     backend::renderer::utils::on_commit_buffer_handler,
     desktop::{
-        layer_map_for_output, Kind, LayerSurface, PopupKeyboardGrab, PopupKind, PopupPointerGrab,
+        layer_map_for_output, Kind, LayerSurface, PopupGrab, PopupKeyboardGrab, PopupKind,
-        PopupUngrabStrategy, Window,
+        PopupPointerGrab, PopupUngrabStrategy, Window,
     },
     reexports::{
         wayland_protocols::xdg_shell::server::xdg_toplevel,
@@ -28,7 +28,9 @@ use smithay::{
         Serial,
     },
 };
-use std::{cell::Cell, rc::Rc, sync::Mutex};
+use std::{cell::Cell, sync::Mutex};
+
+pub type PopupGrabData = Cell<Option<PopupGrab>>;
 
 pub fn init_shell(config: &Config, display: &mut Display) -> super::Shell {
     compositor_init(
@@ -36,14 +38,15 @@ pub fn init_shell(config: &Config, display: &mut Display) -> super::Shell {
         move |surface, mut ddata| {
             on_commit_buffer_handler(&surface);
             let state = ddata.get::<State>().unwrap();
-            state.common.shell.commit(&surface);
+            state
+                .common
+                .shell
+                .commit(&surface, state.common.seats.iter());
             commit(&surface, state)
         },
         None,
     );
 
-    let popup_grab = Rc::new(Cell::new(None));
-    let popup_grab_clone = popup_grab.clone();
     let (_xdg_shell_state, _xdg_global) = xdg_shell_init(
         display,
         move |event, mut ddata| {
@@ -186,7 +189,7 @@ pub fn init_shell(config: &Config, display: &mut Display) -> super::Shell {
                                 grab.ungrab(PopupUngrabStrategy::All);
                                 return;
                             }
-                            keyboard.set_focus(grab.current_grab().as_ref(), serial);
+                            state.set_focus(grab.current_grab().as_ref(), &seat, Some(serial));
                             keyboard.set_grab(PopupKeyboardGrab::new(&grab), serial);
                         }
 
@@ -203,7 +206,12 @@ pub fn init_shell(config: &Config, display: &mut Display) -> super::Shell {
                             pointer.set_grab(PopupPointerGrab::new(&grab), serial, 0);
                         }
 
-                        popup_grab_clone.set(Some(grab));
+                        seat.user_data()
+                            .insert_if_missing(|| PopupGrabData::new(None));
+                        seat.user_data()
+                            .get::<PopupGrabData>()
+                            .unwrap()
+                            .set(Some(grab));
                     }
                 }
                 _ => { /*TODO*/ }
@@ -222,11 +230,11 @@ pub fn init_shell(config: &Config, display: &mut Display) -> super::Shell {
                 ..
             } => {
                 let state = &mut ddata.get::<State>().unwrap().common;
-                let seat = &state.last_active_seat;
+                let seat = state.last_active_seat.clone();
                 let output = wl_output
                     .as_ref()
                     .and_then(Output::from_resource)
-                    .unwrap_or_else(|| active_output(seat, &*state));
+                    .unwrap_or_else(|| active_output(&seat, &*state));
 
                 let focus = surface
                     .get_surface()
@@ -245,7 +253,7 @@ pub fn init_shell(config: &Config, display: &mut Display) -> super::Shell {
                     .unwrap();
 
                 if focus {
-                    state.shell.set_focus(surface.get_surface(), seat);
+                    state.set_focus(surface.get_surface(), &seat, None);
                 }
             }
             _ => {}
@@ -253,7 +261,7 @@ pub fn init_shell(config: &Config, display: &mut Display) -> super::Shell {
         None,
     );
 
-    super::Shell::new(config, popup_grab)
+    super::Shell::new(config)
 }
 
 fn check_grab_preconditions(

src/shell/layout/combined.rs

@@ -63,7 +63,16 @@ impl<A: Layout, B: Layout> Layout for Combined<A, B> {
         self.windows_a.retain(|w| w.toplevel().alive());
         self.windows_b.retain(|w| w.toplevel().alive());
     }
-    //fn unmap_window(&mut self, space: &mut Space, window: &Window);
+
+    fn unmap_window(&mut self, space: &mut Space, window: &Window) {
+        if self.windows_a.contains(window) {
+            self.windows_a.retain(|w| w != window);
+            self.first.unmap_window(space, window)
+        } else if self.windows_b.contains(window) {
+            self.windows_b.retain(|w| w != window);
+            self.second.unmap_window(space, window)
+        }
+    }
 
     fn move_request(
         &mut self,

src/shell/mod.rs

@@ -1,20 +1,27 @@
 // SPDX-License-Identifier: GPL-3.0-only
 
-use crate::config::Config;
+use crate::{config::Config, input::active_output, state::Common};
 pub use smithay::{
     desktop::{PopupGrab, PopupManager, PopupUngrabStrategy, Space, Window},
     reexports::wayland_server::protocol::wl_surface::WlSurface,
     utils::{Logical, Point, Rectangle, Size},
-    wayland::{output::Output, seat::Seat, SERIAL_COUNTER},
+    wayland::{
+        compositor::with_states, output::Output, seat::Seat,
+        shell::xdg::XdgToplevelSurfaceRoleAttributes, Serial, SERIAL_COUNTER,
+    },
+};
+use std::{
+    cell::{Cell, RefCell},
+    mem::MaybeUninit,
+    sync::Mutex,
 };
-use std::{cell::Cell, mem::MaybeUninit, rc::Rc};
 
 pub const MAX_WORKSPACES: usize = 10; // TODO?
 
 mod handler;
 pub mod layout;
 mod workspace;
-pub use self::handler::init_shell;
+pub use self::handler::{init_shell, PopupGrabData};
 pub use self::layout::Layout;
 pub use self::workspace::*;
 
@@ -55,7 +62,6 @@ impl Mode {
 
 pub struct Shell {
     popups: PopupManager,
-    popup_grab: Rc<Cell<Option<PopupGrab>>>,
     mode: Mode,
     outputs: Vec<Output>,
     pub spaces: [Workspace; MAX_WORKSPACES],
@@ -64,10 +70,9 @@ pub struct Shell {
 const UNINIT_SPACE: MaybeUninit<Workspace> = MaybeUninit::uninit();
 
 impl Shell {
-    fn new(config: &Config, popup_grab: Rc<Cell<Option<PopupGrab>>>) -> Self {
+    fn new(config: &Config) -> Self {
         Shell {
             popups: PopupManager::new(None),
-            popup_grab,
             mode: config.workspace_mode,
             outputs: Vec::new(),
             spaces: unsafe {
@@ -348,7 +353,7 @@ impl Shell {
         }
     }
 
-    pub fn commit(&mut self, surface: &WlSurface) {
+    pub fn commit<'a>(&mut self, surface: &WlSurface, seats: impl Iterator<Item = &'a Seat>) {
         let mut new_focus = None;
         for (idx, workspace) in self.spaces.iter_mut().enumerate() {
             if let Some((window, seat)) = workspace
@@ -376,50 +381,11 @@ impl Shell {
             workspace.space.commit(surface)
         }
         if let Some(seat) = new_focus {
-            self.set_focus(Some(surface), &seat)
+            self.set_focus(Some(surface), &seat, seats, None)
         }
         self.popups.commit(&surface);
     }
 
-    pub fn set_focus(&mut self, surface: Option<&WlSurface>, active_seat: &Seat) {
-        // update FocusStack and notify layouts about new focus (if any window)
-        if let Some(surface) = surface {
-            if let Some(workspace) = self.space_for_surface_mut(surface) {
-                if let Some(window) = workspace.space.window_for_surface(surface) {
-                    if Some(window) != workspace.focus_stack.last().as_ref() {
-                        slog_scope::debug!("Focusing window: {:?}", window);
-                        workspace.focus_stack.append(window);
-                        // also remove popup grabs, if we are switching focus
-                        if let Some(mut popup_grab) = self.popup_grab.take() {
-                            if !popup_grab.has_ended() {
-                                popup_grab.ungrab(PopupUngrabStrategy::All);
-                            }
-                        }
-                    }
-                }
-            }
-        }
-
-        // update keyboard focus
-        if let Some(keyboard) = active_seat.get_keyboard() {
-            keyboard.set_focus(surface, SERIAL_COUNTER.next_serial());
-        }
-
-        // update activate status
-        let focused_windows = self
-            .outputs
-            .iter()
-            .flat_map(|o| self.active_space(o).focus_stack.last())
-            .collect::<Vec<_>>();
-
-        for workspace in &self.spaces {
-            for window in workspace.space.windows() {
-                window.set_activated(focused_windows.contains(window));
-                window.configure();
-            }
-        }
-    }
-
     pub fn set_orientation(
         &mut self,
         seat: &Seat,
@@ -429,4 +395,162 @@ impl Shell {
         self.active_space_mut(output)
             .update_orientation(seat, orientation)
     }
+
+    fn set_focus<'a>(
+        &mut self,
+        surface: Option<&WlSurface>,
+        active_seat: &Seat,
+        seats: impl Iterator<Item = &'a Seat>,
+        serial: Option<Serial>,
+    ) {
+        // update FocusStack and notify layouts about new focus (if any window)
+        if let Some(surface) = surface {
+            if let Some(workspace) = self.space_for_surface_mut(surface) {
+                if let Some(window) = workspace.space.window_for_surface(surface) {
+                    let mut focus_stack = workspace.focus_stack_mut(active_seat);
+                    if Some(window) != focus_stack.last().as_ref() {
+                        slog_scope::debug!("Focusing window: {:?}", window);
+                        focus_stack.append(window);
+                        // also remove popup grabs, if we are switching focus
+                        if let Some(mut popup_grab) = active_seat
+                            .user_data()
+                            .get::<PopupGrabData>()
+                            .and_then(|x| x.take())
+                        {
+                            if !popup_grab.has_ended() {
+                                popup_grab.ungrab(PopupUngrabStrategy::All);
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        // update keyboard focus
+        if let Some(keyboard) = active_seat.get_keyboard() {
+            ActiveFocus::set(active_seat, surface.cloned());
+            keyboard.set_focus(
+                surface,
+                serial.unwrap_or_else(|| SERIAL_COUNTER.next_serial()),
+            );
+        }
+
+        self.update_active(seats)
+    }
+
+    fn update_active<'a>(&self, seats: impl Iterator<Item = &'a Seat>) {
+        // update activate status
+        let focused_windows = seats
+            .flat_map(|seat| {
+                self.outputs
+                    .iter()
+                    .flat_map(|o| self.active_space(o).focus_stack(seat).last().clone())
+            })
+            .collect::<Vec<_>>();
+
+        for workspace in &self.spaces {
+            for window in workspace.space.windows() {
+                window.set_activated(focused_windows.contains(window));
+                window.configure();
+            }
+        }
+    }
+}
+
+pub struct ActiveFocus(RefCell<Option<WlSurface>>);
+
+impl ActiveFocus {
+    fn set(seat: &Seat, surface: Option<WlSurface>) {
+        if !seat
+            .user_data()
+            .insert_if_missing(|| ActiveFocus(RefCell::new(surface.clone())))
+        {
+            *seat
+                .user_data()
+                .get::<ActiveFocus>()
+                .unwrap()
+                .0
+                .borrow_mut() = surface;
+        }
+    }
+
+    fn get(seat: &Seat) -> Option<WlSurface> {
+        seat.user_data()
+            .get::<ActiveFocus>()
+            .and_then(|a| a.0.borrow().clone())
+    }
+}
+
+impl Common {
+    pub fn set_focus(
+        &mut self,
+        surface: Option<&WlSurface>,
+        active_seat: &Seat,
+        serial: Option<Serial>,
+    ) {
+        self.shell
+            .set_focus(surface, active_seat, self.seats.iter(), serial)
+    }
+
+    pub fn refresh_focus(&mut self) {
+        for seat in &self.seats {
+            let mut fixup = false;
+            let output = active_output(seat, &self);
+            let last_known_focus = ActiveFocus::get(seat);
+
+            if let Some(surface) = last_known_focus {
+                if surface.as_ref().is_alive() {
+                    let is_toplevel = with_states(&surface, |states| {
+                        states
+                            .data_map
+                            .get::<Mutex<XdgToplevelSurfaceRoleAttributes>>()
+                            .is_some()
+                    })
+                    .unwrap_or(false);
+                    if !is_toplevel {
+                        continue;
+                    }
+
+                    let workspace = self.shell.active_space(&output);
+                    if let Some(window) = workspace.space.window_for_surface(&surface) {
+                        let focus_stack = workspace.focus_stack(&seat);
+                        if focus_stack.last().map(|w| &w == window).unwrap_or(false) {
+                            continue;
+                        } else {
+                            fixup = true;
+                        }
+                    }
+                } else {
+                    fixup = true;
+                }
+            }
+
+            if fixup {
+                // also remove popup grabs, if we are switching focus
+                if let Some(mut popup_grab) = seat
+                    .user_data()
+                    .get::<PopupGrabData>()
+                    .and_then(|x| x.take())
+                {
+                    if !popup_grab.has_ended() {
+                        popup_grab.ungrab(PopupUngrabStrategy::All);
+                    }
+                }
+
+                // update keyboard focus
+                let surface = self
+                    .shell
+                    .active_space(&output)
+                    .focus_stack(seat)
+                    .last()
+                    .and_then(|w| w.toplevel().get_surface().cloned());
+                if let Some(keyboard) = seat.get_keyboard() {
+                    keyboard.set_focus(surface.as_ref(), SERIAL_COUNTER.next_serial());
+                    ActiveFocus::set(seat, surface);
+                }
+            }
+        }
+
+        self.shell.update_active(self.seats.iter())
+    }
 }

src/shell/workspace.rs

@@ -9,13 +9,26 @@ use smithay::{
         Serial,
     },
 };
-pub struct FocusStack(IndexSet<Window>);
+use std::{
+    cell::{Ref, RefCell, RefMut},
+    collections::HashMap,
+};
 
-impl FocusStack {
+pub struct FocusStack<'a>(Ref<'a, IndexSet<Window>>);
-    pub fn new() -> FocusStack {
+pub struct FocusStackMut<'a>(RefMut<'a, IndexSet<Window>>);
-        FocusStack(IndexSet::new())
+
+impl<'a> FocusStack<'a> {
+    pub fn last(&self) -> Option<Window> {
+        self.0.iter().rev().find(|w| w.toplevel().alive()).cloned()
     }
 
+    pub fn iter<'b>(&'b self) -> Box<dyn Iterator<Item = &'b Window> + 'b> {
+        //working around object-safety constraints for trait Layout
+        Box::new(self.0.iter().rev().filter(|w| w.toplevel().alive()))
+    }
+}
+
+impl<'a> FocusStackMut<'a> {
     pub fn append(&mut self, window: &Window) {
         self.0.retain(|w| w.toplevel().alive());
         self.0.shift_remove(window);
@@ -26,36 +39,68 @@ impl FocusStack {
         self.0.iter().rev().find(|w| w.toplevel().alive()).cloned()
     }
 
-    pub fn iter<'a>(&'a self) -> Box<dyn Iterator<Item = &'a Window> + 'a> {
+    pub fn iter<'b>(&'b self) -> Box<dyn Iterator<Item = &'b Window> + 'b> {
         //working around object-safety constraints for trait Layout
         Box::new(self.0.iter().rev().filter(|w| w.toplevel().alive()))
     }
 }
 
+type FocusStackData = RefCell<(HashMap<u8, IndexSet<Window>>, IndexSet<Window>)>;
+
 pub struct Workspace {
+    idx: u8,
     pub space: Space,
     pub(super) layout: Box<dyn Layout>,
-    pub focus_stack: FocusStack,
     pub(super) pending_windows: Vec<(Window, Seat)>,
 }
 
 impl Workspace {
     pub fn new(idx: u8) -> Workspace {
         Workspace {
+            idx,
             space: Space::new(None),
             layout: layout::new_default_layout(idx),
-            focus_stack: FocusStack::new(),
             pending_windows: Vec::new(),
         }
     }
 
+    pub fn focus_stack<'a>(&'a self, seat: &'a Seat) -> FocusStack<'a> {
+        seat.user_data()
+            .insert_if_missing(|| FocusStackData::new((HashMap::new(), IndexSet::new())));
+        FocusStack(Ref::map(
+            seat.user_data().get::<FocusStackData>().unwrap().borrow(),
+            |map| map.0.get(&self.idx).unwrap_or(&map.1), //TODO: workaround until Ref::filter_map goes stable
+        ))
+    }
+
+    pub fn focus_stack_mut<'a>(&'a self, seat: &'a Seat) -> FocusStackMut<'a> {
+        seat.user_data()
+            .insert_if_missing(|| FocusStackData::new((HashMap::new(), IndexSet::new())));
+        FocusStackMut(RefMut::map(
+            seat.user_data()
+                .get::<FocusStackData>()
+                .unwrap()
+                .borrow_mut(),
+            |map| map.0.entry(self.idx).or_insert_with(|| IndexSet::new()),
+        ))
+    }
+
     pub fn pending_window(&mut self, window: Window, seat: &Seat) {
         self.pending_windows.push((window, seat.clone()));
     }
 
     pub(super) fn map_window<'a>(&mut self, window: &Window, seat: &Seat) {
+        seat.user_data()
+            .insert_if_missing(|| FocusStackData::new((HashMap::new(), IndexSet::new())));
+        let focus_stack = FocusStackMut(RefMut::map(
+            seat.user_data()
+                .get::<FocusStackData>()
+                .unwrap()
+                .borrow_mut(),
+            |map| map.0.entry(self.idx).or_insert_with(|| IndexSet::new()),
+        ));
         self.layout
-            .map_window(&mut self.space, window, seat, self.focus_stack.iter())
+            .map_window(&mut self.space, window, seat, focus_stack.iter())
     }
 
     pub fn refresh(&mut self) {
@@ -64,8 +109,17 @@ impl Workspace {
     }
 
     pub fn update_orientation(&mut self, seat: &Seat, orientation: layout::Orientation) {
+        seat.user_data()
+            .insert_if_missing(|| FocusStackData::new((HashMap::new(), IndexSet::new())));
+        let focus_stack = FocusStackMut(RefMut::map(
+            seat.user_data()
+                .get::<FocusStackData>()
+                .unwrap()
+                .borrow_mut(),
+            |map| map.0.entry(self.idx).or_insert_with(|| IndexSet::new()),
+        ));
         self.layout
-            .update_orientation(orientation, seat, &mut self.space, self.focus_stack.iter())
+            .update_orientation(orientation, seat, &mut self.space, focus_stack.iter())
     }
 
     pub fn maximize_request(&mut self, window: &Window, output: &Output) {