From 190fbf66d462748e42e67b6bfe36f1ef28635c23 Mon Sep 17 00:00:00 2001
From: Luis Useche <useche@gmail.com>
Date: Tue, 2 Apr 2024 08:29:57 -0700
Subject: [PATCH] reconfigure tree properly after closing view (#10004)

This patch merges the last child of a container node to the parent. This
avoids the bug where uneven spaced Views would be created. To reproduce:
1. `vsplit`
2. `split`
3. `quit`
4. `vsplit`

With this patch the bug cannot be seen anymore.
---
 helix-view/src/tree.rs | 101 +++++++++++++++++++++++++++++++++--------
 1 file changed, 81 insertions(+), 20 deletions(-)

diff --git a/helix-view/src/tree.rs b/helix-view/src/tree.rs
index 4c9eba0fd125..b7ad5aa9de05 100644
--- a/helix-view/src/tree.rs
+++ b/helix-view/src/tree.rs
@@ -214,33 +214,56 @@ impl Tree {
         node
     }
 
-    pub fn remove(&mut self, index: ViewId) {
-        let mut stack = Vec::new();
+    /// Get a mutable reference to a [Container] by index.
+    /// # Panics
+    /// Panics if `index` is not in self.nodes, or if the node's content is not a [Content::Container].
+    fn container_mut(&mut self, index: ViewId) -> &mut Container {
+        match &mut self.nodes[index] {
+            Node {
+                content: Content::Container(container),
+                ..
+            } => container,
+            _ => unreachable!(),
+        }
+    }
+
+    fn remove_or_replace(&mut self, child: ViewId, replacement: Option<ViewId>) {
+        let parent = self.nodes[child].parent;
+
+        self.nodes.remove(child);
 
+        let container = self.container_mut(parent);
+        let pos = container
+            .children
+            .iter()
+            .position(|&item| item == child)
+            .unwrap();
+
+        if let Some(new) = replacement {
+            container.children[pos] = new;
+            self.nodes[new].parent = parent;
+        } else {
+            container.children.remove(pos);
+        }
+    }
+
+    pub fn remove(&mut self, index: ViewId) {
         if self.focus == index {
             // focus on something else
             self.focus = self.prev();
         }
 
-        stack.push(index);
+        let parent = self.nodes[index].parent;
+        let parent_is_root = parent == self.root;
 
-        while let Some(index) = stack.pop() {
-            let parent_id = self.nodes[index].parent;
-            if let Node {
-                content: Content::Container(container),
-                ..
-            } = &mut self.nodes[parent_id]
-            {
-                if let Some(pos) = container.children.iter().position(|&child| child == index) {
-                    container.children.remove(pos);
-                    // TODO: if container now only has one child, remove it and place child in parent
-                    if container.children.is_empty() && parent_id != self.root {
-                        // if container now empty, remove it
-                        stack.push(parent_id);
-                    }
-                }
-            }
-            self.nodes.remove(index);
+        self.remove_or_replace(index, None);
+
+        let parent_container = self.container_mut(parent);
+        if parent_container.children.len() == 1 && !parent_is_root {
+            // Lets merge the only child back to its grandparent so that Views
+            // are equally spaced.
+            let sibling = parent_container.children.pop().unwrap();
+            self.remove_or_replace(parent, Some(sibling));
         }
 
         self.recalculate()
@@ -873,4 +896,42 @@ mod test {
         assert_eq!(doc_id(&tree, l2), Some(doc_r0));
         assert_eq!(doc_id(&tree, r0), Some(doc_l0));
     }
+
+    #[test]
+    fn all_vertical_views_have_same_width() {
+        let tree_area_width = 180;
+        let mut tree = Tree::new(Rect {
+            x: 0,
+            y: 0,
+            width: tree_area_width,
+            height: 80,
+        });
+        let mut view = View::new(DocumentId::default(), GutterConfig::default());
+        view.area = Rect::new(0, 0, 180, 80);
+        tree.insert(view);
+
+        let view = View::new(DocumentId::default(), GutterConfig::default());
+        tree.split(view, Layout::Vertical);
+
+        let view = View::new(DocumentId::default(), GutterConfig::default());
+        tree.split(view, Layout::Horizontal);
+
+        tree.remove(tree.focus);
+
+        let view = View::new(DocumentId::default(), GutterConfig::default());
+        tree.split(view, Layout::Vertical);
+
+        // Make sure that we only have one level in the tree.
+        assert_eq!(3, tree.views().count());
+        assert_eq!(
+            vec![
+                tree_area_width / 3,
+                tree_area_width / 3,
+                tree_area_width / 3 - 2 // Rounding in `recalculate`.
+            ],
+            tree.views()
+                .map(|(view, _)| view.area.width)
+                .collect::<Vec<_>>()
+        );
+    }
 }