feat(lua): add vim.func._memoize

Memoizes a function, using a custom function to hash the arguments.

Private for now until:

- There are other places in the codebase that could benefit from this
  (e.g. LSP), but might require other changes to accommodate.
- Invalidation of the cache needs to be controllable. Using weak tables
  is an acceptable invalidation policy, but it shouldn't be the only
  one.
- I don't think the story around `hash_fn` is completely thought out. We
  may be able to have a good default hash_fn by hashing each argument,
  so basically a better 'concat'.

runtime/lua/vim/treesitter/_query_linter.lua

@@ -10,20 +10,12 @@ local M = {}
 
 --- @alias vim.treesitter.ParseError {msg: string, range: Range4}
 
---- @private
---- Caches parse results for queries for each language.
---- Entries of parse_cache[lang][query_text] will either be true for successful parse or contain the
---- message and range of the parse error.
---- @type table<string,table<string,vim.treesitter.ParseError|true>>
-local parse_cache = {}
-
 --- Contains language dependent context for the query linter
 --- @class QueryLinterLanguageContext
 --- @field lang string? Current `lang` of the targeted parser
 --- @field parser_info table? Parser info returned by vim.treesitter.language.inspect
 --- @field is_first_lang boolean Whether this is the first language of a linter run checking queries for multiple `langs`
 
---- @private
 --- Adds a diagnostic for node in the query buffer
 --- @param diagnostics Diagnostic[]
 --- @param range Range4
@@ -42,7 +34,6 @@ local function add_lint_for_node(diagnostics, range, lint, lang)
   }
 end
 
---- @private
 --- Determines the target language of a query file by its path: <lang>/<query_type>.scm
 --- @param buf integer
 --- @return string?
@@ -53,7 +44,6 @@ local function guess_query_lang(buf)
   end
 end
 
---- @private
 --- @param buf integer
 --- @param opts QueryLinterOpts|QueryLinterNormalizedOpts|nil
 --- @return QueryLinterNormalizedOpts
@@ -87,7 +77,6 @@ local lint_query = [[;; query
   (ERROR) @error
 ]]
 
---- @private
 --- @param err string
 --- @param node TSNode
 --- @return vim.treesitter.ParseError
@@ -112,38 +101,26 @@ local function get_error_entry(err, node)
   }
 end
 
---- @private
 --- @param node TSNode
 --- @param buf integer
 --- @param lang string
---- @param diagnostics Diagnostic[]
-local function check_toplevel(node, buf, lang, diagnostics)
-  local query_text = vim.treesitter.get_node_text(node, buf)
-
-  if not parse_cache[lang] then
-    parse_cache[lang] = {}
-  end
-
-  local lang_cache = parse_cache[lang]
-
-  if lang_cache[query_text] == nil then
-    local cache_val, err = pcall(vim.treesitter.query.parse, lang, query_text) ---@type boolean|vim.treesitter.ParseError, string|Query
-
-    if not cache_val and type(err) == 'string' then
-      cache_val = get_error_entry(err, node)
-    end
-
-    lang_cache[query_text] = cache_val
-  end
-
-  local cache_entry = lang_cache[query_text]
-
-  if type(cache_entry) ~= 'boolean' then
-    add_lint_for_node(diagnostics, cache_entry.range, cache_entry.msg, lang)
-  end
+local function hash_parse(node, buf, lang)
+  return tostring(node:id()) .. tostring(buf) .. tostring(vim.b[buf].changedtick) .. lang
 end
 
---- @private
+--- @param node TSNode
+--- @param buf integer
+--- @param lang string
+--- @return vim.treesitter.ParseError?
+local parse = vim.func._memoize(hash_parse, function(node, buf, lang)
+  local query_text = vim.treesitter.get_node_text(node, buf)
+  local ok, err = pcall(vim.treesitter.query.parse, lang, query_text) ---@type boolean|vim.treesitter.ParseError, string|Query
+
+  if not ok and type(err) == 'string' then
+    return get_error_entry(err, node)
+  end
+end)
+
 --- @param buf integer
 --- @param match table<integer,TSNode>
 --- @param query Query
@@ -164,7 +141,10 @@ local function lint_match(buf, match, query, lang_context, diagnostics)
 
     -- other checks rely on Neovim parser introspection
     if lang and parser_info and cap_id == 'toplevel' then
-      check_toplevel(node, buf, lang, diagnostics)
+      local err = parse(node, buf, lang)
+      if err then
+        add_lint_for_node(diagnostics, err.range, err.msg, lang)
+      end
     end
   end
 end

runtime/lua/vim/treesitter/languagetree.lua

@@ -738,12 +738,14 @@ local function add_injection(t, tree_index, pattern, lang, combined, ranges)
 end
 
 -- TODO(clason): replace by refactored `ts.has_parser` API (without registering)
----@param lang string parser name
----@return boolean # true if parser for {lang} exists on rtp
-local has_parser = function(lang)
+--- The result of this function is cached to prevent nvim_get_runtime_file from being
+--- called too often
+--- @param lang string parser name
+--- @return boolean # true if parser for {lang} exists on rtp
+local has_parser = vim.func._memoize(1, function(lang)
   return vim._ts_has_language(lang)
     or #vim.api.nvim_get_runtime_file('parser/' .. lang .. '.*', false) > 0
-end
+end)
 
 --- Return parser name for language (if exists) or filetype (if registered and exists).
 --- Also attempts with the input lower-cased.

runtime/lua/vim/treesitter/query.lua

@@ -191,12 +191,6 @@ function M.set(lang, query_name, text)
   explicit_queries[lang][query_name] = M.parse(lang, text)
 end
 
---- `false` if query files didn't exist or were empty
----@type table<string, table<string, Query|false>>
-local query_get_cache = vim.defaulttable(function()
-  return setmetatable({}, { __mode = 'v' })
-end)
-
 ---@deprecated
 function M.get_query(...)
   vim.deprecate('vim.treesitter.query.get_query()', 'vim.treesitter.query.get()', '0.10')
@@ -209,34 +203,19 @@ end
 ---@param query_name string Name of the query (e.g. "highlights")
 ---
 ---@return Query|nil Parsed query
-function M.get(lang, query_name)
+M.get = vim.func._memoize('concat-2', function(lang, query_name)
   if explicit_queries[lang][query_name] then
     return explicit_queries[lang][query_name]
   end
 
-  local cached = query_get_cache[lang][query_name]
-  if cached then
-    return cached
-  elseif cached == false then
-    return nil
-  end
-
   local query_files = M.get_files(lang, query_name)
   local query_string = read_query_files(query_files)
 
   if #query_string == 0 then
-    query_get_cache[lang][query_name] = false
     return nil
   end
 
-  local query = M.parse(lang, query_string)
-  query_get_cache[lang][query_name] = query
-  return query
-end
-
----@type table<string, table<string, Query>>
-local query_parse_cache = vim.defaulttable(function()
-  return setmetatable({}, { __mode = 'v' })
+  return M.parse(lang, query_string)
 end)
 
 ---@deprecated
@@ -262,20 +241,15 @@ end
 ---@param query string Query in s-expr syntax
 ---
 ---@return Query Parsed query
-function M.parse(lang, query)
+M.parse = vim.func._memoize('concat-2', function(lang, query)
   language.add(lang)
-  local cached = query_parse_cache[lang][query]
-  if cached then
-    return cached
-  end
 
   local self = setmetatable({}, Query)
   self.query = vim._ts_parse_query(lang, query)
   self.info = self.query:inspect()
   self.captures = self.info.captures
-  query_parse_cache[lang][query] = self
   return self
-end
+end)
 
 ---@deprecated
 function M.get_range(...)