fix(treesitter): update folds in all relevant windows (#24230)

Problem: When using treesitter foldexpr,
* :diffput/get open diff folds, and
* folds are not updated in other windows that contain the updated
  buffer.

Solution: Update folds in all windows that contain the updated buffer
and use expr foldmethod.

3 files changed, 387 insertions, 89 deletions

diff --git a/runtime/lua/vim/treesitter/_fold.lua b/runtime/lua/vim/treesitter/_fold.lua
index d308657237..a02d0a584d 100644
--- a/runtime/lua/vim/treesitter/_fold.lua
+++ b/runtime/lua/vim/treesitter/_fold.lua
@@ -232,20 +232,40 @@ local M = {}
 ---@type table<integer,TS.FoldInfo>
 local foldinfos = {}
 
-local function recompute_folds()
+--- Update the folds in the windows that contain the buffer and use expr foldmethod (assuming that
+--- the user doesn't use different foldexpr for the same buffer).
+---
+--- Nvim usually automatically updates folds when text changes, but it doesn't work here because
+--- FoldInfo update is scheduled. So we do it manually.
+local function foldupdate(bufnr)
+  local function do_update()
+    for _, win in ipairs(vim.fn.win_findbuf(bufnr)) do
+      api.nvim_win_call(win, function()
+        if vim.wo.foldmethod == 'expr' then
+          vim._foldupdate()
+        end
+      end)
+    end
+  end
+
   if api.nvim_get_mode().mode == 'i' then
     -- foldUpdate() is guarded in insert mode. So update folds on InsertLeave
     api.nvim_create_autocmd('InsertLeave', {
       once = true,
-      callback = vim._foldupdate,
+      callback = do_update,
     })
     return
   end
 
-  vim._foldupdate()
+  do_update()
 end
 
---- Schedule a function only if bufnr is loaded
+--- Schedule a function only if bufnr is loaded.
+--- We schedule fold level computation for the following reasons:
+--- * queries seem to use the old buffer state in on_bytes for some unknown reason;
+--- * to avoid textlock;
+--- * to avoid infinite recursion:
+---   get_folds_levels → parse → _do_callback → on_changedtree → get_folds_levels.
 ---@param bufnr integer
 ---@param fn function
 local function schedule_if_loaded(bufnr, fn)
@@ -261,14 +281,12 @@ end
 ---@param foldinfo TS.FoldInfo
 ---@param tree_changes Range4[]
 local function on_changedtree(bufnr, foldinfo, tree_changes)
-  -- For some reason, queries seem to use the old buffer state in on_bytes.
-  -- Get around this by scheduling and manually updating folds.
   schedule_if_loaded(bufnr, function()
     for _, change in ipairs(tree_changes) do
       local srow, _, erow = Range.unpack4(change)
       get_folds_levels(bufnr, foldinfo, srow, erow)
     end
-    recompute_folds()
+    foldupdate(bufnr)
   end)
 end
 
@@ -289,7 +307,7 @@ local function on_bytes(bufnr, foldinfo, start_row, old_row, new_row)
     end
     schedule_if_loaded(bufnr, function()
       get_folds_levels(bufnr, foldinfo, start_row, end_row_new)
-      recompute_folds()
+      foldupdate(bufnr)
     end)
   end
 end
diff --git a/test/functional/treesitter/fold_spec.lua b/test/functional/treesitter/fold_spec.lua
new file mode 100644
index 0000000000..9ed86e87f1
--- /dev/null
+++ b/test/functional/treesitter/fold_spec.lua
@@ -0,0 +1,361 @@
+local helpers = require('test.functional.helpers')(after_each)
+local clear = helpers.clear
+local eq = helpers.eq
+local insert = helpers.insert
+local exec_lua = helpers.exec_lua
+local command = helpers.command
+local feed = helpers.feed
+local Screen = require('test.functional.ui.screen')
+
+before_each(clear)
+
+describe('treesitter foldexpr', function()
+  clear()
+
+  local test_text = [[
+void ui_refresh(void)
+{
+  int width = INT_MAX, height = INT_MAX;
+  bool ext_widgets[kUIExtCount];
+  for (UIExtension i = 0; (int)i < kUIExtCount; i++) {
+    ext_widgets[i] = true;
+  }
+
+  bool inclusive = ui_override();
+  for (size_t i = 0; i < ui_count; i++) {
+    UI *ui = uis[i];
+    width = MIN(ui->width, width);
+    height = MIN(ui->height, height);
+    foo = BAR(ui->bazaar, bazaar);
+    for (UIExtension j = 0; (int)j < kUIExtCount; j++) {
+      ext_widgets[j] &= (ui->ui_ext[j] || inclusive);
+    }
+  }
+}]]
+
+  local function get_fold_levels()
+    return exec_lua([[
+    local res = {}
+    for i = 1, vim.api.nvim_buf_line_count(0) do
+      res[i] = vim.treesitter.foldexpr(i)
+    end
+    return res
+    ]])
+  end
+
+  it("can compute fold levels", function()
+    insert(test_text)
+
+    exec_lua([[vim.treesitter.get_parser(0, "c")]])
+
+    eq({
+      [1] = '>1',
+      [2] = '1',
+      [3] = '1',
+      [4] = '1',
+      [5] = '>2',
+      [6] = '2',
+      [7] = '2',
+      [8] = '1',
+      [9] = '1',
+      [10] = '>2',
+      [11] = '2',
+      [12] = '2',
+      [13] = '2',
+      [14] = '2',
+      [15] = '>3',
+      [16] = '3',
+      [17] = '3',
+      [18] = '2',
+      [19] = '1' }, get_fold_levels())
+
+  end)
+
+  it("recomputes fold levels after lines are added/removed", function()
+    insert(test_text)
+
+    exec_lua([[vim.treesitter.get_parser(0, "c")]])
+
+    command('1,2d')
+
+    eq({
+      [1] = '0',
+      [2] = '0',
+      [3] = '>1',
+      [4] = '1',
+      [5] = '1',
+      [6] = '0',
+      [7] = '0',
+      [8] = '>1',
+      [9] = '1',
+      [10] = '1',
+      [11] = '1',
+      [12] = '1',
+      [13] = '>2',
+      [14] = '2',
+      [15] = '2',
+      [16] = '1',
+      [17] = '0' }, get_fold_levels())
+
+    command('1put!')
+
+    eq({
+      [1] = '>1',
+      [2] = '1',
+      [3] = '1',
+      [4] = '1',
+      [5] = '>2',
+      [6] = '2',
+      [7] = '2',
+      [8] = '1',
+      [9] = '1',
+      [10] = '>2',
+      [11] = '2',
+      [12] = '2',
+      [13] = '2',
+      [14] = '2',
+      [15] = '>3',
+      [16] = '3',
+      [17] = '3',
+      [18] = '2',
+      [19] = '1' }, get_fold_levels())
+  end)
+
+  it("updates folds in all windows", function()
+    local screen = Screen.new(60, 48)
+    screen:attach()
+    screen:set_default_attr_ids({
+      [1] = {background = Screen.colors.Grey, foreground = Screen.colors.DarkBlue};
+      [2] = {bold = true, foreground = Screen.colors.Blue1};
+      [3] = {bold = true, reverse = true};
+      [4] = {reverse = true};
+    })
+
+    exec_lua([[vim.treesitter.get_parser(0, "c")]])
+    command([[set foldmethod=expr foldexpr=v:lua.vim.treesitter.foldexpr() foldcolumn=1 foldlevel=9]])
+    command('split')
+
+    insert(test_text)
+
+    screen:expect{grid=[[
+      {1:-}void ui_refresh(void)                                      |
+      {1:│}{                                                          |
+      {1:│}  int width = INT_MAX, height = INT_MAX;                   |
+      {1:│}  bool ext_widgets[kUIExtCount];                           |
+      {1:-}  for (UIExtension i = 0; (int)i < kUIExtCount; i++) {     |
+      {1:2}    ext_widgets[i] = true;                                 |
+      {1:2}  }                                                        |
+      {1:│}                                                           |
+      {1:│}  bool inclusive = ui_override();                          |
+      {1:-}  for (size_t i = 0; i < ui_count; i++) {                  |
+      {1:2}    UI *ui = uis[i];                                       |
+      {1:2}    width = MIN(ui->width, width);                         |
+      {1:2}    height = MIN(ui->height, height);                      |
+      {1:2}    foo = BAR(ui->bazaar, bazaar);                         |
+      {1:-}    for (UIExtension j = 0; (int)j < kUIExtCount; j++) {   |
+      {1:3}      ext_widgets[j] &= (ui->ui_ext[j] || inclusive);      |
+      {1:3}    }                                                      |
+      {1:2}  }                                                        |
+      {1:│}^}                                                          |
+      {2:~                                                           }|
+      {2:~                                                           }|
+      {2:~                                                           }|
+      {2:~                                                           }|
+      {3:[No Name] [+]                                               }|
+      {1:-}void ui_refresh(void)                                      |
+      {1:│}{                                                          |
+      {1:│}  int width = INT_MAX, height = INT_MAX;                   |
+      {1:│}  bool ext_widgets[kUIExtCount];                           |
+      {1:-}  for (UIExtension i = 0; (int)i < kUIExtCount; i++) {     |
+      {1:2}    ext_widgets[i] = true;                                 |
+      {1:2}  }                                                        |
+      {1:│}                                                           |
+      {1:│}  bool inclusive = ui_override();                          |
+      {1:-}  for (size_t i = 0; i < ui_count; i++) {                  |
+      {1:2}    UI *ui = uis[i];                                       |
+      {1:2}    width = MIN(ui->width, width);                         |
+      {1:2}    height = MIN(ui->height, height);                      |
+      {1:2}    foo = BAR(ui->bazaar, bazaar);                         |
+      {1:-}    for (UIExtension j = 0; (int)j < kUIExtCount; j++) {   |
+      {1:3}      ext_widgets[j] &= (ui->ui_ext[j] || inclusive);      |
+      {1:3}    }                                                      |
+      {1:2}  }                                                        |
+      {1:│}}                                                          |
+      {2:~                                                           }|
+      {2:~                                                           }|
+      {2:~                                                           }|
+      {4:[No Name] [+]                                               }|
+                                                                  |
+    ]]}
+
+    command('1,2d')
+
+    screen:expect{grid=[[
+      {1: }  ^int width = INT_MAX, height = INT_MAX;                   |
+      {1: }  bool ext_widgets[kUIExtCount];                           |
+      {1:-}  for (UIExtension i = 0; (int)i < kUIExtCount; i++) {     |
+      {1:│}    ext_widgets[i] = true;                                 |
+      {1:│}  }                                                        |
+      {1: }                                                           |
+      {1: }  bool inclusive = ui_override();                          |
+      {1:-}  for (size_t i = 0; i < ui_count; i++) {                  |
+      {1:│}    UI *ui = uis[i];                                       |
+      {1:│}    width = MIN(ui->width, width);                         |
+      {1:│}    height = MIN(ui->height, height);                      |
+      {1:│}    foo = BAR(ui->bazaar, bazaar);                         |
+      {1:-}    for (UIExtension j = 0; (int)j < kUIExtCount; j++) {   |
+      {1:2}      ext_widgets[j] &= (ui->ui_ext[j] || inclusive);      |
+      {1:2}    }                                                      |
+      {1:│}  }                                                        |
+      {1: }}                                                          |
+      {2:~                                                           }|
+      {2:~                                                           }|
+      {2:~                                                           }|
+      {2:~                                                           }|
+      {2:~                                                           }|
+      {2:~                                                           }|
+      {3:[No Name] [+]                                               }|
+      {1: }  int width = INT_MAX, height = INT_MAX;                   |
+      {1: }  bool ext_widgets[kUIExtCount];                           |
+      {1:-}  for (UIExtension i = 0; (int)i < kUIExtCount; i++) {     |
+      {1:│}    ext_widgets[i] = true;                                 |
+      {1:│}  }                                                        |
+      {1: }                                                           |
+      {1: }  bool inclusive = ui_override();                          |
+      {1:-}  for (size_t i = 0; i < ui_count; i++) {                  |
+      {1:│}    UI *ui = uis[i];                                       |
+      {1:│}    width = MIN(ui->width, width);                         |
+      {1:│}    height = MIN(ui->height, height);                      |
+      {1:│}    foo = BAR(ui->bazaar, bazaar);                         |
+      {1:-}    for (UIExtension j = 0; (int)j < kUIExtCount; j++) {   |
+      {1:2}      ext_widgets[j] &= (ui->ui_ext[j] || inclusive);      |
+      {1:2}    }                                                      |
+      {1:│}  }                                                        |
+      {1: }}                                                          |
+      {2:~                                                           }|
+      {2:~                                                           }|
+      {2:~                                                           }|
+      {2:~                                                           }|
+      {2:~                                                           }|
+      {4:[No Name] [+]                                               }|
+                                                                  |
+    ]]}
+
+
+    feed([[O<C-u><C-r>"<BS><Esc>]])
+
+    screen:expect{grid=[[
+      {1:-}void ui_refresh(void)                                      |
+      {1:│}^{                                                          |
+      {1:│}  int width = INT_MAX, height = INT_MAX;                   |
+      {1:│}  bool ext_widgets[kUIExtCount];                           |
+      {1:-}  for (UIExtension i = 0; (int)i < kUIExtCount; i++) {     |
+      {1:2}    ext_widgets[i] = true;                                 |
+      {1:2}  }                                                        |
+      {1:│}                                                           |
+      {1:│}  bool inclusive = ui_override();                          |
+      {1:-}  for (size_t i = 0; i < ui_count; i++) {                  |
+      {1:2}    UI *ui = uis[i];                                       |
+      {1:2}    width = MIN(ui->width, width);                         |
+      {1:2}    height = MIN(ui->height, height);                      |
+      {1:2}    foo = BAR(ui->bazaar, bazaar);                         |
+      {1:-}    for (UIExtension j = 0; (int)j < kUIExtCount; j++) {   |
+      {1:3}      ext_widgets[j] &= (ui->ui_ext[j] || inclusive);      |
+      {1:3}    }                                                      |
+      {1:2}  }                                                        |
+      {1:│}}                                                          |
+      {2:~                                                           }|
+      {2:~                                                           }|
+      {2:~                                                           }|
+      {2:~                                                           }|
+      {3:[No Name] [+]                                               }|
+      {1:-}void ui_refresh(void)                                      |
+      {1:│}{                                                          |
+      {1:│}  int width = INT_MAX, height = INT_MAX;                   |
+      {1:│}  bool ext_widgets[kUIExtCount];                           |
+      {1:-}  for (UIExtension i = 0; (int)i < kUIExtCount; i++) {     |
+      {1:2}    ext_widgets[i] = true;                                 |
+      {1:2}  }                                                        |
+      {1:│}                                                           |
+      {1:│}  bool inclusive = ui_override();                          |
+      {1:-}  for (size_t i = 0; i < ui_count; i++) {                  |
+      {1:2}    UI *ui = uis[i];                                       |
+      {1:2}    width = MIN(ui->width, width);                         |
+      {1:2}    height = MIN(ui->height, height);                      |
+      {1:2}    foo = BAR(ui->bazaar, bazaar);                         |
+      {1:-}    for (UIExtension j = 0; (int)j < kUIExtCount; j++) {   |
+      {1:3}      ext_widgets[j] &= (ui->ui_ext[j] || inclusive);      |
+      {1:3}    }                                                      |
+      {1:2}  }                                                        |
+      {1:│}}                                                          |
+      {2:~                                                           }|
+      {2:~                                                           }|
+      {2:~                                                           }|
+      {4:[No Name] [+]                                               }|
+                                                                  |
+    ]]}
+
+  end)
+
+  it("doesn't open folds in diff mode", function()
+    local screen = Screen.new(60, 36)
+    screen:attach()
+
+    exec_lua([[vim.treesitter.get_parser(0, "c")]])
+    command([[set foldmethod=expr foldexpr=v:lua.vim.treesitter.foldexpr() foldcolumn=1 foldlevel=9]])
+    insert(test_text)
+    command('16d')
+
+    command('new')
+    insert(test_text)
+
+    command('windo diffthis')
+    feed('do')
+
+    screen:expect{grid=[[
+      {1:+ }{2:+--  9 lines: void ui_refresh(void)·······················}|
+      {1:  }  for (size_t i = 0; i < ui_count; i++) {                 |
+      {1:  }    UI *ui = uis[i];                                      |
+      {1:  }    width = MIN(ui->width, width);                        |
+      {1:  }    height = MIN(ui->height, height);                     |
+      {1:  }    foo = BAR(ui->bazaar, bazaar);                        |
+      {1:  }    for (UIExtension j = 0; (int)j < kUIExtCount; j++) {  |
+      {1:  }      ext_widgets[j] &= (ui->ui_ext[j] || inclusive);     |
+      {1:  }    }                                                     |
+      {1:  }  }                                                       |
+      {1:  }}                                                         |
+      {3:~                                                           }|
+      {3:~                                                           }|
+      {3:~                                                           }|
+      {3:~                                                           }|
+      {3:~                                                           }|
+      {3:~                                                           }|
+      {4:[No Name] [+]                                               }|
+      {1:+ }{2:+--  9 lines: void ui_refresh(void)·······················}|
+      {1:  }  for (size_t i = 0; i < ui_count; i++) {                 |
+      {1:  }    UI *ui = uis[i];                                      |
+      {1:  }    width = MIN(ui->width, width);                        |
+      {1:  }    height = MIN(ui->height, height);                     |
+      {1:  }    foo = BAR(ui->bazaar, bazaar);                        |
+      {1:  }    for (UIExtension j = 0; (int)j < kUIExtCount; j++) {  |
+      {1:  }      ext_widgets[j] &= (ui->ui_ext[j] || inclusive);     |
+      {1:  }    ^}                                                     |
+      {1:  }  }                                                       |
+      {1:  }}                                                         |
+      {3:~                                                           }|
+      {3:~                                                           }|
+      {3:~                                                           }|
+      {3:~                                                           }|
+      {3:~                                                           }|
+      {5:[No Name] [+]                                               }|
+                                                                  |
+    ]], attr_ids={
+      [1] = {background = Screen.colors.Grey, foreground = Screen.colors.Blue4};
+      [2] = {background = Screen.colors.LightGrey, foreground = Screen.colors.Blue4};
+      [3] = {foreground = Screen.colors.Blue, bold = true};
+      [4] = {reverse = true};
+      [5] = {reverse = true, bold = true};
+    }}
+  end)
+
+end)
diff --git a/test/functional/treesitter/parser_spec.lua b/test/functional/treesitter/parser_spec.lua
index 7cfe5b69de..ae26b92f52 100644
--- a/test/functional/treesitter/parser_spec.lua
+++ b/test/functional/treesitter/parser_spec.lua
@@ -885,87 +885,6 @@ int x = INT_MAX;
     end)
   end)
 
-  it("can fold via foldexpr", function()
-    insert(test_text)
-
-    local function get_fold_levels()
-      return exec_lua([[
-        local res = {}
-        for i = 1, vim.api.nvim_buf_line_count(0) do
-          res[i] = vim.treesitter.foldexpr(i)
-        end
-        return res
-      ]])
-    end
-
-    exec_lua([[vim.treesitter.get_parser(0, "c")]])
-
-    eq({
-      [1] = '>1',
-      [2] = '1',
-      [3] = '1',
-      [4] = '1',
-      [5] = '>2',
-      [6] = '2',
-      [7] = '2',
-      [8] = '1',
-      [9] = '1',
-      [10] = '>2',
-      [11] = '2',
-      [12] = '2',
-      [13] = '2',
-      [14] = '2',
-      [15] = '>3',
-      [16] = '3',
-      [17] = '3',
-      [18] = '2',
-      [19] = '1' }, get_fold_levels())
-
-    helpers.command('1,2d')
-
-    eq({
-      [1] = '0',
-      [2] = '0',
-      [3] = '>1',
-      [4] = '1',
-      [5] = '1',
-      [6] = '0',
-      [7] = '0',
-      [8] = '>1',
-      [9] = '1',
-      [10] = '1',
-      [11] = '1',
-      [12] = '1',
-      [13] = '>2',
-      [14] = '2',
-      [15] = '2',
-      [16] = '1',
-      [17] = '0' }, get_fold_levels())
-
-    helpers.command('1put!')
-
-    eq({
-      [1] = '>1',
-      [2] = '1',
-      [3] = '1',
-      [4] = '1',
-      [5] = '>2',
-      [6] = '2',
-      [7] = '2',
-      [8] = '1',
-      [9] = '1',
-      [10] = '>2',
-      [11] = '2',
-      [12] = '2',
-      [13] = '2',
-      [14] = '2',
-      [15] = '>3',
-      [16] = '3',
-      [17] = '3',
-      [18] = '2',
-      [19] = '1' }, get_fold_levels())
-  end)
-
   it('tracks the root range properly (#22911)', function()
     insert([[
       int main() {