perf(treesitter): cache query parsing

2 files changed, 41 insertions, 7 deletions

diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index ebed502c92..b3036ea679 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -138,6 +138,13 @@ function M.get_query(lang, query_name)
   end
 end
 
+local query_cache = setmetatable({}, {
+  __index = function(tbl, key)
+    rawset(tbl, key, {})
+    return rawget(tbl, key)
+  end
+})
+
 --- Parse {query} as a string. (If the query is in a file, the caller
 ---       should read the contents into a string before calling).
 ---
@@ -151,17 +158,23 @@ end
 ---   -` info.captures` also points to `captures`.
 ---   - `info.patterns` contains information about predicates.
 ---
----@param lang The language
----@param query A string containing the query (s-expr syntax)
+---@param lang string The language
+---@param query string A string containing the query (s-expr syntax)
 ---
 ---@returns The query
 function M.parse_query(lang, query)
   language.require_language(lang)
-  local self = setmetatable({}, Query)
-  self.query = vim._ts_parse_query(lang, query)
-  self.info = self.query:inspect()
-  self.captures = self.info.captures
-  return self
+  local cached = query_cache[lang][query]
+  if cached then
+    return cached
+  else
+    local self = setmetatable({}, Query)
+    self.query = vim._ts_parse_query(lang, query)
+    self.info = self.query:inspect()
+    self.captures = self.info.captures
+    query_cache[lang][query] = self
+    return self
+  end
 end
 
 --- Gets the text corresponding to a given node
diff --git a/test/functional/treesitter/parser_spec.lua b/test/functional/treesitter/parser_spec.lua
index 911fa017ab..44a7f00d8b 100644
--- a/test/functional/treesitter/parser_spec.lua
+++ b/test/functional/treesitter/parser_spec.lua
@@ -167,6 +167,27 @@ void ui_refresh(void)
     eq('variable', ret)
   end)
 
+  it("supports caching queries", function()
+    local long_query = query:rep(100)
+    local first_run = exec_lua ([[
+      local before = vim.loop.hrtime()
+      cquery = vim.treesitter.parse_query("c", ...)
+      local after = vim.loop.hrtime()
+      return after - before
+    ]], long_query)
+
+    local subsequent_runs = exec_lua ([[
+      local before = vim.loop.hrtime()
+      for i=1,100,1 do
+        cquery = vim.treesitter.parse_query("c", ...)
+      end
+      local after = vim.loop.hrtime()
+      return after - before
+    ]], long_query)
+
+    assert.True(1000 * subsequent_runs < first_run)
+  end)
+
   it('support query and iter by capture', function()
     insert(test_text)