Splits: Preserve perpendicular shares on batch removal

kitty/layout/splits.py

@@ -154,6 +154,24 @@ class Pair:
 
         tune(self)
 
+    def preserve_weights_on_removal(self, removed: Collection[int]) -> bool:
+        # Each perpendicular subtree is one unit along this split axis. If
+        # some of its windows survive, retain its full share on this axis.
+        def tune(child: Pair | int | None, weight: float) -> float:
+            if isinstance(child, Pair):
+                if child.horizontal != self.horizontal:
+                    return weight if child.preserve_weights_on_removal(removed) else 0.0
+                if child.is_redundant:
+                    return tune(child.one or child.two, weight)
+                one = tune(child.one, weight * child.bias)
+                two = tune(child.two, weight * (1 - child.bias))
+                if one > 0 and two > 0:
+                    child.bias = one / (one + two)
+                return one + two
+            return weight if child is not None and child not in removed else 0.0
+
+        return tune(self, 1.0) > 0
+
     def remove_windows(self, window_ids: Collection[int]) -> None:
         if isinstance(self.one, int) and self.one in window_ids:
             self.one = None
@@ -618,15 +636,8 @@ class Splits(Layout):
 
     def remove_windows(self, *windows_to_remove: int) -> None:
         root = self.pairs_root
-        plans: dict[tuple[bool, frozenset[int]], dict[int, float]] = {}
         if self.layout_opts.proportional:
-            for group_id in windows_to_remove:
-                pair = root.pair_for_window(group_id)
-                if pair is not None and (container := self.proportional_container(group_id, pair.horizontal)) is not None:
-                    weights = container.window_weights()
-                    survivors = frozenset(weights).difference(windows_to_remove)
-                    if len(survivors) > 1:
-                        plans[container.horizontal, survivors] = weights
+            root.preserve_weights_on_removal(frozenset(windows_to_remove))
         for pair in root.self_and_descendants():
             pair.remove_windows(windows_to_remove)
         root.collapse_redundant_pairs()
@@ -634,11 +645,6 @@ class Splits(Layout):
             q = root.one or root.two
             if isinstance(q, Pair):
                 self.pairs_root = q
-        for (horizontal, survivors), weights in plans.items():
-            for pair in self.pairs_root.self_and_descendants():
-                if pair.horizontal == horizontal and frozenset(pair.all_window_ids()) == survivors:
-                    pair.set_window_weights(weights)
-                    break
 
     def proportional_container(self, group_id: int, horizontal: bool) -> Pair | None:
         root = self.pairs_root

kitty_tests/layout.py

@@ -858,11 +858,11 @@ class TestLayout(BaseTest):
 
 
 class TestProportionalSplits(BaseTest):
-    def make_layout(self, shape=None, **options):
+    def make_layout(self, shape=None, num=3, **options):
         q = create_layout(Splits)
         q.layout_opts = SplitsLayoutOpts({'proportional': 'yes', **options})
         windows = create_windows(q, num=0)
-        for i in range(1, 4 if shape else 2):
+        for i in range(1, num + 1 if shape else 2):
             q.add_window(windows, Window(i), location='vsplit')
         if shape:
             q.pairs_root.unserialize(shape, lambda x: x)
@@ -933,6 +933,19 @@ class TestProportionalSplits(BaseTest):
         q.remove_windows(1, 3, 5)
         self.check_weights(q, {2: 0.5, 4: 0.5})
 
+    def test_proportional_batch_close_mixed_axes(self):
+        for horizontal in (True, False):
+            shape = {
+                'horizontal': horizontal,
+                'one': 1,
+                'two': {'horizontal': horizontal, 'one': {'horizontal': not horizontal, 'one': 2, 'two': 3}, 'two': 4},
+            }
+            q, windows = self.make_layout(shape, num=4)
+            q.remove_windows(1, 2)
+            # The surviving center pane keeps its column's width, even though
+            # the other pane in that column was removed at the same time.
+            self.check_weights(q, {3: 0.5, 4: 0.5})
+
     def test_proportional_mixed_axes(self):
         q, windows = self.make_layout({'bias': 0.4, 'one': 1, 'two': {'horizontal': False, 'bias': 0.25, 'one': 2, 'two': 3}})
         q.add_window(windows, Window(4), location='hsplit', next_to=windows.id_map[2])