fix(ui): make `winbar` properly equalize window heights for local value

Fixes `'winbar'` not properly equalizing window heights for splits when
the global value is empty and a window-local value is set instead.

2 files changed, 54 insertions, 5 deletions

diff --git a/src/nvim/window.c b/src/nvim/window.c
index 51208e7178..1bff5933ea 100644
--- a/src/nvim/window.c
+++ b/src/nvim/window.c
@@ -2028,6 +2028,37 @@ void win_move_after(win_T *win1, win_T *win2)
   win2->w_pos_changed = true;
 }
 
+/// Compute maximum number of windows that can fit within "height" in frame "fr".
+static int get_maximum_wincount(frame_T *fr, int height)
+{
+  if (fr->fr_layout != FR_COL) {
+    return (height / (p_wmh + STATUS_HEIGHT + frame2win(fr)->w_winbar_height));
+  } else if (global_winbar_height()) {
+    // If winbar is globally enabled, no need to check each window for it.
+    return (height / (p_wmh + STATUS_HEIGHT + 1));
+  }
+
+  frame_T *frp;
+  int total_wincount = 0;
+
+  // First, try to fit all child frames of "fr" into "height"
+  FOR_ALL_FRAMES(frp, fr->fr_child) {
+    win_T *wp = frame2win(frp);
+
+    if (height < (p_wmh + STATUS_HEIGHT + wp->w_winbar_height)) {
+      break;
+    }
+    height -= p_wmh + STATUS_HEIGHT + wp->w_winbar_height;
+    total_wincount += 1;
+  }
+
+  // If we still have enough room for more windows, just use the default winbar height (which is 0)
+  // in order to get the amount of windows that'd fit in the remaining space
+  total_wincount += height / (p_wmh + STATUS_HEIGHT);
+
+  return total_wincount;
+}
+
 /// Make all windows the same height.
 ///'next_curwin' will soon be the current window, make sure it has enough rows.
 ///
@@ -2227,7 +2258,7 @@ static void win_equal_rec(win_T *next_curwin, bool current, frame_T *topfr, int
       } else {
         extra_sep = 0;
       }
-      totwincount = (n + extra_sep) / (p_wmh + STATUS_HEIGHT + global_winbar_height());
+      totwincount = get_maximum_wincount(topfr, n + extra_sep);
       has_next_curwin = frame_has_win(topfr, next_curwin);
 
       /*
@@ -2261,8 +2292,7 @@ static void win_equal_rec(win_T *next_curwin, bool current, frame_T *topfr, int
             }
           } else {
             // These windows don't use up room.
-            totwincount -= (n + (fr->fr_next == NULL
-                                 ? extra_sep : 0)) / (p_wmh + STATUS_HEIGHT + global_winbar_height());
+            totwincount -= get_maximum_wincount(fr, (n + (fr->fr_next == NULL ? extra_sep : 0)));
           }
           room -= new_size - n;
           if (room < 0) {
@@ -2307,8 +2337,7 @@ static void win_equal_rec(win_T *next_curwin, bool current, frame_T *topfr, int
       } else {
         // Compute the maximum number of windows vert. in "fr".
         n = frame_minheight(fr, NOWIN);
-        wincount = (n + (fr->fr_next == NULL ? extra_sep : 0))
-                   / (p_wmh + STATUS_HEIGHT + global_winbar_height());
+        wincount = get_maximum_wincount(fr, (n + (fr->fr_next == NULL ? extra_sep : 0)));
         m = frame_minheight(fr, next_curwin);
         if (has_next_curwin) {
           hnc = frame_has_win(fr, next_curwin);
diff --git a/test/functional/ui/winbar_spec.lua b/test/functional/ui/winbar_spec.lua
index abc6088801..e5cdc53b93 100644
--- a/test/functional/ui/winbar_spec.lua
+++ b/test/functional/ui/winbar_spec.lua
@@ -194,4 +194,24 @@ describe('winbar', function()
     meths.input_mouse('left', 'press', '', 0, 5, 1)
     eq({5, 1}, meths.win_get_cursor(0))
   end)
+  it('properly equalizes window height for window-local value', function()
+    command('set equalalways | set winbar= | setlocal winbar=a | split')
+    command('setlocal winbar= | split')
+    command('setlocal winbar=b | split')
+    screen:expect([[
+      {1:b                                                           }|
+      ^                                                            |
+      {4:[No Name]                                                   }|
+      {1:b                                                           }|
+                                                                  |
+      {2:[No Name]                                                   }|
+                                                                  |
+      {3:~                                                           }|
+      {2:[No Name]                                                   }|
+      {1:a                                                           }|
+                                                                  |
+      {2:[No Name]                                                   }|
+                                                                  |
+    ]])
+  end)
 end)