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>

1 files changed, 75 insertions, 0 deletions

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)