monolith

public

Star0

Watch0

109 commits 1 branch 0 tags 0 releases1.5 MiB

use monolith_planner::Rect;
use slotmap::{SlotMap, new_key_type};
use std::collections::HashMap;
use std::hash::Hash;

new_key_type! { pub struct NodeId; }

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[cfg_attr(feature = "serde", serde(rename_all = "kebab-case"))]
pub enum Axis {
    Horizontal,
    Vertical,
}

impl Axis {
    pub fn other(self) -> Self {
        match self {
            Self::Horizontal => Self::Vertical,
            Self::Vertical => Self::Horizontal,
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[cfg_attr(feature = "serde", serde(rename_all = "kebab-case"))]
pub enum Dir {
    Left,
    Right,
    Up,
    Down,
}

impl Dir {
    pub fn axis(self) -> Axis {
        match self {
            Self::Left | Self::Right => Axis::Horizontal,
            Self::Up | Self::Down => Axis::Vertical,
        }
    }

    pub fn forward(self) -> bool {
        matches!(self, Self::Right | Self::Down)
    }
}

const MIN_SHARE: f64 = 0.05;

#[derive(Clone, Debug)]
struct Child {
    node: NodeId,
    weight: f64,
}

#[derive(Clone, Debug)]
enum Kind<W> {
    Leaf(W),
    Container {
        axis: Axis,
        children: Vec<Child>,
        expanded: Option<NodeId>,
    },
}

#[derive(Clone, Debug)]
struct Node<W> {
    parent: Option<NodeId>,
    kind: Kind<W>,
}

/// One output's tiling tree. Every tiled window of every tag lives here, the
/// caller's visibility filter decides which ones take space.
///
/// Children run on the axis opposite their parent. Containers hold at least
/// two children, hidden ones included, and a container left with one child
/// is replaced by it.
#[derive(Clone, Debug)]
pub struct Tree<W> {
    nodes: SlotMap<NodeId, Node<W>>,
    root: Option<NodeId>,
    leaves: HashMap<W, NodeId>,
    stack: Axis,
}

impl<W: Copy + Eq + Hash> Default for Tree<W> {
    fn default() -> Self {
        Self {
            nodes: SlotMap::with_key(),
            root: None,
            leaves: HashMap::new(),
            stack: Axis::Vertical,
        }
    }
}

impl<W: Copy + Eq + Hash> Tree<W> {
    /// The axis tag rows stack along.
    pub fn stack_axis(&self) -> Axis {
        self.stack
    }

    pub fn contains(&self, w: W) -> bool {
        self.leaves.contains_key(&w)
    }

    pub fn is_empty(&self) -> bool {
        self.root.is_none()
    }

    pub fn windows(&self) -> impl Iterator<Item = W> + '_ {
        self.leaves.keys().copied()
    }

    /// Inserts `w` as the next sibling of `anchor`, or as a new row if
    /// `anchor` is not in the tree.
    pub fn insert_after(&mut self, w: W, anchor: W) {
        let Some(&a) = self.leaves.get(&anchor) else {
            return self.insert_row(w);
        };
        let n = self.new_leaf(w);
        match self.position(a) {
            Some((p, i)) if Some(p) == self.root && self.axis(p) == Some(self.stack) => {
                let weight = self.unlink(p, i).weight;
                let c = self.new_container(self.stack.other());
                self.attach(c, 0, a, 1.0);
                self.attach(c, 1, n, 1.0);
                self.attach(p, i, c, weight);
            }
            Some((p, i)) => {
                let weight = self.mean_weight(p);
                self.attach(p, i + 1, n, weight);
            }
            None => {
                let c = self.new_container(self.stack.other());
                self.attach(c, 0, a, 1.0);
                self.attach(c, 1, n, 1.0);
                self.root = Some(c);
            }
        }
    }

    /// Appends `w` as a row of its own along the stacking axis.
    pub fn insert_row(&mut self, w: W) {
        let n = self.new_leaf(w);
        match self.root {
            None => self.root = Some(n),
            Some(r) if self.axis(r) == Some(self.stack) => {
                let weight = self.mean_weight(r);
                let end = self.children(r).len();
                self.attach(r, end, n, weight);
            }
            Some(r) => {
                let c = self.new_container(self.stack);
                self.attach(c, 0, r, 1.0);
                self.attach(c, 1, n, 1.0);
                self.root = Some(c);
            }
        }
    }

    pub fn remove(&mut self, w: W) -> bool {
        let Some(n) = self.leaves.remove(&w) else {
            return false;
        };
        self.detach(n);
        self.nodes.remove(n);
        true
    }

    pub fn retain(&mut self, mut keep: impl FnMut(W) -> bool) {
        let gone: Vec<W> = self.windows().filter(|&w| !keep(w)).collect();
        for w in gone {
            self.remove(w);
        }
    }

    /// Moves a window one step in `dir`. Leaves swap slots with the next
    /// visible sibling, enter a sibling container, or step out beside their
    /// container. Hidden windows never move.
    pub fn move_window(&mut self, w: W, dir: Dir, vis: impl Fn(W) -> bool) -> bool {
        let Some(&n) = self.leaves.get(&w) else {
            return false;
        };
        self.move_node(n, dir, &vis, true)
    }

    /// Moves the container holding `w` as a rigid block. A window sitting
    /// directly in the root is its own block.
    pub fn move_container(&mut self, w: W, dir: Dir, vis: impl Fn(W) -> bool) -> bool {
        let Some(&n) = self.leaves.get(&w) else {
            return false;
        };
        let block = self.nodes[n]
            .parent
            .filter(|&p| Some(p) != self.root)
            .unwrap_or(n);
        self.move_node(block, dir, &vis, false)
    }

    /// Moves a window to the end of the next visible lane in `dir`, stepping
    /// up a level when its container has none that way. A lone window there
    /// becomes a stack with it. With no lane anywhere that way, it leaves its
    /// stack as `move_window` would.
    pub fn merge_window(&mut self, w: W, dir: Dir, vis: impl Fn(W) -> bool) -> bool {
        let Some(&n) = self.leaves.get(&w) else {
            return false;
        };
        let mut cur = n;
        while let Some((p, i)) = self.position(cur) {
            if self.axis(p) == Some(dir.axis())
                && let Some(j) = self.next_shown(p, i, dir, &vis)
            {
                self.join_lane(n, p, j);
                return true;
            }
            cur = p;
        }
        self.move_node(n, dir, &vis, true)
    }

    /// Whether every container above `w` is expanded towards it.
    pub fn maximized(&self, w: W) -> bool {
        let Some(&leaf) = self.leaves.get(&w) else {
            return false;
        };
        let mut cur = leaf;
        let mut any = false;
        while let Some(p) = self.nodes[cur].parent {
            if self.expanded(p) != Some(cur) {
                return false;
            }
            any = true;
            cur = p;
        }
        any
    }

    /// Transposes the whole tree.
    pub fn flip(&mut self) {
        self.stack = self.stack.other();
        for node in self.nodes.values_mut() {
            if let Kind::Container { axis, .. } = &mut node.kind {
                *axis = axis.other();
            }
        }
    }

    /// Toggles expansion for `count` levels of `w`'s ancestry, starting
    /// `skip` levels up. Level 0 expands the window inside its container.
    pub fn toggle_expand(&mut self, w: W, skip: usize, count: usize) -> bool {
        let Some(&leaf) = self.leaves.get(&w) else {
            return false;
        };
        let mut path = Vec::new();
        let mut cur = leaf;
        while let Some(p) = self.nodes[cur].parent {
            path.push((p, cur));
            cur = p;
        }
        let path: Vec<_> = path.into_iter().skip(skip).take(count).collect();
        if path.is_empty() {
            return false;
        }
        let all = path.iter().all(|&(p, c)| self.expanded(p) == Some(c));
        for (p, c) in path {
            if let Kind::Container { expanded, .. } = &mut self.nodes[p].kind {
                *expanded = (!all).then_some(c);
            }
        }
        true
    }

    /// Grows `w` by `delta`, a fraction of the visible length of the nearest
    /// container along `axis` where it has a visible sibling. Hidden siblings
    /// keep their weights, and with them their share.
    pub fn resize(&mut self, w: W, axis: Axis, delta: f64, vis: impl Fn(W) -> bool) -> bool {
        let Some(&leaf) = self.leaves.get(&w) else {
            return false;
        };
        let mut cur = leaf;
        while let Some((p, i)) = self.position(cur) {
            if self.axis(p) == Some(axis) {
                let shown: Vec<usize> = (0..self.children(p).len())
                    .filter(|&j| self.visible(self.children(p)[j].node, &vis))
                    .collect();
                if shown.len() >= 2 {
                    let children = self.children_mut(p);
                    let sum: f64 = shown.iter().map(|&j| children[j].weight).sum();
                    let old = children[i].weight;
                    let max = 1.0 - MIN_SHARE * (shown.len() - 1) as f64;
                    let new = (old / sum + delta).clamp(MIN_SHARE, max) * sum;
                    let scale = (sum - new) / (sum - old);
                    for j in shown {
                        children[j].weight = if j == i {
                            new
                        } else {
                            children[j].weight * scale
                        };
                    }
                    return true;
                }
            }
            cur = p;
        }
        false
    }

    /// Lays the tree out in `area`. Hidden windows come back as zero-size
    /// rects on the boundary where they would reappear.
    pub fn arrange(&self, area: Rect, gap: f64, vis: impl Fn(W) -> bool) -> Vec<(W, Rect)> {
        let mut out = Vec::with_capacity(self.leaves.len());
        if let Some(r) = self.root {
            let area = if self.visible(r, &vis) {
                area
            } else {
                Rect::new(area.x + area.w / 2.0, area.y, 0.0, area.h)
            };
            self.place(r, area, gap, &vis, &mut out);
        }
        out
    }

    fn place(
        &self,
        n: NodeId,
        r: Rect,
        gap: f64,
        vis: &impl Fn(W) -> bool,
        out: &mut Vec<(W, Rect)>,
    ) {
        let (axis, children, expanded) = match &self.nodes[n].kind {
            Kind::Leaf(w) => return out.push((*w, r)),
            Kind::Container {
                axis,
                children,
                expanded,
            } => (*axis, children, *expanded),
        };
        let shown: Vec<bool> = children.iter().map(|c| self.visible(c.node, vis)).collect();
        let (start, len) = span(r, axis);
        if let Some(e) = expanded.and_then(|e| children.iter().position(|c| c.node == e))
            && shown[e]
        {
            for (i, c) in children.iter().enumerate() {
                let sub = match i.cmp(&e) {
                    std::cmp::Ordering::Equal => r,
                    std::cmp::Ordering::Less => with_span(r, axis, start, 0.0),
                    std::cmp::Ordering::Greater => with_span(r, axis, start + len, 0.0),
                };
                self.place(c.node, sub, gap, vis, out);
            }
            return;
        }
        let count = shown.iter().filter(|s| **s).count();
        let total: f64 = children
            .iter()
            .zip(&shown)
            .filter(|(_, s)| **s)
            .map(|(c, _)| c.weight)
            .sum();
        let gaps = count.saturating_sub(1) as f64;
        let gap = if gaps > 0.0 { gap.min(len / gaps) } else { gap };
        let avail = len - gap * gaps;
        let last = shown.iter().rposition(|s| *s);
        let mut pos = start;
        let mut seen = false;
        for (i, c) in children.iter().enumerate() {
            if shown[i] {
                if seen {
                    pos += gap;
                }
                let l = avail * c.weight / total;
                self.place(c.node, with_span(r, axis, pos, l), gap, vis, out);
                pos += l;
                seen = true;
            } else {
                let between = seen && last.is_some_and(|l| l > i);
                let at = if between { pos + gap / 2.0 } else { pos };
                self.place(c.node, with_span(r, axis, at, 0.0), gap, vis, out);
            }
        }
    }

    fn move_node(&mut self, n: NodeId, dir: Dir, vis: &impl Fn(W) -> bool, enter: bool) -> bool {
        let mut cur = n;
        while let Some((p, i)) = self.position(cur) {
            if self.axis(p) == Some(dir.axis()) {
                let hidden_lane = self
                    .adjacent(p, i, dir)
                    .filter(|&j| enter && !self.visible(self.children(p)[j].node, vis));
                if cur != n {
                    if let Some(j) = hidden_lane {
                        self.join_lane(n, p, j);
                    } else {
                        let at = if dir.forward() { i + 1 } else { i };
                        self.relocate(n, p, at);
                    }
                    return true;
                }
                if let Some(j) = self.next_shown(p, i, dir, vis) {
                    let sib = self.children(p)[j].node;
                    if enter && self.axis(sib).is_some() {
                        let end = self.children(sib).len();
                        self.relocate(n, sib, end);
                    } else {
                        let children = self.children_mut(p);
                        children[i].node = sib;
                        children[j].node = n;
                    }
                    return true;
                }
                if let Some(j) = hidden_lane {
                    self.join_lane(n, p, j);
                    return true;
                }
            }
            cur = p;
        }
        if cur == n || self.axis(cur) == Some(dir.axis()) {
            return false;
        }
        let (old, i) = self.position(n).expect("n is below the root");
        self.unlink(old, i);
        self.normalize(old);
        let r = self.root.expect("the rest of the tree survives");
        if self.axis(r) == Some(dir.axis()) {
            let weight = self.mean_weight(r);
            let at = if dir.forward() {
                self.children(r).len()
            } else {
                0
            };
            self.attach(r, at, n, weight);
        } else {
            let c = self.new_container(dir.axis());
            self.attach(c, 0, r, 1.0);
            let at = if dir.forward() {
                self.children(c).len()
            } else {
                0
            };
            self.attach(c, at, n, 1.0);
            self.root = Some(c);
        }
        true
    }

    fn next_shown(&self, p: NodeId, i: usize, dir: Dir, vis: &impl Fn(W) -> bool) -> Option<usize> {
        let children = self.children(p);
        if dir.forward() {
            (i + 1..children.len()).find(|&j| self.visible(children[j].node, vis))
        } else {
            (0..i).rev().find(|&j| self.visible(children[j].node, vis))
        }
    }

    fn adjacent(&self, p: NodeId, i: usize, dir: Dir) -> Option<usize> {
        let j = if dir.forward() {
            i + 1
        } else {
            i.checked_sub(1)?
        };
        (j < self.children(p).len()).then_some(j)
    }

    /// Moves `n` to the end of the lane at `p`'s child `j`. A lone window
    /// there becomes a lane of its own, perpendicular to `p`.
    fn join_lane(&mut self, n: NodeId, p: NodeId, j: usize) {
        let lane = self.children(p)[j].node;
        let lane = if self.axis(lane).is_some() {
            lane
        } else {
            let c = self.new_container(self.axis(p).expect("p is a container").other());
            self.children_mut(p)[j].node = c;
            self.nodes[c].parent = Some(p);
            if let Kind::Container { expanded, .. } = &mut self.nodes[p].kind
                && *expanded == Some(lane)
            {
                *expanded = Some(c);
            }
            self.attach(c, 0, lane, 1.0);
            c
        };
        let end = self.children(lane).len();
        self.relocate(n, lane, end);
    }

    fn relocate(&mut self, n: NodeId, p: NodeId, at: usize) {
        let (old, i) = self.position(n).expect("relocated nodes have a parent");
        let weight = self.mean_weight(p);
        let at = if old == p && i < at { at - 1 } else { at };
        self.unlink(old, i);
        self.attach(p, at, n, weight);
        self.normalize(old);
    }

    /// Places `node` among `p`'s children. A container running along `p`'s
    /// axis is dissolved into `p` so axes keep alternating.
    fn attach(&mut self, p: NodeId, at: usize, node: NodeId, weight: f64) {
        if self.axis(node) == self.axis(p) {
            let Some(Node {
                kind: Kind::Container { children, .. },
                ..
            }) = self.nodes.remove(node)
            else {
                unreachable!("only containers have an axis");
            };
            let total: f64 = children.iter().map(|c| c.weight).sum();
            for (k, c) in children.into_iter().enumerate() {
                self.nodes[c.node].parent = Some(p);
                let weight = weight * c.weight / total;
                self.children_mut(p).insert(
                    at + k,
                    Child {
                        node: c.node,
                        weight,
                    },
                );
            }
        } else {
            self.nodes[node].parent = Some(p);
            self.children_mut(p).insert(at, Child { node, weight });
        }
    }

    fn unlink(&mut self, p: NodeId, i: usize) -> Child {
        let Kind::Container {
            children, expanded, ..
        } = &mut self.nodes[p].kind
        else {
            unreachable!("parents are containers");
        };
        let child = children.remove(i);
        if *expanded == Some(child.node) {
            *expanded = None;
        }
        self.nodes[child.node].parent = None;
        child
    }

    fn detach(&mut self, n: NodeId) {
        match self.position(n) {
            Some((p, i)) => {
                self.unlink(p, i);
                self.normalize(p);
            }
            None => self.root = None,
        }
    }

    /// Repairs a container that lost a child.
    fn normalize(&mut self, c: NodeId) {
        if self.children(c).len() >= 2 {
            return;
        }
        let pos = self.position(c);
        let was_expanded = pos.is_some_and(|(p, _)| self.expanded(p) == Some(c));
        let Some(Node {
            kind: Kind::Container { children, .. },
            ..
        }) = self.nodes.remove(c)
        else {
            unreachable!("only containers are normalized");
        };
        let only = children.into_iter().next();
        match (only, pos) {
            (None, Some((p, i))) => {
                self.unlink_dead(p, i);
                self.normalize(p);
            }
            (None, None) => self.root = None,
            (Some(only), Some((p, i))) => {
                let weight = self.unlink_dead(p, i).weight;
                let leaf_like = self.axis(only.node) != self.axis(p);
                self.attach(p, i, only.node, weight);
                if was_expanded
                    && leaf_like
                    && let Kind::Container { expanded, .. } = &mut self.nodes[p].kind
                {
                    *expanded = Some(only.node);
                }
            }
            (Some(only), None) => {
                self.nodes[only.node].parent = None;
                self.root = Some(only.node);
            }
        }
    }

    /// Like `unlink`, for a child whose node is already gone.
    fn unlink_dead(&mut self, p: NodeId, i: usize) -> Child {
        let Kind::Container {
            children, expanded, ..
        } = &mut self.nodes[p].kind
        else {
            unreachable!("parents are containers");
        };
        let child = children.remove(i);
        if *expanded == Some(child.node) {
            *expanded = None;
        }
        child
    }

    fn new_leaf(&mut self, w: W) -> NodeId {
        let n = self.nodes.insert(Node {
            parent: None,
            kind: Kind::Leaf(w),
        });
        self.leaves.insert(w, n);
        n
    }

    fn new_container(&mut self, axis: Axis) -> NodeId {
        self.nodes.insert(Node {
            parent: None,
            kind: Kind::Container {
                axis,
                children: Vec::new(),
                expanded: None,
            },
        })
    }

    fn position(&self, n: NodeId) -> Option<(NodeId, usize)> {
        let p = self.nodes[n].parent?;
        let i = self.children(p).iter().position(|c| c.node == n)?;
        Some((p, i))
    }

    fn axis(&self, n: NodeId) -> Option<Axis> {
        match &self.nodes[n].kind {
            Kind::Container { axis, .. } => Some(*axis),
            Kind::Leaf(_) => None,
        }
    }

    fn expanded(&self, n: NodeId) -> Option<NodeId> {
        match &self.nodes[n].kind {
            Kind::Container { expanded, .. } => *expanded,
            Kind::Leaf(_) => None,
        }
    }

    fn children(&self, n: NodeId) -> &[Child] {
        match &self.nodes[n].kind {
            Kind::Container { children, .. } => children,
            Kind::Leaf(_) => &[],
        }
    }

    fn children_mut(&mut self, n: NodeId) -> &mut Vec<Child> {
        match &mut self.nodes[n].kind {
            Kind::Container { children, .. } => children,
            Kind::Leaf(_) => unreachable!("leaves have no children"),
        }
    }

    fn mean_weight(&self, p: NodeId) -> f64 {
        let children = self.children(p);
        if children.is_empty() {
            1.0
        } else {
            children.iter().map(|c| c.weight).sum::<f64>() / children.len() as f64
        }
    }

    fn visible(&self, n: NodeId, vis: &impl Fn(W) -> bool) -> bool {
        match &self.nodes[n].kind {
            Kind::Leaf(w) => vis(*w),
            Kind::Container { children, .. } => children.iter().any(|c| self.visible(c.node, vis)),
        }
    }

    #[cfg(test)]
    pub(crate) fn check(&self) {
        let reachable = match self.root {
            Some(r) => {
                assert!(self.nodes[r].parent.is_none());
                self.check_node(r)
            }
            None => 0,
        };
        assert_eq!(reachable, self.nodes.len(), "orphaned nodes");
        for (&w, &n) in &self.leaves {
            assert!(matches!(self.nodes[n].kind, Kind::Leaf(x) if x == w));
        }
    }

    #[cfg(test)]
    fn check_node(&self, n: NodeId) -> usize {
        match &self.nodes[n].kind {
            Kind::Leaf(_) => 1,
            Kind::Container {
                axis,
                children,
                expanded,
            } => {
                assert!(
                    children.len() >= 2,
                    "container with {} children",
                    children.len()
                );
                if let Some(e) = expanded {
                    assert!(children.iter().any(|c| c.node == *e));
                }
                let mut count = 1;
                for c in children {
                    assert!(c.weight > 0.0 && c.weight.is_finite());
                    assert_eq!(self.nodes[c.node].parent, Some(n));
                    assert_ne!(self.axis(c.node), Some(*axis), "axes must alternate");
                    count += self.check_node(c.node);
                }
                count
            }
        }
    }
}

fn span(r: Rect, axis: Axis) -> (f64, f64) {
    match axis {
        Axis::Horizontal => (r.x, r.w),
        Axis::Vertical => (r.y, r.h),
    }
}

fn with_span(r: Rect, axis: Axis, start: f64, len: f64) -> Rect {
    match axis {
        Axis::Horizontal => Rect::new(start, r.y, len, r.h),
        Axis::Vertical => Rect::new(r.x, start, r.w, len),
    }
}

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