Move PartialEq from LayoutElement to an associated type

7 files changed, 116 insertions, 105 deletions

diff --git a/niri-visual-tests/src/cases/layout.rs b/niri-visual-tests/src/cases/layout.rs
index be611cc5..e6807228 100644
--- a/niri-visual-tests/src/cases/layout.rs
+++ b/niri-visual-tests/src/cases/layout.rs
@@ -2,7 +2,7 @@ use std::collections::HashMap;
 use std::time::Duration;
 
 use niri::layout::workspace::ColumnWidth;
-use niri::layout::Options;
+use niri::layout::{LayoutElement as _, Options};
 use niri::utils::get_monotonic_time;
 use niri_config::Color;
 use smithay::backend::renderer::element::RenderElement;
@@ -73,12 +73,12 @@ impl Layout {
 
         rv.add_window(TestWindow::freeform(0), Some(ColumnWidth::Proportion(0.3)));
         rv.add_window(TestWindow::freeform(1), Some(ColumnWidth::Proportion(0.3)));
-        rv.layout.activate_window(&rv.windows[0]);
+        rv.layout.activate_window(&0);
 
         rv.add_step(500, |l| {
             let win = TestWindow::freeform(2);
             l.add_window(win.clone(), Some(ColumnWidth::Proportion(0.3)));
-            l.layout.start_open_animation_for_window(&win);
+            l.layout.start_open_animation_for_window(win.id());
         });
 
         rv
@@ -91,7 +91,7 @@ impl Layout {
             rv.add_step(delay, move |l| {
                 let win = TestWindow::freeform(delay as usize);
                 l.add_window(win.clone(), Some(ColumnWidth::Proportion(0.3)));
-                l.layout.start_open_animation_for_window(&win);
+                l.layout.start_open_animation_for_window(win.id());
             });
         }
 
@@ -105,7 +105,7 @@ impl Layout {
             rv.add_step(delay, move |l| {
                 let win = TestWindow::freeform(delay as usize);
                 l.add_window(win.clone(), Some(ColumnWidth::Proportion(0.5)));
-                l.layout.start_open_animation_for_window(&win);
+                l.layout.start_open_animation_for_window(win.id());
             });
         }
 
@@ -122,7 +122,7 @@ impl Layout {
             let win = TestWindow::freeform(2);
             let right_of = l.windows[0].clone();
             l.add_window_right_of(&right_of, win.clone(), Some(ColumnWidth::Proportion(0.3)));
-            l.layout.start_open_animation_for_window(&win);
+            l.layout.start_open_animation_for_window(win.id());
         });
 
         rv
@@ -138,7 +138,7 @@ impl Layout {
             let win = TestWindow::freeform(2);
             let right_of = l.windows[0].clone();
             l.add_window_right_of(&right_of, win.clone(), Some(ColumnWidth::Proportion(0.5)));
-            l.layout.start_open_animation_for_window(&win);
+            l.layout.start_open_animation_for_window(win.id());
         });
 
         rv
@@ -147,7 +147,7 @@ impl Layout {
     fn add_window(&mut self, window: TestWindow, width: Option<ColumnWidth>) {
         self.layout.add_window(window.clone(), width, false);
         if window.communicate() {
-            self.layout.update_window(&window);
+            self.layout.update_window(window.id());
         }
         self.windows.push(window);
     }
@@ -159,9 +159,9 @@ impl Layout {
         width: Option<ColumnWidth>,
     ) {
         self.layout
-            .add_window_right_of(right_of, window.clone(), width, false);
+            .add_window_right_of(right_of.id(), window.clone(), width, false);
         if window.communicate() {
-            self.layout.update_window(&window);
+            self.layout.update_window(window.id());
         }
         self.windows.push(window);
     }
@@ -183,7 +183,7 @@ impl TestCase for Layout {
         self.layout.update_output_size(&self.output);
         for win in &self.windows {
             if win.communicate() {
-                self.layout.update_window(win);
+                self.layout.update_window(win.id());
             }
         }
     }
diff --git a/niri-visual-tests/src/test_window.rs b/niri-visual-tests/src/test_window.rs
index 7e50023a..a5b10b20 100644
--- a/niri-visual-tests/src/test_window.rs
+++ b/niri-visual-tests/src/test_window.rs
@@ -12,7 +12,6 @@ use smithay::utils::{Logical, Point, Scale, Size, Transform};
 
 #[derive(Debug)]
 struct TestWindowInner {
-    id: usize,
     size: Size<i32, Logical>,
     requested_size: Option<Size<i32, Logical>>,
     min_size: Size<i32, Logical>,
@@ -24,7 +23,10 @@ struct TestWindowInner {
 }
 
 #[derive(Debug, Clone)]
-pub struct TestWindow(Rc<RefCell<TestWindowInner>>);
+pub struct TestWindow {
+    id: usize,
+    inner: Rc<RefCell<TestWindowInner>>,
+}
 
 impl TestWindow {
     pub fn freeform(id: usize) -> Self {
@@ -33,17 +35,19 @@ impl TestWindow {
         let max_size = Size::from((0, 0));
         let buffer = SolidColorBuffer::new(size, [0.15, 0.64, 0.41, 1.]);
 
-        Self(Rc::new(RefCell::new(TestWindowInner {
+        Self {
             id,
-            size,
-            requested_size: None,
-            min_size,
-            max_size,
-            buffer,
-            pending_fullscreen: false,
-            csd_shadow_width: 0,
-            csd_shadow_buffer: SolidColorBuffer::new((0, 0), [0., 0., 0., 0.3]),
-        })))
+            inner: Rc::new(RefCell::new(TestWindowInner {
+                size,
+                requested_size: None,
+                min_size,
+                max_size,
+                buffer,
+                pending_fullscreen: false,
+                csd_shadow_width: 0,
+                csd_shadow_buffer: SolidColorBuffer::new((0, 0), [0., 0., 0., 0.3]),
+            })),
+        }
     }
 
     pub fn fixed_size(id: usize) -> Self {
@@ -56,24 +60,24 @@ impl TestWindow {
     }
 
     pub fn set_min_size(&self, size: Size<i32, Logical>) {
-        self.0.borrow_mut().min_size = size;
+        self.inner.borrow_mut().min_size = size;
     }
 
     pub fn set_max_size(&self, size: Size<i32, Logical>) {
-        self.0.borrow_mut().max_size = size;
+        self.inner.borrow_mut().max_size = size;
     }
 
     pub fn set_color(&self, color: [f32; 4]) {
-        self.0.borrow_mut().buffer.set_color(color);
+        self.inner.borrow_mut().buffer.set_color(color);
     }
 
     pub fn set_csd_shadow_width(&self, width: i32) {
-        self.0.borrow_mut().csd_shadow_width = width;
+        self.inner.borrow_mut().csd_shadow_width = width;
     }
 
     pub fn communicate(&self) -> bool {
         let mut rv = false;
-        let mut inner = self.0.borrow_mut();
+        let mut inner = self.inner.borrow_mut();
 
         let mut new_size = inner.size;
 
@@ -117,15 +121,15 @@ impl TestWindow {
     }
 }
 
-impl PartialEq for TestWindow {
-    fn eq(&self, other: &Self) -> bool {
-        self.0.borrow().id == other.0.borrow().id
+impl LayoutElement for TestWindow {
+    type Id = usize;
+
+    fn id(&self) -> &Self::Id {
+        &self.id
     }
-}
 
-impl LayoutElement for TestWindow {
     fn size(&self) -> Size<i32, Logical> {
-        self.0.borrow().size
+        self.inner.borrow().size
     }
 
     fn buf_loc(&self) -> Point<i32, Logical> {
@@ -142,7 +146,7 @@ impl LayoutElement for TestWindow {
         location: Point<i32, Logical>,
         scale: Scale<f64>,
     ) -> Vec<LayoutElementRenderElement<R>> {
-        let inner = self.0.borrow();
+        let inner = self.inner.borrow();
 
         vec![
             SolidColorRenderElement::from_buffer(
@@ -166,20 +170,20 @@ impl LayoutElement for TestWindow {
     }
 
     fn request_size(&self, size: Size<i32, Logical>) {
-        self.0.borrow_mut().requested_size = Some(size);
-        self.0.borrow_mut().pending_fullscreen = false;
+        self.inner.borrow_mut().requested_size = Some(size);
+        self.inner.borrow_mut().pending_fullscreen = false;
     }
 
     fn request_fullscreen(&self, _size: Size<i32, Logical>) {
-        self.0.borrow_mut().pending_fullscreen = true;
+        self.inner.borrow_mut().pending_fullscreen = true;
     }
 
     fn min_size(&self) -> Size<i32, Logical> {
-        self.0.borrow().min_size
+        self.inner.borrow().min_size
     }
 
     fn max_size(&self) -> Size<i32, Logical> {
-        self.0.borrow().max_size
+        self.inner.borrow().max_size
     }
 
     fn is_wl_surface(&self, _wl_surface: &WlSurface) -> bool {
@@ -203,6 +207,6 @@ impl LayoutElement for TestWindow {
     }
 
     fn is_pending_fullscreen(&self) -> bool {
-        self.0.borrow().pending_fullscreen
+        self.inner.borrow().pending_fullscreen
     }
 }
diff --git a/src/handlers/mod.rs b/src/handlers/mod.rs
index e5df8b43..4d3a606f 100644
--- a/src/handlers/mod.rs
+++ b/src/handlers/mod.rs
@@ -365,7 +365,7 @@ impl ForeignToplevelHandler for State {
                 if &requested_output != current_output {
                     self.niri
                         .layout
-                        .move_window_to_output(window.clone(), &requested_output);
+                        .move_window_to_output(&window, &requested_output);
                 }
             }
 
diff --git a/src/handlers/xdg_shell.rs b/src/handlers/xdg_shell.rs
index ec40bb7e..66c3dea4 100644
--- a/src/handlers/xdg_shell.rs
+++ b/src/handlers/xdg_shell.rs
@@ -289,7 +289,7 @@ impl XdgShellHandler for State {
                 if &requested_output != current_output {
                     self.niri
                         .layout
-                        .move_window_to_output(window.clone(), &requested_output);
+                        .move_window_to_output(&window, &requested_output);
                 }
             }
 
diff --git a/src/layout/mod.rs b/src/layout/mod.rs
index 55c1e45d..09c8abed 100644
--- a/src/layout/mod.rs
+++ b/src/layout/mod.rs
@@ -69,7 +69,13 @@ niri_render_elements! {
     }
 }
 
-pub trait LayoutElement: PartialEq {
+pub trait LayoutElement {
+    /// Type that can be used as a unique ID of this element.
+    type Id: PartialEq;
+
+    /// Unique ID of this element.
+    fn id(&self) -> &Self::Id;
+
     /// Visual size of the element.
     ///
     /// This is what the user would consider the size, i.e. excluding CSD shadows and whatnot.
@@ -217,6 +223,12 @@ impl Options {
 }
 
 impl LayoutElement for Window {
+    type Id = Self;
+
+    fn id(&self) -> &Self::Id {
+        self
+    }
+
     fn size(&self) -> Size<i32, Logical> {
         self.geometry().size
     }
@@ -599,7 +611,7 @@ impl<W: LayoutElement> Layout<W> {
     /// Returns an output that the window was added to, if there were any outputs.
     pub fn add_window_right_of(
         &mut self,
-        right_of: &W,
+        right_of: &W::Id,
         window: W,
         width: Option<ColumnWidth>,
         is_full_width: bool,
@@ -681,13 +693,15 @@ impl<W: LayoutElement> Layout<W> {
         );
     }
 
-    pub fn remove_window(&mut self, window: &W) {
+    pub fn remove_window(&mut self, window: &W::Id) -> Option<W> {
+        let mut rv = None;
+
         match &mut self.monitor_set {
             MonitorSet::Normal { monitors, .. } => {
                 for mon in monitors {
                     for (idx, ws) in mon.workspaces.iter_mut().enumerate() {
                         if ws.has_window(window) {
-                            ws.remove_window(window);
+                            rv = Some(ws.remove_window(window));
 
                             // Clean up empty workspaces that are not active and not last.
                             if !ws.has_windows()
@@ -710,7 +724,7 @@ impl<W: LayoutElement> Layout<W> {
             MonitorSet::NoOutputs { workspaces, .. } => {
                 for (idx, ws) in workspaces.iter_mut().enumerate() {
                     if ws.has_window(window) {
-                        ws.remove_window(window);
+                        rv = Some(ws.remove_window(window));
 
                         // Clean up empty workspaces.
                         if !ws.has_windows() {
@@ -722,9 +736,11 @@ impl<W: LayoutElement> Layout<W> {
                 }
             }
         }
+
+        rv
     }
 
-    pub fn update_window(&mut self, window: &W) {
+    pub fn update_window(&mut self, window: &W::Id) {
         match &mut self.monitor_set {
             MonitorSet::Normal { monitors, .. } => {
                 for mon in monitors {
@@ -761,7 +777,7 @@ impl<W: LayoutElement> Layout<W> {
         None
     }
 
-    pub fn window_y(&self, window: &W) -> Option<i32> {
+    pub fn window_y(&self, window: &W::Id) -> Option<i32> {
         match &self.monitor_set {
             MonitorSet::Normal { monitors, .. } => {
                 for mon in monitors {
@@ -810,7 +826,7 @@ impl<W: LayoutElement> Layout<W> {
         }
     }
 
-    pub fn activate_window(&mut self, window: &W) {
+    pub fn activate_window(&mut self, window: &W::Id) {
         let MonitorSet::Normal {
             monitors,
             active_monitor_idx,
@@ -1493,7 +1509,7 @@ impl<W: LayoutElement> Layout<W> {
         }
     }
 
-    pub fn move_window_to_output(&mut self, window: W, output: &Output) {
+    pub fn move_window_to_output(&mut self, window: &W::Id, output: &Output) {
         let mut width = None;
         let mut is_full_width = false;
 
@@ -1501,7 +1517,7 @@ impl<W: LayoutElement> Layout<W> {
             for mon in &*monitors {
                 for ws in &mon.workspaces {
                     for col in &ws.columns {
-                        if col.contains(&window) {
+                        if col.contains(window) {
                             width = Some(col.width);
                             is_full_width = col.is_full_width;
                             break;
@@ -1513,7 +1529,7 @@ impl<W: LayoutElement> Layout<W> {
 
         let Some(width) = width else { return };
 
-        self.remove_window(&window);
+        let window = self.remove_window(window).unwrap();
 
         if let MonitorSet::Normal { monitors, .. } = &mut self.monitor_set {
             let new_idx = monitors
@@ -1570,7 +1586,7 @@ impl<W: LayoutElement> Layout<W> {
         *active_monitor_idx = target_idx;
     }
 
-    pub fn set_fullscreen(&mut self, window: &W, is_fullscreen: bool) {
+    pub fn set_fullscreen(&mut self, window: &W::Id, is_fullscreen: bool) {
         match &mut self.monitor_set {
             MonitorSet::Normal { monitors, .. } => {
                 for mon in monitors {
@@ -1593,7 +1609,7 @@ impl<W: LayoutElement> Layout<W> {
         }
     }
 
-    pub fn toggle_fullscreen(&mut self, window: &W) {
+    pub fn toggle_fullscreen(&mut self, window: &W::Id) {
         match &mut self.monitor_set {
             MonitorSet::Normal { monitors, .. } => {
                 for mon in monitors {
@@ -1746,14 +1762,14 @@ impl<W: LayoutElement> Layout<W> {
         monitor.move_workspace_up();
     }
 
-    pub fn start_open_animation_for_window(&mut self, window: &W) {
+    pub fn start_open_animation_for_window(&mut self, window: &W::Id) {
         match &mut self.monitor_set {
             MonitorSet::Normal { monitors, .. } => {
                 for mon in monitors {
                     for ws in &mut mon.workspaces {
                         for col in &mut ws.columns {
                             for tile in &mut col.tiles {
-                                if tile.window() == window {
+                                if tile.window().id() == window {
                                     tile.start_open_animation();
                                     return;
                                 }
@@ -1764,13 +1780,11 @@ impl<W: LayoutElement> Layout<W> {
             }
             MonitorSet::NoOutputs { workspaces, .. } => {
                 for ws in workspaces {
-                    if ws.has_window(window) {
-                        for col in &mut ws.columns {
-                            for tile in &mut col.tiles {
-                                if tile.window() == window {
-                                    tile.start_open_animation();
-                                    return;
-                                }
+                    for col in &mut ws.columns {
+                        for tile in &mut col.tiles {
+                            if tile.window().id() == window {
+                                tile.start_open_animation();
+                                return;
                             }
                         }
                     }
@@ -1889,13 +1903,13 @@ mod tests {
         }
     }
 
-    impl PartialEq for TestWindow {
-        fn eq(&self, other: &Self) -> bool {
-            self.0.id == other.0.id
+    impl LayoutElement for TestWindow {
+        type Id = usize;
+
+        fn id(&self) -> &Self::Id {
+            &self.0.id
         }
-    }
 
-    impl LayoutElement for TestWindow {
         fn size(&self) -> Size<i32, Logical> {
             self.0.bbox.get().size
         }
@@ -2208,24 +2222,14 @@ mod tests {
                         return;
                     }
 
-                    let right_of_win = TestWindow::new(
-                        right_of_id,
-                        Rectangle::default(),
-                        Size::default(),
-                        Size::default(),
-                    );
                     let win = TestWindow::new(id, bbox, min_max_size.0, min_max_size.1);
-                    layout.add_window_right_of(&right_of_win, win, None, false);
+                    layout.add_window_right_of(&right_of_id, win, None, false);
                 }
                 Op::CloseWindow(id) => {
-                    let dummy =
-                        TestWindow::new(id, Rectangle::default(), Size::default(), Size::default());
-                    layout.remove_window(&dummy);
+                    layout.remove_window(&id);
                 }
                 Op::FullscreenWindow(id) => {
-                    let dummy =
-                        TestWindow::new(id, Rectangle::default(), Size::default(), Size::default());
-                    layout.toggle_fullscreen(&dummy);
+                    layout.toggle_fullscreen(&id);
                 }
                 Op::FocusColumnLeft => layout.focus_left(),
                 Op::FocusColumnRight => layout.focus_right(),
@@ -2284,7 +2288,7 @@ mod tests {
                 Op::SetColumnWidth(change) => layout.set_column_width(change),
                 Op::SetWindowHeight(change) => layout.set_window_height(change),
                 Op::Communicate(id) => {
-                    let mut window = None;
+                    let mut update = false;
                     match &mut layout.monitor_set {
                         MonitorSet::Normal { monitors, .. } => {
                             'outer: for mon in monitors {
@@ -2292,7 +2296,7 @@ mod tests {
                                     for win in ws.windows() {
                                         if win.0.id == id {
                                             if win.communicate() {
-                                                window = Some(win.clone());
+                                                update = true;
                                             }
                                             break 'outer;
                                         }
@@ -2305,7 +2309,7 @@ mod tests {
                                 for win in ws.windows() {
                                     if win.0.id == id {
                                         if win.communicate() {
-                                            window = Some(win.clone());
+                                            update = true;
                                         }
                                         break 'outer;
                                     }
@@ -2314,8 +2318,8 @@ mod tests {
                         }
                     }
 
-                    if let Some(win) = window {
-                        layout.update_window(&win);
+                    if update {
+                        layout.update_window(&id);
                     }
                 }
                 Op::MoveWorkspaceToOutput(id) => {
diff --git a/src/layout/monitor.rs b/src/layout/monitor.rs
index dd5709dd..fbea52cd 100644
--- a/src/layout/monitor.rs
+++ b/src/layout/monitor.rs
@@ -159,7 +159,7 @@ impl<W: LayoutElement> Monitor<W> {
 
     pub fn add_window_right_of(
         &mut self,
-        right_of: &W,
+        right_of: &W::Id,
         window: W,
         width: ColumnWidth,
         is_full_width: bool,
diff --git a/src/layout/workspace.rs b/src/layout/workspace.rs
index e2a95ad4..b5d1b220 100644
--- a/src/layout/workspace.rs
+++ b/src/layout/workspace.rs
@@ -577,8 +577,8 @@ impl<W: LayoutElement> Workspace<W> {
         self.windows().next().is_some()
     }
 
-    pub fn has_window(&self, window: &W) -> bool {
-        self.windows().any(|win| win == window)
+    pub fn has_window(&self, window: &W::Id) -> bool {
+        self.windows().any(|win| win.id() == window)
     }
 
     pub fn find_wl_surface(&self, wl_surface: &WlSurface) -> Option<&W> {
@@ -658,7 +658,7 @@ impl<W: LayoutElement> Workspace<W> {
 
     pub fn add_window_right_of(
         &mut self,
-        right_of: &W,
+        right_of: &W::Id,
         window: W,
         width: ColumnWidth,
         is_full_width: bool,
@@ -835,7 +835,7 @@ impl<W: LayoutElement> Workspace<W> {
         column
     }
 
-    pub fn remove_window(&mut self, window: &W) {
+    pub fn remove_window(&mut self, window: &W::Id) -> W {
         let column_idx = self
             .columns
             .iter()
@@ -844,10 +844,10 @@ impl<W: LayoutElement> Workspace<W> {
         let column = &self.columns[column_idx];
 
         let window_idx = column.position(window).unwrap();
-        self.remove_window_by_idx(column_idx, window_idx);
+        self.remove_window_by_idx(column_idx, window_idx)
     }
 
-    pub fn update_window(&mut self, window: &W) {
+    pub fn update_window(&mut self, window: &W::Id) {
         let (idx, column) = self
             .columns
             .iter_mut()
@@ -869,7 +869,7 @@ impl<W: LayoutElement> Workspace<W> {
         }
     }
 
-    pub fn activate_window(&mut self, window: &W) {
+    pub fn activate_window(&mut self, window: &W::Id) {
         let column_idx = self
             .columns
             .iter()
@@ -1233,7 +1233,7 @@ impl<W: LayoutElement> Workspace<W> {
         self.columns[self.active_column_idx].set_window_height(change);
     }
 
-    pub fn set_fullscreen(&mut self, window: &W, is_fullscreen: bool) {
+    pub fn set_fullscreen(&mut self, window: &W::Id, is_fullscreen: bool) {
         let (mut col_idx, tile_idx) = self
             .columns
             .iter()
@@ -1275,7 +1275,7 @@ impl<W: LayoutElement> Workspace<W> {
         col.set_fullscreen(is_fullscreen);
     }
 
-    pub fn toggle_fullscreen(&mut self, window: &W) {
+    pub fn toggle_fullscreen(&mut self, window: &W::Id) {
         let col = self
             .columns
             .iter_mut()
@@ -1654,18 +1654,21 @@ impl<W: LayoutElement> Column<W> {
         self.tiles.iter().any(Tile::are_animations_ongoing)
     }
 
-    pub fn contains(&self, window: &W) -> bool {
-        self.tiles.iter().map(Tile::window).any(|win| win == window)
+    pub fn contains(&self, window: &W::Id) -> bool {
+        self.tiles
+            .iter()
+            .map(Tile::window)
+            .any(|win| win.id() == window)
     }
 
-    pub fn position(&self, window: &W) -> Option<usize> {
+    pub fn position(&self, window: &W::Id) -> Option<usize> {
         self.tiles
             .iter()
             .map(Tile::window)
-            .position(|win| win == window)
+            .position(|win| win.id() == window)
     }
 
-    fn activate_window(&mut self, window: &W) {
+    fn activate_window(&mut self, window: &W::Id) {
         let idx = self.position(window).unwrap();
         self.active_tile_idx = idx;
     }
@@ -1678,11 +1681,11 @@ impl<W: LayoutElement> Column<W> {
         self.update_tile_sizes();
     }
 
-    fn update_window(&mut self, window: &W) {
+    fn update_window(&mut self, window: &W::Id) {
         let tile = self
             .tiles
             .iter_mut()
-            .find(|tile| tile.window() == window)
+            .find(|tile| tile.window().id() == window)
             .unwrap();
         tile.update_window();
     }