Use weak tables in tree-sitter code (#17117)

feat(treesitter): use weak tables when possible

Also add the defaulttable function to create a table whose values are created when a key is missing.

5 files changed, 69 insertions, 10 deletions

diff --git a/runtime/doc/lua.txt b/runtime/doc/lua.txt
index a634cc1e93..35badb13b1 100644
--- a/runtime/doc/lua.txt
+++ b/runtime/doc/lua.txt
@@ -1581,6 +1581,28 @@ deepcopy({orig})                                              *vim.deepcopy()*
     Return: ~
         (table) Table of copied keys and (nested) values.
 
+defaulttable({create})                                    *vim.defaulttable()*
+    Creates a table whose members are automatically created when accessed, if
+    they don't already exist.
+
+    They mimic defaultdict in python.
+
+    If `create` is `nil`, this will create a defaulttable whose constructor
+    function is this function, effectively allowing to create nested tables on
+    the fly:
+>
+
+    local a = vim.defaulttable()
+    a.b.c = 1
+<
+
+    Parameters: ~
+        {create}  (function|nil) The function called to create a missing
+                  value.
+
+    Return: ~
+        (table) Empty table with metamethod
+
 endswith({s}, {suffix})                                       *vim.endswith()*
     Tests if `s` ends with `suffix`.
 
diff --git a/runtime/lua/vim/shared.lua b/runtime/lua/vim/shared.lua
index e1b4ed4ea9..59cb669609 100644
--- a/runtime/lua/vim/shared.lua
+++ b/runtime/lua/vim/shared.lua
@@ -715,5 +715,30 @@ function vim.is_callable(f)
   return type(m.__call) == 'function'
 end
 
+--- Creates a table whose members are automatically created when accessed, if they don't already
+--- exist.
+---
+--- They mimic defaultdict in python.
+---
+--- If @p create is @c nil, this will create a defaulttable whose constructor function is
+--- this function, effectively allowing to create nested tables on the fly:
+---
+--- <pre>
+--- local a = vim.defaulttable()
+--- a.b.c = 1
+--- </pre>
+---
+---@param create function|nil The function called to create a missing value.
+---@return table Empty table with metamethod
+function vim.defaulttable(create)
+  create = create or vim.defaulttable
+  return setmetatable({}, {
+    __index = function(tbl, key)
+      rawset(tbl, key, create())
+      return rawget(tbl, key)
+    end,
+  })
+end
+
 return vim
 -- vim:sw=2 ts=2 et
diff --git a/runtime/lua/vim/treesitter.lua b/runtime/lua/vim/treesitter.lua
index 5ebccff8f7..cde0491b12 100644
--- a/runtime/lua/vim/treesitter.lua
+++ b/runtime/lua/vim/treesitter.lua
@@ -3,10 +3,7 @@ local query = require('vim.treesitter.query')
 local language = require('vim.treesitter.language')
 local LanguageTree = require('vim.treesitter.languagetree')
 
--- TODO(bfredl): currently we retain parsers for the lifetime of the buffer.
--- Consider use weak references to release parser if all plugins are done with
--- it.
-local parsers = {}
+local parsers = setmetatable({}, { __mode = 'v' })
 
 local M = vim.tbl_extend('error', query, language)
 
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index 697e2e7691..78042a74bf 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -140,12 +140,9 @@ 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,
-})
+local query_cache = vim.defaulttable(function()
+  return setmetatable({}, { __mode = 'v' })
+end)
 
 --- Parse {query} as a string. (If the query is in a file, the caller
 --- should read the contents into a string before calling).
diff --git a/test/functional/lua/vim_spec.lua b/test/functional/lua/vim_spec.lua
index f2fb661b70..cd3240cd30 100644
--- a/test/functional/lua/vim_spec.lua
+++ b/test/functional/lua/vim_spec.lua
@@ -2754,6 +2754,24 @@ describe('lua stdlib', function()
 
   end)
 
+  describe("vim.defaulttable", function()
+    it("creates nested table by default", function()
+      eq({ b = {c = 1 } }, exec_lua[[
+        local a = vim.defaulttable()
+        a.b.c = 1
+        return a
+      ]])
+    end)
+
+    it("allows to create default objects", function()
+      eq({ b = 1 }, exec_lua[[
+        local a = vim.defaulttable(function() return 0 end)
+        a.b = a.b + 1
+        return a
+      ]])
+    end)
+  end)
+
 end)
 
 describe('lua: builtin modules', function()