perf(treesitter): don't fetch parser for each fold line

**Problem:** The treesitter `foldexpr` calls `get_parser()` for each
line in the buffer when calculating folds. This can be incredibly slow
for buffers where a parser cannot be found (because the result is not
cached), and exponentially more so when the user has many
`runtimepath`s.

**Solution:** Only fetch the parser when it is needed; that is, only
when initializing fold data for a buffer.

Co-authored-by: Jongwook Choi <wookayin@gmail.com>
Co-authored-by: Justin M. Keyes <justinkz@gmail.com>

3 files changed, 102 insertions, 9 deletions

diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index b7606e65f5..17f85154c2 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -287,6 +287,8 @@ PERFORMANCE
   highlighting.
 • LSP diagnostics and inlay hints are de-duplicated (new requests cancel
   inflight requests). This greatly improves performance with slow LSP servers.
+• 10x speedup for |vim.treesitter.foldexpr()| (when no parser exists for the
+  buffer).
 
 PLUGINS
 
diff --git a/runtime/lua/vim/treesitter/_fold.lua b/runtime/lua/vim/treesitter/_fold.lua
index 207ac1ab67..d16013eca2 100644
--- a/runtime/lua/vim/treesitter/_fold.lua
+++ b/runtime/lua/vim/treesitter/_fold.lua
@@ -19,14 +19,19 @@ local api = vim.api
 ---The range on which to evaluate foldexpr.
 ---When in insert mode, the evaluation is deferred to InsertLeave.
 ---@field foldupdate_range? Range2
+---
+---The treesitter parser associated with this buffer.
+---@field parser? vim.treesitter.LanguageTree
 local FoldInfo = {}
 FoldInfo.__index = FoldInfo
 
 ---@private
-function FoldInfo.new()
+---@param bufnr integer
+function FoldInfo.new(bufnr)
   return setmetatable({
     levels0 = {},
     levels = {},
+    parser = ts.get_parser(bufnr, nil, { error = false }),
   }, FoldInfo)
 end
 
@@ -69,7 +74,10 @@ local function compute_folds_levels(bufnr, info, srow, erow, parse_injections)
   srow = srow or 0
   erow = erow or api.nvim_buf_line_count(bufnr)
 
-  local parser = assert(ts.get_parser(bufnr, nil, { error = false }))
+  local parser = info.parser
+  if not parser then
+    return
+  end
 
   parser:parse(parse_injections and { srow, erow } or nil)
 
@@ -347,13 +355,21 @@ function M.foldexpr(lnum)
   lnum = lnum or vim.v.lnum
   local bufnr = api.nvim_get_current_buf()
 
-  local parser = ts.get_parser(bufnr, nil, { error = false })
-  if not parser then
-    return '0'
-  end
-
   if not foldinfos[bufnr] then
-    foldinfos[bufnr] = FoldInfo.new()
+    foldinfos[bufnr] = FoldInfo.new(bufnr)
+    api.nvim_create_autocmd('BufUnload', {
+      buffer = bufnr,
+      once = true,
+      callback = function()
+        foldinfos[bufnr] = nil
+      end,
+    })
+
+    local parser = foldinfos[bufnr].parser
+    if not parser then
+      return '0'
+    end
+
     compute_folds_levels(bufnr, foldinfos[bufnr])
 
     parser:register_cbs({
@@ -383,7 +399,7 @@ api.nvim_create_autocmd('OptionSet', {
       or foldinfos[buf] and { buf }
       or {}
     for _, bufnr in ipairs(bufs) do
-      foldinfos[bufnr] = FoldInfo.new()
+      foldinfos[bufnr] = FoldInfo.new(bufnr)
       api.nvim_buf_call(bufnr, function()
         compute_folds_levels(bufnr, foldinfos[bufnr])
       end)
diff --git a/test/functional/treesitter/fold_spec.lua b/test/functional/treesitter/fold_spec.lua
index e38e58ff92..9f7fdf529f 100644
--- a/test/functional/treesitter/fold_spec.lua
+++ b/test/functional/treesitter/fold_spec.lua
@@ -5,6 +5,7 @@ local Screen = require('test.functional.ui.screen')
 local clear = n.clear
 local eq = t.eq
 local insert = n.insert
+local write_file = t.write_file
 local exec_lua = n.exec_lua
 local command = n.command
 local feed = n.feed
@@ -767,4 +768,78 @@ t2]])
     ]],
     }
   end)
+
+  it("doesn't call get_parser too often when parser is not available", function()
+    -- spy on vim.treesitter.get_parser() to keep track of how many times it is called
+    exec_lua(function()
+      _G.count = 0
+      vim.treesitter.get_parser = (function(wrapped)
+        return function(...)
+          _G.count = _G.count + 1
+          return wrapped(...)
+        end
+      end)(vim.treesitter.get_parser)
+    end)
+
+    insert(test_text)
+    command [[
+      set filetype=some_filetype_without_treesitter_parser
+      set foldmethod=expr foldexpr=v:lua.vim.treesitter.foldexpr() foldcolumn=1 foldlevel=0
+    ]]
+
+    -- 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))
+    end
+
+    eq(
+      1,
+      exec_lua [[ return _G.count ]],
+      'count should not be as high as the # of lines; actually only once for the buffer.'
+    )
+  end)
+
+  it('can detect a new parser and refresh folds accordingly', function()
+    write_file('test_fold_file.txt', test_text)
+    command [[
+      e test_fold_file.txt
+      set filetype=some_filetype_without_treesitter_parser
+      set foldmethod=expr foldexpr=v:lua.vim.treesitter.foldexpr() foldcolumn=1 foldlevel=0
+    ]]
+
+    -- 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))
+    end
+
+    -- reload buffer as c filetype to simulate new parser being found
+    feed('GA// vim: ft=c<Esc>')
+    command([[w | e]])
+
+    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)
 end)