fix(treesitter): avoid computing fold levels for empty buffer

Problem:  Computing fold levels for an empty buffer (somehow) breaks the
          parser state, resulting in a broken highlighter and foldexpr.
          Cached foldexpr parser is invalid after filetype has changed.
Solution: Avoid computing fold levels for empty buffer.
          Clear cached foldinfos upon `FileType`.

3 files changed, 79 insertions, 9 deletions

diff --git a/runtime/lua/vim/treesitter/_fold.lua b/runtime/lua/vim/treesitter/_fold.lua
index 38318347a7..1064004320 100644
--- a/runtime/lua/vim/treesitter/_fold.lua
+++ b/runtime/lua/vim/treesitter/_fold.lua
@@ -75,7 +75,15 @@ local function compute_folds_levels(bufnr, info, srow, erow, callback)
   erow = erow or api.nvim_buf_line_count(bufnr)
 
   local parser = info.parser
-  if not parser then
+  if
+    not parser
+    -- Parsing an empty buffer results in problems with the parsing state,
+    -- resulting in both a broken highlighter and foldexpr.
+    or api.nvim_buf_line_count(bufnr) == 1
+      and api.nvim_buf_call(bufnr, function()
+        return vim.fn.line2byte(1) <= 0
+      end)
+  then
     return
   end
 
@@ -380,7 +388,7 @@ function M.foldexpr(lnum)
 
   if not foldinfos[bufnr] then
     foldinfos[bufnr] = FoldInfo.new(bufnr)
-    api.nvim_create_autocmd({ 'BufUnload', 'VimEnter' }, {
+    api.nvim_create_autocmd({ 'BufUnload', 'VimEnter', 'FileType' }, {
       buffer = bufnr,
       once = true,
       callback = function()
diff --git a/test/functional/ex_cmds/swapfile_preserve_recover_spec.lua b/test/functional/ex_cmds/swapfile_preserve_recover_spec.lua
index 2974564f70..848db7d088 100644
--- a/test/functional/ex_cmds/swapfile_preserve_recover_spec.lua
+++ b/test/functional/ex_cmds/swapfile_preserve_recover_spec.lua
@@ -150,6 +150,58 @@ describe('swapfile detection', function()
     rmdir(swapdir)
   end)
 
+  it('redrawing during prompt does not break treesitter', function()
+    local testfile = 'Xtest_swapredraw.lua'
+    write_file(
+      testfile,
+      [[
+vim.o.foldmethod = 'expr'
+vim.o.foldexpr = 'v:lua.vim.treesitter.foldexpr()'
+vim.defer_fn(function()
+  vim.api.nvim__redraw({ valid = false })
+end, 500)
+pcall(vim.cmd.edit, 'Xtest_swapredraw.lua')
+    ]]
+    )
+    exec(init)
+    command('edit! ' .. testfile)
+    command('preserve')
+    local nvim2 = n.new_session(true, { args = { '--clean', '--embed' }, merge = false })
+    set_session(nvim2)
+    local screen2 = Screen.new(100, 40)
+    screen2:add_extra_attr_ids({
+      [100] = { foreground = Screen.colors.NvimLightGrey2 },
+      [101] = { foreground = Screen.colors.NvimLightGreen },
+      [102] = {
+        foreground = Screen.colors.NvimLightGrey4,
+        background = Screen.colors.NvimDarkGrey1,
+      },
+      [104] = { foreground = Screen.colors.NvimLightCyan },
+      [105] = { foreground = Screen.colors.NvimDarkGrey4 },
+      [106] = {
+        foreground = Screen.colors.NvimDarkGrey3,
+        background = Screen.colors.NvimLightGrey3,
+      },
+    })
+    exec(init)
+    command('autocmd! nvim.swapfile') -- Delete the default handler (which skips the dialog).
+    feed(':edit ' .. testfile .. '<CR>')
+    feed('E:source<CR>')
+    screen2:sleep(1000)
+    feed('E')
+    screen2:expect([[
+      {100:^vim.o.foldmethod} {100:=} {101:'expr'}                                                                           |
+      {100:vim.o.foldexpr} {100:=} {101:'v:lua.vim.treesitter.foldexpr()'}                                                  |
+      {102:+--  3 lines: vim.defer_fn(function()·······························································}|
+      {104:pcall}{100:(vim.cmd.edit,} {101:'Xtest_swapredraw.lua'}{100:)}                                                         |
+                                                                                                          |
+      {105:~                                                                                                   }|*33
+      {106:Xtest_swapredraw.lua                                                              1,1            All}|
+                                                                                                          |
+    ]])
+    nvim2:close()
+  end)
+
   it('always show swapfile dialog #8840 #9027', function()
     local testfile = 'Xtest_swapdialog_file1'
 
diff --git a/test/functional/treesitter/fold_spec.lua b/test/functional/treesitter/fold_spec.lua
index ac58df4bba..1212212d62 100644
--- a/test/functional/treesitter/fold_spec.lua
+++ b/test/functional/treesitter/fold_spec.lua
@@ -811,17 +811,19 @@ t2]])
     ]]
 
     -- foldexpr will return '0' for all lines
-    local levels = get_fold_levels() ---@type integer[]
-    eq(19, #levels)
-    for lnum, level in ipairs(levels) do
-      eq('0', level, string.format("foldlevel[%d] == %s; expected '0'", lnum, level))
+    local function expect_no_folds()
+      local levels = get_fold_levels() ---@type integer[]
+      eq(19, #levels)
+      for lnum, level in ipairs(levels) do
+        eq('0', level, string.format("foldlevel[%d] == %s; expected '0'", lnum, level))
+      end
     end
+    expect_no_folds()
 
     -- reload buffer as c filetype to simulate new parser being found
     feed('GA// vim: ft=c<Esc>')
     command([[write | edit]])
-
-    eq({
+    local foldlevels = {
       [1] = '>1',
       [2] = '1',
       [3] = '1',
@@ -841,6 +843,14 @@ t2]])
       [17] = '3',
       [18] = '2',
       [19] = '1',
-    }, get_fold_levels())
+    }
+    eq(foldlevels, get_fold_levels())
+
+    -- only changing filetype should change the parser again
+    command('set ft=some_filetype_without_treesitter_parser')
+    expect_no_folds()
+
+    command('set ft=c')
+    eq(foldlevels, get_fold_levels())
   end)
 end)