Rust · 13558 bytes Raw Blame History
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