| 1 | use x11rb::protocol::xproto::Window as XWindow; |
| 2 | |
| 3 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| 4 | pub enum SplitDirection { |
| 5 | Horizontal, |
| 6 | Vertical, |
| 7 | } |
| 8 | |
| 9 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] |
| 10 | pub enum Direction { |
| 11 | Left, |
| 12 | Right, |
| 13 | Up, |
| 14 | Down, |
| 15 | } |
| 16 | |
| 17 | #[derive(Debug)] |
| 18 | pub enum Node { |
| 19 | Internal { |
| 20 | split: SplitDirection, |
| 21 | ratio: f32, |
| 22 | left: Box<Node>, |
| 23 | right: Box<Node>, |
| 24 | }, |
| 25 | Leaf { |
| 26 | window: Option<XWindow>, |
| 27 | }, |
| 28 | } |
| 29 | |
| 30 | impl Default for Node { |
| 31 | fn default() -> Self { |
| 32 | Self::empty() |
| 33 | } |
| 34 | } |
| 35 | |
| 36 | impl Node { |
| 37 | pub fn empty() -> Self { |
| 38 | Node::Leaf { window: None } |
| 39 | } |
| 40 | |
| 41 | pub fn with_window(window: XWindow) -> Self { |
| 42 | Node::Leaf { |
| 43 | window: Some(window), |
| 44 | } |
| 45 | } |
| 46 | |
| 47 | pub fn is_empty(&self) -> bool { |
| 48 | matches!(self, Node::Leaf { window: None }) |
| 49 | } |
| 50 | |
| 51 | pub fn window_count(&self) -> usize { |
| 52 | match self { |
| 53 | Node::Leaf { window: Some(_) } => 1, |
| 54 | Node::Leaf { window: None } => 0, |
| 55 | Node::Internal { left, right, .. } => left.window_count() + right.window_count(), |
| 56 | } |
| 57 | } |
| 58 | |
| 59 | pub fn windows(&self) -> Vec<XWindow> { |
| 60 | match self { |
| 61 | Node::Leaf { window: Some(w) } => vec![*w], |
| 62 | Node::Leaf { window: None } => vec![], |
| 63 | Node::Internal { left, right, .. } => { |
| 64 | let mut windows = left.windows(); |
| 65 | windows.extend(right.windows()); |
| 66 | windows |
| 67 | } |
| 68 | } |
| 69 | } |
| 70 | |
| 71 | pub fn contains(&self, window: XWindow) -> bool { |
| 72 | match self { |
| 73 | Node::Leaf { window: Some(w) } => *w == window, |
| 74 | Node::Leaf { window: None } => false, |
| 75 | Node::Internal { left, right, .. } => left.contains(window) || right.contains(window), |
| 76 | } |
| 77 | } |
| 78 | |
| 79 | /// Determine optimal split direction based on container dimensions. |
| 80 | /// If wider than tall, split vertically (side-by-side). |
| 81 | /// If taller than wide, split horizontally (stacked). |
| 82 | pub fn smart_split_direction(rect: &Rect) -> SplitDirection { |
| 83 | if rect.width > rect.height { |
| 84 | SplitDirection::Vertical |
| 85 | } else { |
| 86 | SplitDirection::Horizontal |
| 87 | } |
| 88 | } |
| 89 | |
| 90 | /// Insert a window into the tree with smart splitting. |
| 91 | /// If the tree is empty, the window becomes the root. |
| 92 | /// Split direction is determined by container dimensions. |
| 93 | pub fn insert(&mut self, new_window: XWindow, target: Option<XWindow>) { |
| 94 | self.insert_with_rect(new_window, target, Rect::new(0, 0, 1920, 1080)); |
| 95 | } |
| 96 | |
| 97 | /// Insert with known container rect for smart split calculation. |
| 98 | pub fn insert_with_rect(&mut self, new_window: XWindow, target: Option<XWindow>, rect: Rect) { |
| 99 | match self { |
| 100 | Node::Leaf { window: None } => { |
| 101 | // Empty tree - just place the window here |
| 102 | *self = Node::with_window(new_window); |
| 103 | } |
| 104 | Node::Leaf { window: Some(existing) } => { |
| 105 | // Split this leaf using smart split direction |
| 106 | let direction = Self::smart_split_direction(&rect); |
| 107 | let existing_window = *existing; |
| 108 | *self = Node::Internal { |
| 109 | split: direction, |
| 110 | ratio: 0.5, |
| 111 | left: Box::new(Node::with_window(existing_window)), |
| 112 | right: Box::new(Node::with_window(new_window)), |
| 113 | }; |
| 114 | } |
| 115 | Node::Internal { split, ratio, left, right } => { |
| 116 | // Find the target window and insert next to it |
| 117 | let (left_rect, right_rect) = rect.split(*split, *ratio); |
| 118 | |
| 119 | if let Some(target) = target { |
| 120 | if left.contains(target) { |
| 121 | left.insert_with_rect(new_window, Some(target), left_rect); |
| 122 | } else if right.contains(target) { |
| 123 | right.insert_with_rect(new_window, Some(target), right_rect); |
| 124 | } else { |
| 125 | // Target not found, insert in right subtree |
| 126 | right.insert_with_rect(new_window, None, right_rect); |
| 127 | } |
| 128 | } else { |
| 129 | // No target specified, insert in right subtree |
| 130 | right.insert_with_rect(new_window, None, right_rect); |
| 131 | } |
| 132 | } |
| 133 | } |
| 134 | } |
| 135 | |
| 136 | /// Remove a window from the tree. |
| 137 | /// Returns true if the window was found and removed. |
| 138 | pub fn remove(&mut self, window: XWindow) -> bool { |
| 139 | match self { |
| 140 | Node::Leaf { window: Some(w) } if *w == window => { |
| 141 | *self = Node::empty(); |
| 142 | true |
| 143 | } |
| 144 | Node::Leaf { .. } => false, |
| 145 | Node::Internal { left, right, .. } => { |
| 146 | if left.remove(window) { |
| 147 | // Left child removed the window, check if we need to collapse |
| 148 | if left.is_empty() { |
| 149 | // Promote right child |
| 150 | *self = std::mem::take(right.as_mut()); |
| 151 | } |
| 152 | true |
| 153 | } else if right.remove(window) { |
| 154 | // Right child removed the window, check if we need to collapse |
| 155 | if right.is_empty() { |
| 156 | // Promote left child |
| 157 | *self = std::mem::take(left.as_mut()); |
| 158 | } |
| 159 | true |
| 160 | } else { |
| 161 | false |
| 162 | } |
| 163 | } |
| 164 | } |
| 165 | } |
| 166 | |
| 167 | /// Calculate geometries for all windows in the tree. |
| 168 | pub fn calculate_geometries(&self, rect: Rect) -> Vec<(XWindow, Rect)> { |
| 169 | match self { |
| 170 | Node::Leaf { window: Some(w) } => vec![(*w, rect)], |
| 171 | Node::Leaf { window: None } => vec![], |
| 172 | Node::Internal { |
| 173 | split, |
| 174 | ratio, |
| 175 | left, |
| 176 | right, |
| 177 | } => { |
| 178 | let (left_rect, right_rect) = rect.split(*split, *ratio); |
| 179 | let mut geometries = left.calculate_geometries(left_rect); |
| 180 | geometries.extend(right.calculate_geometries(right_rect)); |
| 181 | geometries |
| 182 | } |
| 183 | } |
| 184 | } |
| 185 | |
| 186 | /// Find the first window in the tree (leftmost leaf). |
| 187 | pub fn first_window(&self) -> Option<XWindow> { |
| 188 | match self { |
| 189 | Node::Leaf { window } => *window, |
| 190 | Node::Internal { left, right, .. } => { |
| 191 | left.first_window().or_else(|| right.first_window()) |
| 192 | } |
| 193 | } |
| 194 | } |
| 195 | |
| 196 | /// Equalize all split ratios to 0.5. |
| 197 | pub fn equalize(&mut self) { |
| 198 | if let Node::Internal { |
| 199 | ratio, left, right, .. |
| 200 | } = self |
| 201 | { |
| 202 | *ratio = 0.5; |
| 203 | left.equalize(); |
| 204 | right.equalize(); |
| 205 | } |
| 206 | } |
| 207 | |
| 208 | /// Resize the split affecting a window in the given direction. |
| 209 | /// Returns true if a resize was performed. |
| 210 | pub fn resize(&mut self, window: XWindow, direction: Direction, delta: f32) -> bool { |
| 211 | match self { |
| 212 | Node::Leaf { .. } => false, |
| 213 | Node::Internal { |
| 214 | split, |
| 215 | ratio, |
| 216 | left, |
| 217 | right, |
| 218 | } => { |
| 219 | // Check if this split is in the right orientation for the direction |
| 220 | let dominated = match (split, direction) { |
| 221 | (SplitDirection::Vertical, Direction::Left | Direction::Right) => true, |
| 222 | (SplitDirection::Horizontal, Direction::Up | Direction::Down) => true, |
| 223 | _ => false, |
| 224 | }; |
| 225 | |
| 226 | if dominated { |
| 227 | // Check if the window is on the appropriate side |
| 228 | let in_left = left.contains(window); |
| 229 | let in_right = right.contains(window); |
| 230 | |
| 231 | if in_left || in_right { |
| 232 | // Adjust ratio based on direction |
| 233 | let adjustment = match direction { |
| 234 | Direction::Right | Direction::Down => { |
| 235 | if in_left { delta } else { -delta } |
| 236 | } |
| 237 | Direction::Left | Direction::Up => { |
| 238 | if in_left { -delta } else { delta } |
| 239 | } |
| 240 | }; |
| 241 | |
| 242 | *ratio = (*ratio + adjustment).clamp(0.1, 0.9); |
| 243 | return true; |
| 244 | } |
| 245 | } |
| 246 | |
| 247 | // Recurse into children |
| 248 | if left.contains(window) { |
| 249 | left.resize(window, direction, delta) |
| 250 | } else if right.contains(window) { |
| 251 | right.resize(window, direction, delta) |
| 252 | } else { |
| 253 | false |
| 254 | } |
| 255 | } |
| 256 | } |
| 257 | } |
| 258 | |
| 259 | /// Swap two windows in the tree. |
| 260 | pub fn swap(&mut self, a: XWindow, b: XWindow) -> bool { |
| 261 | // Find and swap the windows |
| 262 | let mut found_a = false; |
| 263 | let mut found_b = false; |
| 264 | |
| 265 | self.swap_impl(a, b, &mut found_a, &mut found_b); |
| 266 | found_a && found_b |
| 267 | } |
| 268 | |
| 269 | fn swap_impl(&mut self, a: XWindow, b: XWindow, found_a: &mut bool, found_b: &mut bool) { |
| 270 | match self { |
| 271 | Node::Leaf { window: Some(w) } => { |
| 272 | if *w == a { |
| 273 | *w = b; |
| 274 | *found_a = true; |
| 275 | } else if *w == b { |
| 276 | *w = a; |
| 277 | *found_b = true; |
| 278 | } |
| 279 | } |
| 280 | Node::Leaf { window: None } => {} |
| 281 | Node::Internal { left, right, .. } => { |
| 282 | left.swap_impl(a, b, found_a, found_b); |
| 283 | right.swap_impl(a, b, found_a, found_b); |
| 284 | } |
| 285 | } |
| 286 | } |
| 287 | |
| 288 | /// Find adjacent window in a direction, given geometries. |
| 289 | /// If `preferred` is Some and is a valid candidate, it will be returned. |
| 290 | /// This enables "window memory" - remembering which window was last focused in a direction. |
| 291 | pub fn find_adjacent( |
| 292 | geometries: &[(XWindow, Rect)], |
| 293 | from: XWindow, |
| 294 | direction: Direction, |
| 295 | preferred: Option<XWindow>, |
| 296 | ) -> Option<XWindow> { |
| 297 | let from_rect = geometries.iter().find(|(w, _)| *w == from)?.1; |
| 298 | |
| 299 | // Find the center point of the source window |
| 300 | let from_cx = from_rect.x as i32 + from_rect.width as i32 / 2; |
| 301 | let from_cy = from_rect.y as i32 + from_rect.height as i32 / 2; |
| 302 | |
| 303 | let candidates: Vec<_> = geometries |
| 304 | .iter() |
| 305 | .filter(|(w, rect)| { |
| 306 | if *w == from { |
| 307 | return false; |
| 308 | } |
| 309 | |
| 310 | let cx = rect.x as i32 + rect.width as i32 / 2; |
| 311 | let cy = rect.y as i32 + rect.height as i32 / 2; |
| 312 | |
| 313 | match direction { |
| 314 | Direction::Left => cx < from_cx, |
| 315 | Direction::Right => cx > from_cx, |
| 316 | Direction::Up => cy < from_cy, |
| 317 | Direction::Down => cy > from_cy, |
| 318 | } |
| 319 | }) |
| 320 | .collect(); |
| 321 | |
| 322 | // If preferred window is a valid candidate, use it (window memory) |
| 323 | if let Some(pref) = preferred { |
| 324 | if candidates.iter().any(|(w, _)| *w == pref) { |
| 325 | return Some(pref); |
| 326 | } |
| 327 | } |
| 328 | |
| 329 | // Find the closest window in the direction |
| 330 | // Prioritize alignment perpendicular to movement, then distance in movement direction |
| 331 | candidates |
| 332 | .into_iter() |
| 333 | .min_by_key(|(_, rect)| { |
| 334 | let cx = rect.x as i32 + rect.width as i32 / 2; |
| 335 | let cy = rect.y as i32 + rect.height as i32 / 2; |
| 336 | let dx = (cx - from_cx).abs(); |
| 337 | let dy = (cy - from_cy).abs(); |
| 338 | |
| 339 | match direction { |
| 340 | // For left/right: prioritize same row (small dy), then closest x |
| 341 | Direction::Left | Direction::Right => dy * 100 + dx, |
| 342 | // For up/down: prioritize same column (small dx), then closest y |
| 343 | Direction::Up | Direction::Down => dx * 100 + dy, |
| 344 | } |
| 345 | }) |
| 346 | .map(|(w, _)| *w) |
| 347 | } |
| 348 | } |
| 349 | |
| 350 | #[derive(Debug, Clone, Copy, Default)] |
| 351 | pub struct Rect { |
| 352 | pub x: i16, |
| 353 | pub y: i16, |
| 354 | pub width: u16, |
| 355 | pub height: u16, |
| 356 | } |
| 357 | |
| 358 | impl Rect { |
| 359 | pub fn new(x: i16, y: i16, width: u16, height: u16) -> Self { |
| 360 | Self { |
| 361 | x, |
| 362 | y, |
| 363 | width, |
| 364 | height, |
| 365 | } |
| 366 | } |
| 367 | |
| 368 | pub fn split(&self, direction: SplitDirection, ratio: f32) -> (Rect, Rect) { |
| 369 | match direction { |
| 370 | SplitDirection::Horizontal => { |
| 371 | let split_y = (self.height as f32 * ratio) as u16; |
| 372 | ( |
| 373 | Rect::new(self.x, self.y, self.width, split_y), |
| 374 | Rect::new( |
| 375 | self.x, |
| 376 | self.y + split_y as i16, |
| 377 | self.width, |
| 378 | self.height.saturating_sub(split_y), |
| 379 | ), |
| 380 | ) |
| 381 | } |
| 382 | SplitDirection::Vertical => { |
| 383 | let split_x = (self.width as f32 * ratio) as u16; |
| 384 | ( |
| 385 | Rect::new(self.x, self.y, split_x, self.height), |
| 386 | Rect::new( |
| 387 | self.x + split_x as i16, |
| 388 | self.y, |
| 389 | self.width.saturating_sub(split_x), |
| 390 | self.height, |
| 391 | ), |
| 392 | ) |
| 393 | } |
| 394 | } |
| 395 | } |
| 396 | } |
| 397 |