Refactor everything, add initial tiling code

2 files changed, 176 insertions, 114 deletions

diff --git a/src/handlers/compositor.rs b/src/handlers/compositor.rs
index f5a955e1..d9e25b8f 100644
--- a/src/handlers/compositor.rs
+++ b/src/handlers/compositor.rs
@@ -1,4 +1,7 @@
-use smithay::backend::renderer::utils::on_commit_buffer_handler;
+use std::collections::hash_map::Entry;
+
+use smithay::backend::renderer::utils::{on_commit_buffer_handler, with_renderer_surface_state};
+use smithay::desktop::find_popup_root_surface;
 use smithay::reexports::wayland_server::protocol::wl_buffer;
 use smithay::reexports::wayland_server::protocol::wl_surface::WlSurface;
 use smithay::reexports::wayland_server::Client;
@@ -9,6 +12,7 @@ use smithay::wayland::compositor::{
 use smithay::wayland::shm::{ShmHandler, ShmState};
 use smithay::{delegate_compositor, delegate_shm};
 
+use super::xdg_shell;
 use crate::grabs::resize_grab;
 use crate::niri::ClientState;
 use crate::Niri;
@@ -28,24 +32,95 @@ impl CompositorHandler for Niri {
             .message("client commit", 0);
 
         on_commit_buffer_handler::<Self>(surface);
-        if !is_sync_subsurface(surface) {
-            let mut root = surface.clone();
-            while let Some(parent) = get_parent(&root) {
-                root = parent;
+
+        if is_sync_subsurface(surface) {
+            return;
+        }
+
+        let mut root_surface = surface.clone();
+        while let Some(parent) = get_parent(&root_surface) {
+            root_surface = parent;
+        }
+
+        if surface == &root_surface {
+            // This is a root surface commit. It might have mapped a previously-unmapped toplevel.
+            if let Entry::Occupied(entry) = self.unmapped_windows.entry(surface.clone()) {
+                let is_mapped =
+                    with_renderer_surface_state(surface, |state| state.buffer().is_some());
+
+                if is_mapped {
+                    // The toplevel got mapped.
+                    let window = entry.remove();
+                    window.on_commit();
+
+                    let output = self.monitor_set.active_output().unwrap().clone();
+                    self.monitor_set.add_window_to_output(&output, window, true);
+                    self.update_focus();
+
+                    self.queue_redraw(output);
+                    return;
+                }
+
+                // The toplevel remains unmapped.
+                let window = entry.get();
+                xdg_shell::send_initial_configure_if_needed(window);
+                return;
             }
-            if let Some(window) = self
-                .space
-                .elements()
-                .find(|w| w.toplevel().wl_surface() == &root)
-            {
+
+            // This is a commit of a previously-mapped root or a non-toplevel root.
+            if let Some((window, space)) = self.monitor_set.find_window_and_space(surface) {
+                // This is a commit of a previously-mapped toplevel.
+                let output = space.outputs().next().unwrap().clone();
+
                 window.on_commit();
+
+                // This is a commit of a previously-mapped toplevel.
+                let is_mapped =
+                    with_renderer_surface_state(surface, |state| state.buffer().is_some());
+
+                if !is_mapped {
+                    // The toplevel got unmapped.
+                    let window = window.clone();
+                    self.monitor_set.remove_window(&window);
+                    self.unmapped_windows.insert(surface.clone(), window);
+                    self.update_focus();
+
+                    self.queue_redraw(output);
+                    return;
+                }
+
+                // The toplevel remains mapped.
+                resize_grab::handle_commit(&window);
+                self.monitor_set.update_window(&window);
+
+                self.queue_redraw(output);
+                return;
             }
-        };
 
-        self.xdg_handle_commit(surface);
-        resize_grab::handle_commit(&mut self.space, surface);
+            // This is a commit of a non-toplevel root.
+        }
 
-        self.queue_redraw();
+        // This is a commit of a non-root or a non-toplevel root.
+        let root_window_space = self.monitor_set.find_window_and_space(&root_surface);
+        if let Some((window, space)) = root_window_space {
+            let output = space.outputs().next().unwrap().clone();
+            window.on_commit();
+            self.monitor_set.update_window(&window);
+            self.queue_redraw(output);
+            return;
+        }
+
+        // This might be a popup.
+        self.popups_handle_commit(surface);
+        if let Some(popup) = self.popups.find_popup(surface) {
+            if let Ok(root) = find_popup_root_surface(&popup) {
+                let root_window_space = self.monitor_set.find_window_and_space(&root);
+                if let Some((_window, space)) = root_window_space {
+                    let output = space.outputs().next().unwrap().clone();
+                    self.queue_redraw(output);
+                }
+            }
+        }
     }
 }
 
diff --git a/src/handlers/xdg_shell.rs b/src/handlers/xdg_shell.rs
index 0e1ad760..c7df18c7 100644
--- a/src/handlers/xdg_shell.rs
+++ b/src/handlers/xdg_shell.rs
@@ -2,20 +2,19 @@ use smithay::delegate_xdg_shell;
 use smithay::desktop::{PopupKind, Window};
 use smithay::input::pointer::{Focus, GrabStartData as PointerGrabStartData};
 use smithay::input::Seat;
-use smithay::output::Output;
 use smithay::reexports::wayland_protocols::xdg::shell::server::xdg_toplevel;
 use smithay::reexports::wayland_server::protocol::wl_surface::WlSurface;
 use smithay::reexports::wayland_server::protocol::{wl_output, wl_seat};
 use smithay::reexports::wayland_server::Resource;
 use smithay::utils::{Rectangle, Serial};
 use smithay::wayland::compositor::with_states;
-use smithay::wayland::seat::WaylandFocus;
 use smithay::wayland::shell::xdg::{
     PopupSurface, PositionerState, ToplevelSurface, XdgPopupSurfaceData, XdgShellHandler,
     XdgShellState, XdgToplevelSurfaceData,
 };
 
 use crate::grabs::{MoveSurfaceGrab, ResizeSurfaceGrab};
+use crate::layout::MonitorSet;
 use crate::Niri;
 
 impl XdgShellHandler for Niri {
@@ -24,8 +23,16 @@ impl XdgShellHandler for Niri {
     }
 
     fn new_toplevel(&mut self, surface: ToplevelSurface) {
+        let wl_surface = surface.wl_surface().clone();
         let window = Window::new(surface);
-        self.space.map_element(window, (0, 0), false);
+
+        // Tell the surface the preferred size and bounds for its likely output.
+        let output = self.monitor_set.active_output().unwrap();
+        MonitorSet::configure_new_window(output, &window);
+
+        // At the moment of creation, xdg toplevels must have no buffer.
+        let existing = self.unmapped_windows.insert(wl_surface, window);
+        assert!(existing.is_none());
     }
 
     fn new_popup(&mut self, surface: PopupSurface, positioner: PositionerState) {
@@ -49,17 +56,12 @@ impl XdgShellHandler for Niri {
         if let Some(start_data) = check_grab(&seat, wl_surface, serial) {
             let pointer = seat.get_pointer().unwrap();
 
-            let window = self
-                .space
-                .elements()
-                .find(|w| w.toplevel().wl_surface() == wl_surface)
-                .unwrap()
-                .clone();
-            let initial_window_location = self.space.element_location(&window).unwrap();
+            let (window, space) = self.monitor_set.find_window_and_space(wl_surface).unwrap();
+            let initial_window_location = space.element_location(&window).unwrap();
 
             let grab = MoveSurfaceGrab {
                 start_data,
-                window,
+                window: window.clone(),
                 initial_window_location,
             };
 
@@ -81,13 +83,8 @@ impl XdgShellHandler for Niri {
         if let Some(start_data) = check_grab(&seat, wl_surface, serial) {
             let pointer = seat.get_pointer().unwrap();
 
-            let window = self
-                .space
-                .elements()
-                .find(|w| w.toplevel().wl_surface() == wl_surface)
-                .unwrap()
-                .clone();
-            let initial_window_location = self.space.element_location(&window).unwrap();
+            let (window, space) = self.monitor_set.find_window_and_space(wl_surface).unwrap();
+            let initial_window_location = space.element_location(&window).unwrap();
             let initial_window_size = window.geometry().size;
 
             surface.with_pending_state(|state| {
@@ -98,7 +95,7 @@ impl XdgShellHandler for Niri {
 
             let grab = ResizeSurfaceGrab::start(
                 start_data,
-                window,
+                window.clone(),
                 edges.into(),
                 Rectangle::from_loc_and_size(initial_window_location, initial_window_size),
             );
@@ -126,7 +123,7 @@ impl XdgShellHandler for Niri {
     }
 
     fn grab(&mut self, _surface: PopupSurface, _seat: wl_seat::WlSeat, _serial: Serial) {
-        // TODO popup grabs
+        // FIXME popup grabs
     }
 
     fn maximize_request(&mut self, surface: ToplevelSurface) {
@@ -135,23 +132,17 @@ impl XdgShellHandler for Niri {
             .capabilities
             .contains(xdg_toplevel::WmCapabilities::Maximize)
         {
-            let wl_surface = surface.wl_surface();
-            let window = self
-                .space
-                .elements()
-                .find(|w| w.toplevel().wl_surface() == wl_surface)
-                .unwrap()
-                .clone();
-            let geometry = self
-                .space
-                .output_geometry(self.output.as_ref().unwrap())
-                .unwrap();
-
-            surface.with_pending_state(|state| {
-                state.states.set(xdg_toplevel::State::Maximized);
-                state.size = Some(geometry.size);
-            });
-            self.space.map_element(window, geometry.loc, true);
+            // let wl_surface = surface.wl_surface();
+            // let (window, space) = self.monitor_set.find_window_and_space(wl_surface).unwrap();
+            // let geometry = space
+            //     .output_geometry(space.outputs().next().unwrap())
+            //     .unwrap();
+
+            // surface.with_pending_state(|state| {
+            //     state.states.set(xdg_toplevel::State::Maximized);
+            //     state.size = Some(geometry.size);
+            // });
+            // space.map_element(window.clone(), geometry.loc, true);
         }
 
         // The protocol demands us to always reply with a configure,
@@ -185,44 +176,32 @@ impl XdgShellHandler for Niri {
             .capabilities
             .contains(xdg_toplevel::WmCapabilities::Fullscreen)
         {
-            // NOTE: This is only one part of the solution. We can set the
-            // location and configure size here, but the surface should be rendered fullscreen
-            // independently from its buffer size
-            let wl_surface = surface.wl_surface();
-
-            let output = wl_output
-                .as_ref()
-                .and_then(Output::from_resource)
-                .or_else(|| {
-                    let w = self
-                        .space
-                        .elements()
-                        .find(|window| {
-                            window
-                                .wl_surface()
-                                .map(|s| s == *wl_surface)
-                                .unwrap_or(false)
-                        })
-                        .cloned();
-                    w.and_then(|w| self.space.outputs_for_element(&w).get(0).cloned())
-                });
-
-            if let Some(output) = output {
-                let geometry = self.space.output_geometry(&output).unwrap();
-
-                surface.with_pending_state(|state| {
-                    state.states.set(xdg_toplevel::State::Fullscreen);
-                    state.size = Some(geometry.size);
-                });
-
-                let window = self
-                    .space
-                    .elements()
-                    .find(|w| w.toplevel().wl_surface() == wl_surface)
-                    .unwrap()
-                    .clone();
-                self.space.map_element(window, geometry.loc, true);
-            }
+            // // NOTE: This is only one part of the solution. We can set the
+            // // location and configure size here, but the surface should be rendered fullscreen
+            // // independently from its buffer size
+            // let wl_surface = surface.wl_surface();
+
+            // let output = wl_output
+            //     .as_ref()
+            //     .and_then(Output::from_resource)
+            //     .or_else(|| {
+            //         self.monitor_set
+            //             .find_window_and_space(wl_surface)
+            //             .and_then(|(_window, space)| space.outputs().next().cloned())
+            //     });
+
+            // if let Some(output) = output {
+            //     let (window, space) =
+            // self.monitor_set.find_window_and_space(wl_surface).unwrap();
+            //     let geometry = space.output_geometry(&output).unwrap();
+
+            //     surface.with_pending_state(|state| {
+            //         state.states.set(xdg_toplevel::State::Fullscreen);
+            //         state.size = Some(geometry.size);
+            //     });
+
+            //     space.map_element(window.clone(), geometry.loc, true);
+            // }
         }
 
         // The protocol demands us to always reply with a configure,
@@ -247,12 +226,25 @@ impl XdgShellHandler for Niri {
         surface.send_pending_configure();
     }
 
-    fn toplevel_destroyed(&mut self, _surface: ToplevelSurface) {
-        self.queue_redraw();
+    fn toplevel_destroyed(&mut self, surface: ToplevelSurface) {
+        if self.unmapped_windows.remove(surface.wl_surface()).is_some() {
+            // An unmapped toplevel got destroyed.
+            return;
+        }
+
+        let (window, space) = self
+            .monitor_set
+            .find_window_and_space(surface.wl_surface())
+            .unwrap();
+        let output = space.outputs().next().unwrap().clone();
+        self.monitor_set.remove_window(&window);
+        self.update_focus();
+        self.queue_redraw(output);
     }
 
     fn popup_destroyed(&mut self, _surface: PopupSurface) {
-        self.queue_redraw();
+        // FIXME granular
+        self.queue_redraw_all();
     }
 }
 
@@ -282,32 +274,27 @@ fn check_grab(
     Some(start_data)
 }
 
+pub fn send_initial_configure_if_needed(window: &Window) {
+    let initial_configure_sent = with_states(window.toplevel().wl_surface(), |states| {
+        states
+            .data_map
+            .get::<XdgToplevelSurfaceData>()
+            .unwrap()
+            .lock()
+            .unwrap()
+            .initial_configure_sent
+    });
+
+    if !initial_configure_sent {
+        window.toplevel().send_configure();
+    }
+}
+
 impl Niri {
     /// Should be called on `WlSurface::commit`
-    pub fn xdg_handle_commit(&mut self, surface: &WlSurface) {
+    pub fn popups_handle_commit(&mut self, surface: &WlSurface) {
         self.popups.commit(surface);
 
-        if let Some(window) = self
-            .space
-            .elements()
-            .find(|w| w.toplevel().wl_surface() == surface)
-            .cloned()
-        {
-            let initial_configure_sent = with_states(surface, |states| {
-                states
-                    .data_map
-                    .get::<XdgToplevelSurfaceData>()
-                    .unwrap()
-                    .lock()
-                    .unwrap()
-                    .initial_configure_sent
-            });
-
-            if !initial_configure_sent {
-                window.toplevel().send_configure();
-            }
-        }
-
         if let Some(popup) = self.popups.find_popup(surface) {
             let PopupKind::Xdg(ref popup) = popup;
             let initial_configure_sent = with_states(surface, |states| {