perf(treesitter): cache queries strongly

**Problem:** Query parsing uses a weak cache which is invalidated
frequently

**Solution:** Make the cache strong, and invalidate it manually when
necessary (that is, when `rtp` is changed or `query.set()` is called)

Co-authored-by: Christian Clason <c.clason@uni-graz.at>

3 files changed, 48 insertions, 5 deletions

diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 810f40180a..4f4bfe9ecc 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -294,6 +294,9 @@ PERFORMANCE
   inflight requests). This greatly improves performance with slow LSP servers.
 • 10x speedup for |vim.treesitter.foldexpr()| (when no parser exists for the
   buffer).
+• Strong |treesitter-query| caching makes repeat  |vim.treesitter.query.get()|
+  and |vim.treesitter.query.parse()| calls significantly faster for large
+  queries.
 
 PLUGINS
 
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index b9bcbe9a80..b0b0fecd38 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -262,6 +262,7 @@ local explicit_queries = setmetatable({}, {
 ---@param query_name string Name of the query (e.g., "highlights")
 ---@param text string Query text (unparsed).
 function M.set(lang, query_name, text)
+  M.get:clear(lang, query_name)
   explicit_queries[lang][query_name] = M.parse(lang, text)
 end
 
@@ -284,7 +285,15 @@ M.get = memoize('concat-2', function(lang, query_name)
   end
 
   return M.parse(lang, query_string)
-end)
+end, false)
+
+api.nvim_create_autocmd('OptionSet', {
+  pattern = { 'runtimepath' },
+  group = api.nvim_create_augroup('ts_query_cache_reset', { clear = true }),
+  callback = function()
+    M.get:clear()
+  end,
+})
 
 --- Parses a {query} string and returns a `Query` object (|lua-treesitter-query|), which can be used
 --- to search the tree for the query patterns (via |Query:iter_captures()|, |Query:iter_matches()|),
@@ -316,7 +325,7 @@ M.parse = memoize('concat-2', function(lang, query)
   assert(language.add(lang))
   local ts_query = vim._ts_parse_query(lang, query)
   return Query.new(lang, ts_query)
-end)
+end, false)
 
 --- Implementations of predicates that can optionally be prefixed with "any-".
 ---
diff --git a/test/functional/treesitter/query_spec.lua b/test/functional/treesitter/query_spec.lua
index 6e21ed1d99..6bab171ee8 100644
--- a/test/functional/treesitter/query_spec.lua
+++ b/test/functional/treesitter/query_spec.lua
@@ -86,7 +86,7 @@ void ui_refresh(void)
         local before = vim.api.nvim__stats().ts_query_parse_count
         collectgarbage('stop')
         for _ = 1, _n, 1 do
-          vim.treesitter.query.parse('c', long_query, _n)
+          vim.treesitter.query.parse('c', long_query)
         end
         collectgarbage('restart')
         collectgarbage('collect')
@@ -96,8 +96,39 @@ void ui_refresh(void)
     end
 
     eq(1, q(1))
-    -- cache is cleared by garbage collection even if valid "cquery" reference is kept around
-    eq(1, q(100))
+    -- cache is retained even after garbage collection
+    eq(0, q(100))
+  end)
+
+  it('cache is cleared upon runtimepath changes, or setting query manually', function()
+    ---@return number
+    exec_lua(function()
+      _G.query_parse_count = _G.query_parse_count or 0
+      local parse = vim.treesitter.query.parse
+      vim.treesitter.query.parse = function(...)
+        _G.query_parse_count = _G.query_parse_count + 1
+        return parse(...)
+      end
+    end)
+
+    local function q(_n)
+      return exec_lua(function()
+        for _ = 1, _n, 1 do
+          vim.treesitter.query.get('c', 'highlights')
+        end
+        return _G.query_parse_count
+      end)
+    end
+
+    eq(1, q(10))
+    exec_lua(function()
+      vim.opt.rtp:prepend('/another/dir')
+    end)
+    eq(2, q(100))
+    exec_lua(function()
+      vim.treesitter.query.set('c', 'highlights', [[; test]])
+    end)
+    eq(3, q(100))
   end)
 
   it('supports query and iter by capture (iter_captures)', function()