Improve master stack layout and tags

Cargo.toml

@@ -20,6 +20,7 @@ rmp-serde = { version = "1.1.1" }
 calloop = { version = "0.10.1", features = ["executor", "futures-io"] }
 futures-lite = { version = "1.13.0" }
 async-process = { version = "1.7.0" }
+itertools = { version = "0.11.0" }
 
 [features]
 default = ["egl", "winit", "udev"]

src/handlers.rs

@@ -247,7 +247,7 @@ impl<B: Backend> XdgShellHandler for State<B> {
         });
 
         self.windows.push(window.clone());
-        self.space.map_element(window.clone(), (0, 0), true);
+        // self.space.map_element(window.clone(), (0, 0), true);
         let clone = window.clone();
         self.loop_handle.insert_idle(move |data| {
             data.state
@@ -264,20 +264,25 @@ impl<B: Backend> XdgShellHandler for State<B> {
         self.loop_handle.insert_idle(move |data| {
             if let Some(focused_output) = &data.state.focus_state.focused_output {
                 OutputState::with(focused_output, |state| {
-                    if let Some(id) = state.focused_tags.iter().next() {
-                        // TODO: make it work with more than one active tag
-                        let tag = data
-                            .state
-                            .tag_state
-                            .tags
-                            .iter_mut()
-                            .find(|tag| &tag.id == id)
-                            .unwrap();
-                        tag.windows.as_master_stack().add(
-                            &data.state.space,
-                            focused_output,
-                            window.clone(),
-                        );
+                    let window = window.clone();
+                    let mut tags = data
+                        .state
+                        .tag_state
+                        .tags
+                        .iter_mut()
+                        .filter(|tg| state.focused_tags.contains(&tg.id));
+
+                    if let Some(first) = tags.next() {
+                        let mut layout = first.windows.as_master_stack();
+
+                        for tg in tags {
+                            layout = layout.chain_with(&mut tg.windows);
+                        }
+
+                        layout.add(&data.state.space, focused_output, window);
+                    }
+                    for tag in data.state.tag_state.tags.iter() {
+                        tracing::debug!("tag {:?}, {}", tag.id, tag.windows.len());
                     }
                 });
             }
@@ -293,20 +298,36 @@ impl<B: Backend> XdgShellHandler for State<B> {
             .unwrap();
         if let Some(focused_output) = self.focus_state.focused_output.as_ref() {
             OutputState::with(focused_output, |state| {
-                if let Some(id) = state.focused_tags.iter().next() {
-                    // TODO: make it work with more than one active tag
-                    let tag = self
-                        .tag_state
-                        .tags
-                        .iter_mut()
-                        .find(|tag| &tag.id == id)
-                        .unwrap();
-                    tag.windows
-                        .as_master_stack()
-                        .remove(&self.space, focused_output, window);
+                let mut tags = self
+                    .tag_state
+                    .tags
+                    .iter_mut()
+                    .filter(|tg| state.focused_tags.contains(&tg.id));
+
+                if let Some(first) = tags.next() {
+                    tracing::debug!("first tag: {:?}", first.id);
+                    let mut layout = first.windows.as_master_stack();
+
+                    for tg in tags {
+                        tracing::debug!("tag: {:?}", tg.id);
+                        layout = layout.chain_with(&mut tg.windows);
+                    }
+
+                    // This will only remove the window from focused tags...
+                    layout.remove(&self.space, focused_output, window);
+                }
+
+                // ...so here we remove the window from any tag that isn't focused
+                for tag in self.tag_state.tags.iter_mut() {
+                    tag.windows.retain(|el| el != window);
                 }
             });
         }
+
+        for tag in self.tag_state.tags.iter() {
+            tracing::debug!("tag {:?}, {}", tag.id, tag.windows.len());
+        }
+
         self.windows.retain(|window| window.toplevel() != &surface);
         // let mut windows: Vec<Window> = self.space.elements().cloned().collect();
         // windows.retain(|window| window.toplevel() != &surface);

src/layout.rs

@@ -4,6 +4,7 @@
 //
 // SPDX-License-Identifier: MPL-2.0
 
+use itertools::Itertools;
 use smithay::{
     desktop::{space::SpaceElement, Space, Window},
     output::Output,
@@ -25,10 +26,10 @@ pub enum Direction {
 }
 
 pub struct MasterStack<'a, S: SpaceElement> {
-    inner: &'a mut Vec<S>,
+    inner: Vec<&'a mut Vec<S>>,
 }
 
-pub trait Layout<S: SpaceElement> {
+pub trait Layout<'a, S: SpaceElement> {
     /// Add a [`SpaceElement`] to this layout and update positions.
     fn add(&mut self, space: &Space<S>, output: &Output, elem: S);
     /// Remove a [`SpaceElement`] from this layout and update positions.
@@ -40,15 +41,21 @@ pub trait Layout<S: SpaceElement> {
 
     /// Perform a full layout with all elements. Use this when you are switching from another layout.
     fn layout(&self, space: &Space<S>, output: &Output);
+
+    fn chain_with(self, vec: &'a mut Vec<S>) -> Self;
 }
 
-impl<S: SpaceElement + Eq> MasterStack<'_, S> {
-    pub fn master(&self) -> Option<&S> {
-        self.inner.first()
+impl MasterStack<'_, Window> {
+    pub fn master(&self) -> Option<&Window> {
+        self.inner.iter().flat_map(|vec| vec.iter()).next()
     }
 
-    pub fn stack(&self) -> impl Iterator<Item = &S> {
-        self.inner.iter().skip(1)
+    pub fn stack(&self) -> impl Iterator<Item = &Window> {
+        self.inner
+            .iter()
+            .flat_map(|vec| vec.iter())
+            .unique()
+            .skip(1)
     }
 }
 
@@ -104,9 +111,11 @@ pub fn swap_window_positions(space: &Space<Window>, win1: &Window, win2: &Window
     });
 }
 
-impl Layout<Window> for MasterStack<'_, Window> {
+impl<'a> Layout<'a, Window> for MasterStack<'a, Window> {
     fn add(&mut self, space: &Space<Window>, output: &Output, elem: Window) {
-        self.inner.push(elem);
+        for vec in self.inner.iter_mut() {
+            vec.push(elem.clone());
+        }
 
         if self.stack().count() == 0 {
             let Some(master) = self.master() else { unreachable!() };
@@ -147,7 +156,9 @@ impl Layout<Window> for MasterStack<'_, Window> {
     }
 
     fn remove(&mut self, space: &Space<Window>, output: &Output, elem: &Window) {
-        self.inner.retain(|el| el != elem);
+        for vec in self.inner.iter_mut() {
+            vec.retain(|el| el != elem);
+        }
 
         let Some(master) = self.master() else { return };
 
@@ -173,7 +184,7 @@ impl Layout<Window> for MasterStack<'_, Window> {
     }
 
     fn swap(&mut self, space: &Space<Window>, elem1: &Window, elem2: &Window) {
-        let mut elems = self.inner.iter_mut();
+        let mut elems = self.inner.iter_mut().flat_map(|vec| vec.iter_mut());
         let first = elems.find(|elem| *elem == elem1);
         let second = elems.find(|elem| *elem == elem2);
         if let Some(first) = first {
@@ -223,6 +234,12 @@ impl Layout<Window> for MasterStack<'_, Window> {
             self.layout_stack(space, output);
         }
     }
+
+    /// Chain another tag's windows to this one to be layed out.
+    fn chain_with(mut self, vec: &'a mut Vec<Window>) -> Self {
+        self.inner.push(vec);
+        self
+    }
 }
 
 pub trait LayoutVec<S: SpaceElement> {
@@ -233,6 +250,6 @@ pub trait LayoutVec<S: SpaceElement> {
 
 impl<S: SpaceElement> LayoutVec<S> for Vec<S> {
     fn as_master_stack(&mut self) -> MasterStack<S> {
-        MasterStack { inner: self }
+        MasterStack { inner: vec![self] }
     }
 }