atagen

monolith

public
Star0
Watch0
60f00dd4f2bec6086695b8f8830432ee38a0242e
1
ssh://brewery@pls.lobotomise.me:33/atagen/monolith.git
crates/layout/src/tree.rsRust687 lines21.9 KB
1
use monolith_planner::Rect;
2
use slotmap::{SlotMap, new_key_type};
3
use std::collections::HashMap;
4
use std::hash::Hash;
5
6
new_key_type! { pub struct NodeId; }
7
8
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
9
pub enum Axis {
10
    Horizontal,
11
    Vertical,
12
}
13
14
impl Axis {
15
    pub fn other(self) -> Self {
16
        match self {
17
            Self::Horizontal => Self::Vertical,
18
            Self::Vertical => Self::Horizontal,
19
        }
20
    }
21
}
22
23
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
24
pub enum Dir {
25
    Left,
26
    Right,
27
    Up,
28
    Down,
29
}
30
31
impl Dir {
32
    pub fn axis(self) -> Axis {
33
        match self {
34
            Self::Left | Self::Right => Axis::Horizontal,
35
            Self::Up | Self::Down => Axis::Vertical,
36
        }
37
    }
38
39
    pub fn forward(self) -> bool {
40
        matches!(self, Self::Right | Self::Down)
41
    }
42
}
43
44
const MIN_SHARE: f64 = 0.05;
45
46
#[derive(Clone, Debug)]
47
struct Child {
48
    node: NodeId,
49
    weight: f64,
50
}
51
52
#[derive(Clone, Debug)]
53
enum Kind<W> {
54
    Leaf(W),
55
    Container {
56
        axis: Axis,
57
        children: Vec<Child>,
58
        expanded: Option<NodeId>,
59
    },
60
}
61
62
#[derive(Clone, Debug)]
63
struct Node<W> {
64
    parent: Option<NodeId>,
65
    kind: Kind<W>,
66
}
67
68
/// One output's tiling tree. Every tiled window of every tag lives here, the
69
/// caller's visibility filter decides which ones take space.
70
///
71
/// Children run on the axis opposite their parent. Containers hold at least
72
/// two children, hidden ones included, and a container left with one child
73
/// is replaced by it.
74
#[derive(Clone, Debug)]
75
pub struct Tree<W> {
76
    nodes: SlotMap<NodeId, Node<W>>,
77
    root: Option<NodeId>,
78
    leaves: HashMap<W, NodeId>,
79
    stack: Axis,
80
}
81
82
impl<W: Copy + Eq + Hash> Default for Tree<W> {
83
    fn default() -> Self {
84
        Self {
85
            nodes: SlotMap::with_key(),
86
            root: None,
87
            leaves: HashMap::new(),
88
            stack: Axis::Vertical,
89
        }
90
    }
91
}
92
93
impl<W: Copy + Eq + Hash> Tree<W> {
94
    /// The axis tag rows stack along.
95
    pub fn stack_axis(&self) -> Axis {
96
        self.stack
97
    }
98
99
    pub fn contains(&self, w: W) -> bool {
100
        self.leaves.contains_key(&w)
101
    }
102
103
    pub fn is_empty(&self) -> bool {
104
        self.root.is_none()
105
    }
106
107
    pub fn windows(&self) -> impl Iterator<Item = W> + '_ {
108
        self.leaves.keys().copied()
109
    }
110
111
    /// Inserts `w` as the next sibling of `anchor`, or as a new row if
112
    /// `anchor` is not in the tree.
113
    pub fn insert_after(&mut self, w: W, anchor: W) {
114
        let Some(&a) = self.leaves.get(&anchor) else {
115
            return self.insert_row(w);
116
        };
117
        let n = self.new_leaf(w);
118
        match self.position(a) {
119
            Some((p, i)) if Some(p) == self.root && self.axis(p) == Some(self.stack) => {
120
                let weight = self.unlink(p, i).weight;
121
                let c = self.new_container(self.stack.other());
122
                self.attach(c, 0, a, 1.0);
123
                self.attach(c, 1, n, 1.0);
124
                self.attach(p, i, c, weight);
125
            }
126
            Some((p, i)) => {
127
                let weight = self.mean_weight(p);
128
                self.attach(p, i + 1, n, weight);
129
            }
130
            None => {
131
                let c = self.new_container(self.stack.other());
132
                self.attach(c, 0, a, 1.0);
133
                self.attach(c, 1, n, 1.0);
134
                self.root = Some(c);
135
            }
136
        }
137
    }
138
139
    /// Appends `w` as a row of its own along the stacking axis.
140
    pub fn insert_row(&mut self, w: W) {
141
        let n = self.new_leaf(w);
142
        match self.root {
143
            None => self.root = Some(n),
144
            Some(r) if self.axis(r) == Some(self.stack) => {
145
                let weight = self.mean_weight(r);
146
                let end = self.children(r).len();
147
                self.attach(r, end, n, weight);
148
            }
149
            Some(r) => {
150
                let c = self.new_container(self.stack);
151
                self.attach(c, 0, r, 1.0);
152
                self.attach(c, 1, n, 1.0);
153
                self.root = Some(c);
154
            }
155
        }
156
    }
157
158
    pub fn remove(&mut self, w: W) -> bool {
159
        let Some(n) = self.leaves.remove(&w) else {
160
            return false;
161
        };
162
        self.detach(n);
163
        self.nodes.remove(n);
164
        true
165
    }
166
167
    pub fn retain(&mut self, mut keep: impl FnMut(W) -> bool) {
168
        let gone: Vec<W> = self.windows().filter(|&w| !keep(w)).collect();
169
        for w in gone {
170
            self.remove(w);
171
        }
172
    }
173
174
    /// Moves a window one step in `dir`. Leaves swap slots with the next
175
    /// visible sibling, enter a sibling container, or step out beside their
176
    /// container. Hidden windows never move.
177
    pub fn move_window(&mut self, w: W, dir: Dir, vis: impl Fn(W) -> bool) -> bool {
178
        let Some(&n) = self.leaves.get(&w) else {
179
            return false;
180
        };
181
        self.move_node(n, dir, &vis, true)
182
    }
183
184
    /// Moves the container holding `w` as a rigid block. A window sitting
185
    /// directly in the root is its own block.
186
    pub fn move_block(&mut self, w: W, dir: Dir, vis: impl Fn(W) -> bool) -> bool {
187
        let Some(&n) = self.leaves.get(&w) else {
188
            return false;
189
        };
190
        let block = self.nodes[n]
191
            .parent
192
            .filter(|&p| Some(p) != self.root)
193
            .unwrap_or(n);
194
        self.move_node(block, dir, &vis, false)
195
    }
196
197
    /// Transposes the whole tree.
198
    pub fn flip(&mut self) {
199
        self.stack = self.stack.other();
200
        for node in self.nodes.values_mut() {
201
            if let Kind::Container { axis, .. } = &mut node.kind {
202
                *axis = axis.other();
203
            }
204
        }
205
    }
206
207
    /// Toggles expansion for `count` levels of `w`'s ancestry, starting
208
    /// `skip` levels up. Level 0 expands the window inside its container.
209
    pub fn toggle_expand(&mut self, w: W, skip: usize, count: usize) -> bool {
210
        let Some(&leaf) = self.leaves.get(&w) else {
211
            return false;
212
        };
213
        let mut path = Vec::new();
214
        let mut cur = leaf;
215
        while let Some(p) = self.nodes[cur].parent {
216
            path.push((p, cur));
217
            cur = p;
218
        }
219
        let path: Vec<_> = path.into_iter().skip(skip).take(count).collect();
220
        if path.is_empty() {
221
            return false;
222
        }
223
        let all = path.iter().all(|&(p, c)| self.expanded(p) == Some(c));
224
        for (p, c) in path {
225
            if let Kind::Container { expanded, .. } = &mut self.nodes[p].kind {
226
                *expanded = (!all).then_some(c);
227
            }
228
        }
229
        true
230
    }
231
232
    /// Grows `w` by `delta`, a fraction of the visible length of the nearest
233
    /// container along `axis` where it has a visible sibling. Hidden siblings
234
    /// keep their weights, and with them their share.
235
    pub fn resize(&mut self, w: W, axis: Axis, delta: f64, vis: impl Fn(W) -> bool) -> bool {
236
        let Some(&leaf) = self.leaves.get(&w) else {
237
            return false;
238
        };
239
        let mut cur = leaf;
240
        while let Some((p, i)) = self.position(cur) {
241
            if self.axis(p) == Some(axis) {
242
                let shown: Vec<usize> = (0..self.children(p).len())
243
                    .filter(|&j| self.visible(self.children(p)[j].node, &vis))
244
                    .collect();
245
                if shown.len() >= 2 {
246
                    let children = self.children_mut(p);
247
                    let sum: f64 = shown.iter().map(|&j| children[j].weight).sum();
248
                    let old = children[i].weight;
249
                    let max = 1.0 - MIN_SHARE * (shown.len() - 1) as f64;
250
                    let new = (old / sum + delta).clamp(MIN_SHARE, max) * sum;
251
                    let scale = (sum - new) / (sum - old);
252
                    for j in shown {
253
                        children[j].weight = if j == i {
254
                            new
255
                        } else {
256
                            children[j].weight * scale
257
                        };
258
                    }
259
                    return true;
260
                }
261
            }
262
            cur = p;
263
        }
264
        false
265
    }
266
267
    /// Lays the tree out in `area`. Hidden windows come back as zero-size
268
    /// rects on the boundary where they would reappear.
269
    pub fn arrange(&self, area: Rect, gap: f64, vis: impl Fn(W) -> bool) -> Vec<(W, Rect)> {
270
        let mut out = Vec::with_capacity(self.leaves.len());
271
        if let Some(r) = self.root {
272
            self.place(r, area, gap, &vis, &mut out);
273
        }
274
        out
275
    }
276
277
    fn place(
278
        &self,
279
        n: NodeId,
280
        r: Rect,
281
        gap: f64,
282
        vis: &impl Fn(W) -> bool,
283
        out: &mut Vec<(W, Rect)>,
284
    ) {
285
        let (axis, children, expanded) = match &self.nodes[n].kind {
286
            Kind::Leaf(w) => return out.push((*w, r)),
287
            Kind::Container {
288
                axis,
289
                children,
290
                expanded,
291
            } => (*axis, children, *expanded),
292
        };
293
        let shown: Vec<bool> = children.iter().map(|c| self.visible(c.node, vis)).collect();
294
        let (start, len) = span(r, axis);
295
        if let Some(e) = expanded.and_then(|e| children.iter().position(|c| c.node == e))
296
            && shown[e]
297
        {
298
            for (i, c) in children.iter().enumerate() {
299
                let sub = match i.cmp(&e) {
300
                    std::cmp::Ordering::Equal => r,
301
                    std::cmp::Ordering::Less => with_span(r, axis, start, 0.0),
302
                    std::cmp::Ordering::Greater => with_span(r, axis, start + len, 0.0),
303
                };
304
                self.place(c.node, sub, gap, vis, out);
305
            }
306
            return;
307
        }
308
        let count = shown.iter().filter(|s| **s).count();
309
        let total: f64 = children
310
            .iter()
311
            .zip(&shown)
312
            .filter(|(_, s)| **s)
313
            .map(|(c, _)| c.weight)
314
            .sum();
315
        let gaps = count.saturating_sub(1) as f64;
316
        let gap = if gaps > 0.0 { gap.min(len / gaps) } else { gap };
317
        let avail = len - gap * gaps;
318
        let last = shown.iter().rposition(|s| *s);
319
        let mut pos = start;
320
        let mut seen = false;
321
        for (i, c) in children.iter().enumerate() {
322
            if shown[i] {
323
                if seen {
324
                    pos += gap;
325
                }
326
                let l = avail * c.weight / total;
327
                self.place(c.node, with_span(r, axis, pos, l), gap, vis, out);
328
                pos += l;
329
                seen = true;
330
            } else {
331
                let between = seen && last.is_some_and(|l| l > i);
332
                let at = if between { pos + gap / 2.0 } else { pos };
333
                self.place(c.node, with_span(r, axis, at, 0.0), gap, vis, out);
334
            }
335
        }
336
    }
337
338
    fn move_node(&mut self, n: NodeId, dir: Dir, vis: &impl Fn(W) -> bool, enter: bool) -> bool {
339
        let mut cur = n;
340
        while let Some((p, i)) = self.position(cur) {
341
            if self.axis(p) == Some(dir.axis()) {
342
                let hidden_lane = self
343
                    .adjacent(p, i, dir)
344
                    .filter(|&j| enter && !self.visible(self.children(p)[j].node, vis));
345
                if cur != n {
346
                    if let Some(j) = hidden_lane {
347
                        self.join_lane(n, p, j);
348
                    } else {
349
                        let at = if dir.forward() { i + 1 } else { i };
350
                        self.relocate(n, p, at);
351
                    }
352
                    return true;
353
                }
354
                if let Some(j) = self.next_shown(p, i, dir, vis) {
355
                    let sib = self.children(p)[j].node;
356
                    if enter && self.axis(sib).is_some() {
357
                        let end = self.children(sib).len();
358
                        self.relocate(n, sib, end);
359
                    } else {
360
                        let children = self.children_mut(p);
361
                        children[i].node = sib;
362
                        children[j].node = n;
363
                    }
364
                    return true;
365
                }
366
                if let Some(j) = hidden_lane {
367
                    self.join_lane(n, p, j);
368
                    return true;
369
                }
370
            }
371
            cur = p;
372
        }
373
        if cur == n || self.axis(cur) == Some(dir.axis()) {
374
            return false;
375
        }
376
        let (old, i) = self.position(n).expect("n is below the root");
377
        self.unlink(old, i);
378
        self.normalize(old);
379
        let r = self.root.expect("the rest of the tree survives");
380
        if self.axis(r) == Some(dir.axis()) {
381
            let weight = self.mean_weight(r);
382
            let at = if dir.forward() {
383
                self.children(r).len()
384
            } else {
385
                0
386
            };
387
            self.attach(r, at, n, weight);
388
        } else {
389
            let c = self.new_container(dir.axis());
390
            self.attach(c, 0, r, 1.0);
391
            let at = if dir.forward() {
392
                self.children(c).len()
393
            } else {
394
                0
395
            };
396
            self.attach(c, at, n, 1.0);
397
            self.root = Some(c);
398
        }
399
        true
400
    }
401
402
    fn next_shown(&self, p: NodeId, i: usize, dir: Dir, vis: &impl Fn(W) -> bool) -> Option<usize> {
403
        let children = self.children(p);
404
        if dir.forward() {
405
            (i + 1..children.len()).find(|&j| self.visible(children[j].node, vis))
406
        } else {
407
            (0..i).rev().find(|&j| self.visible(children[j].node, vis))
408
        }
409
    }
410
411
    fn adjacent(&self, p: NodeId, i: usize, dir: Dir) -> Option<usize> {
412
        let j = if dir.forward() {
413
            i + 1
414
        } else {
415
            i.checked_sub(1)?
416
        };
417
        (j < self.children(p).len()).then_some(j)
418
    }
419
420
    /// Moves `n` to the end of the lane at `p`'s child `j`, which holds only
421
    /// hidden windows, so it shares that lane when they come back instead of
422
    /// opening a new one beside it. A lone window becomes a lane of its own.
423
    fn join_lane(&mut self, n: NodeId, p: NodeId, j: usize) {
424
        let lane = self.children(p)[j].node;
425
        let lane = if self.axis(lane).is_some() {
426
            lane
427
        } else {
428
            let c = self.new_container(self.axis(p).expect("p is a container").other());
429
            self.children_mut(p)[j].node = c;
430
            self.nodes[c].parent = Some(p);
431
            if let Kind::Container { expanded, .. } = &mut self.nodes[p].kind
432
                && *expanded == Some(lane)
433
            {
434
                *expanded = Some(c);
435
            }
436
            self.attach(c, 0, lane, 1.0);
437
            c
438
        };
439
        let end = self.children(lane).len();
440
        self.relocate(n, lane, end);
441
    }
442
443
    fn relocate(&mut self, n: NodeId, p: NodeId, at: usize) {
444
        let (old, i) = self.position(n).expect("relocated nodes have a parent");
445
        let weight = self.mean_weight(p);
446
        let at = if old == p && i < at { at - 1 } else { at };
447
        self.unlink(old, i);
448
        self.attach(p, at, n, weight);
449
        self.normalize(old);
450
    }
451
452
    /// Places `node` among `p`'s children. A container running along `p`'s
453
    /// axis is dissolved into `p` so axes keep alternating.
454
    fn attach(&mut self, p: NodeId, at: usize, node: NodeId, weight: f64) {
455
        if self.axis(node) == self.axis(p) {
456
            let Some(Node {
457
                kind: Kind::Container { children, .. },
458
                ..
459
            }) = self.nodes.remove(node)
460
            else {
461
                unreachable!("only containers have an axis");
462
            };
463
            let total: f64 = children.iter().map(|c| c.weight).sum();
464
            for (k, c) in children.into_iter().enumerate() {
465
                self.nodes[c.node].parent = Some(p);
466
                let weight = weight * c.weight / total;
467
                self.children_mut(p).insert(
468
                    at + k,
469
                    Child {
470
                        node: c.node,
471
                        weight,
472
                    },
473
                );
474
            }
475
        } else {
476
            self.nodes[node].parent = Some(p);
477
            self.children_mut(p).insert(at, Child { node, weight });
478
        }
479
    }
480
481
    fn unlink(&mut self, p: NodeId, i: usize) -> Child {
482
        let Kind::Container {
483
            children, expanded, ..
484
        } = &mut self.nodes[p].kind
485
        else {
486
            unreachable!("parents are containers");
487
        };
488
        let child = children.remove(i);
489
        if *expanded == Some(child.node) {
490
            *expanded = None;
491
        }
492
        self.nodes[child.node].parent = None;
493
        child
494
    }
495
496
    fn detach(&mut self, n: NodeId) {
497
        match self.position(n) {
498
            Some((p, i)) => {
499
                self.unlink(p, i);
500
                self.normalize(p);
501
            }
502
            None => self.root = None,
503
        }
504
    }
505
506
    /// Repairs a container that lost a child.
507
    fn normalize(&mut self, c: NodeId) {
508
        if self.children(c).len() >= 2 {
509
            return;
510
        }
511
        let pos = self.position(c);
512
        let was_expanded = pos.is_some_and(|(p, _)| self.expanded(p) == Some(c));
513
        let Some(Node {
514
            kind: Kind::Container { children, .. },
515
            ..
516
        }) = self.nodes.remove(c)
517
        else {
518
            unreachable!("only containers are normalized");
519
        };
520
        let only = children.into_iter().next();
521
        match (only, pos) {
522
            (None, Some((p, i))) => {
523
                self.unlink_dead(p, i);
524
                self.normalize(p);
525
            }
526
            (None, None) => self.root = None,
527
            (Some(only), Some((p, i))) => {
528
                let weight = self.unlink_dead(p, i).weight;
529
                let leaf_like = self.axis(only.node) != self.axis(p);
530
                self.attach(p, i, only.node, weight);
531
                if was_expanded
532
                    && leaf_like
533
                    && let Kind::Container { expanded, .. } = &mut self.nodes[p].kind
534
                {
535
                    *expanded = Some(only.node);
536
                }
537
            }
538
            (Some(only), None) => {
539
                self.nodes[only.node].parent = None;
540
                self.root = Some(only.node);
541
            }
542
        }
543
    }
544
545
    /// Like `unlink`, for a child whose node is already gone.
546
    fn unlink_dead(&mut self, p: NodeId, i: usize) -> Child {
547
        let Kind::Container {
548
            children, expanded, ..
549
        } = &mut self.nodes[p].kind
550
        else {
551
            unreachable!("parents are containers");
552
        };
553
        let child = children.remove(i);
554
        if *expanded == Some(child.node) {
555
            *expanded = None;
556
        }
557
        child
558
    }
559
560
    fn new_leaf(&mut self, w: W) -> NodeId {
561
        let n = self.nodes.insert(Node {
562
            parent: None,
563
            kind: Kind::Leaf(w),
564
        });
565
        self.leaves.insert(w, n);
566
        n
567
    }
568
569
    fn new_container(&mut self, axis: Axis) -> NodeId {
570
        self.nodes.insert(Node {
571
            parent: None,
572
            kind: Kind::Container {
573
                axis,
574
                children: Vec::new(),
575
                expanded: None,
576
            },
577
        })
578
    }
579
580
    fn position(&self, n: NodeId) -> Option<(NodeId, usize)> {
581
        let p = self.nodes[n].parent?;
582
        let i = self.children(p).iter().position(|c| c.node == n)?;
583
        Some((p, i))
584
    }
585
586
    fn axis(&self, n: NodeId) -> Option<Axis> {
587
        match &self.nodes[n].kind {
588
            Kind::Container { axis, .. } => Some(*axis),
589
            Kind::Leaf(_) => None,
590
        }
591
    }
592
593
    fn expanded(&self, n: NodeId) -> Option<NodeId> {
594
        match &self.nodes[n].kind {
595
            Kind::Container { expanded, .. } => *expanded,
596
            Kind::Leaf(_) => None,
597
        }
598
    }
599
600
    fn children(&self, n: NodeId) -> &[Child] {
601
        match &self.nodes[n].kind {
602
            Kind::Container { children, .. } => children,
603
            Kind::Leaf(_) => &[],
604
        }
605
    }
606
607
    fn children_mut(&mut self, n: NodeId) -> &mut Vec<Child> {
608
        match &mut self.nodes[n].kind {
609
            Kind::Container { children, .. } => children,
610
            Kind::Leaf(_) => unreachable!("leaves have no children"),
611
        }
612
    }
613
614
    fn mean_weight(&self, p: NodeId) -> f64 {
615
        let children = self.children(p);
616
        if children.is_empty() {
617
            1.0
618
        } else {
619
            children.iter().map(|c| c.weight).sum::<f64>() / children.len() as f64
620
        }
621
    }
622
623
    fn visible(&self, n: NodeId, vis: &impl Fn(W) -> bool) -> bool {
624
        match &self.nodes[n].kind {
625
            Kind::Leaf(w) => vis(*w),
626
            Kind::Container { children, .. } => children.iter().any(|c| self.visible(c.node, vis)),
627
        }
628
    }
629
630
    #[cfg(test)]
631
    pub(crate) fn check(&self) {
632
        let reachable = match self.root {
633
            Some(r) => {
634
                assert!(self.nodes[r].parent.is_none());
635
                self.check_node(r)
636
            }
637
            None => 0,
638
        };
639
        assert_eq!(reachable, self.nodes.len(), "orphaned nodes");
640
        for (&w, &n) in &self.leaves {
641
            assert!(matches!(self.nodes[n].kind, Kind::Leaf(x) if x == w));
642
        }
643
    }
644
645
    #[cfg(test)]
646
    fn check_node(&self, n: NodeId) -> usize {
647
        match &self.nodes[n].kind {
648
            Kind::Leaf(_) => 1,
649
            Kind::Container {
650
                axis,
651
                children,
652
                expanded,
653
            } => {
654
                assert!(
655
                    children.len() >= 2,
656
                    "container with {} children",
657
                    children.len()
658
                );
659
                if let Some(e) = expanded {
660
                    assert!(children.iter().any(|c| c.node == *e));
661
                }
662
                let mut count = 1;
663
                for c in children {
664
                    assert!(c.weight > 0.0 && c.weight.is_finite());
665
                    assert_eq!(self.nodes[c.node].parent, Some(n));
666
                    assert_ne!(self.axis(c.node), Some(*axis), "axes must alternate");
667
                    count += self.check_node(c.node);
668
                }
669
                count
670
            }
671
        }
672
    }
673
}
674
675
fn span(r: Rect, axis: Axis) -> (f64, f64) {
676
    match axis {
677
        Axis::Horizontal => (r.x, r.w),
678
        Axis::Vertical => (r.y, r.h),
679
    }
680
}
681
682
fn with_span(r: Rect, axis: Axis, start: f64, len: f64) -> Rect {
683
    match axis {
684
        Axis::Horizontal => Rect::new(start, r.y, len, r.h),
685
        Axis::Vertical => Rect::new(r.x, start, r.w, len),
686
    }
687
}