Merge pull request #330 from steelsojka/fix-do-not-reattach

fix(modules): do not reattach if already attached

10 files changed, 76 insertions, 42 deletions

diff --git a/lua/nvim-treesitter/caching.lua b/lua/nvim-treesitter/caching.lua
new file mode 100644
index 000000000..bfda39622
--- /dev/null
+++ b/lua/nvim-treesitter/caching.lua
@@ -0,0 +1,48 @@
+local api = vim.api
+
+local M = {}
+
+--- Creates a cache table for buffers keyed by a type name.
+--- Cache entries attach to the buffer and cleanup entries
+--- as buffers are detached.
+function M.create_buffer_cache()
+  local cache = {}
+
+  local items = setmetatable({}, {
+    __index = function(tbl, key)
+      rawset(tbl, key, {})
+      return rawget(tbl, key)
+    end
+  })
+
+  function cache.set(type_name, bufnr, value)
+    if not cache.has(type_name, bufnr) then
+      -- Clean up the cache if the buffer is detached
+      -- to avoid memory leaks
+      api.nvim_buf_attach(bufnr, false, {
+        on_detach = function()
+          cache.remove(type_name, bufnr)
+          return true
+        end
+      })
+    end
+
+    items[type_name][bufnr] = value
+  end
+
+  function cache.get(type_name, bufnr)
+    return items[type_name][bufnr]
+  end
+
+  function cache.has(type_name, bufnr)
+    return cache.get(type_name, bufnr) ~= nil
+  end
+
+  function cache.remove(type_name, bufnr)
+    items[type_name][bufnr] = nil
+  end
+
+  return cache
+end
+
+return M
diff --git a/lua/nvim-treesitter/configs.lua b/lua/nvim-treesitter/configs.lua
index ba4ae3f8d..51d982c79 100644
--- a/lua/nvim-treesitter/configs.lua
+++ b/lua/nvim-treesitter/configs.lua
@@ -3,6 +3,7 @@ local api = vim.api
 local queries = require'nvim-treesitter.query'
 local parsers = require'nvim-treesitter.parsers'
 local utils = require'nvim-treesitter.utils'
+local caching = require'nvim-treesitter.caching'
 
 local M = {}
 
@@ -109,6 +110,8 @@ local builtin_modules = {
   }
 }
 
+local attached_buffers_by_module = caching.create_buffer_cache()
+
 -- Resolves a module by requiring the `module_path` or using the module definition.
 local function resolve_module(mod_name)
   local config_mod = M.get_module(mod_name)
@@ -367,8 +370,13 @@ end
 -- @param lang the language of the buffer
 function M.attach_module(mod_name, bufnr, lang)
   local resolved_mod = resolve_module(mod_name)
+  local bufnr = bufnr or api.nvim_get_current_buf()
 
-  if resolved_mod and parsers.has_parser(lang) then
+  if resolved_mod
+    and parsers.has_parser(lang)
+    and not attached_buffers_by_module.has(mod_name, bufnr)
+  then
+    attached_buffers_by_module.set(mod_name, bufnr, true)
     resolved_mod.attach(bufnr, lang)
   end
 end
@@ -378,8 +386,10 @@ end
 -- @param bufnr the bufnr
 function M.detach_module(mod_name, bufnr)
   local resolved_mod = resolve_module(mod_name)
+  local bufnr = bufnr or api.nvim_get_current_buf()
 
   if resolved_mod then
+    attached_buffers_by_module.remove(mod_name, bufnr)
     resolved_mod.detach(bufnr)
   end
 end
diff --git a/lua/nvim-treesitter/highlight.lua b/lua/nvim-treesitter/highlight.lua
index 1fbc6621d..a9549fc06 100644
--- a/lua/nvim-treesitter/highlight.lua
+++ b/lua/nvim-treesitter/highlight.lua
@@ -69,7 +69,6 @@ hlmap["text.literal"] = "TSLiteral"
 hlmap["text.uri"] = "TSURI"
 
 function M.attach(bufnr, lang)
-  local bufnr = bufnr or api.nvim_get_current_buf()
   local lang = lang or parsers.get_buf_lang(bufnr)
   local config = configs.get_module('highlight')
 
@@ -84,12 +83,11 @@ function M.attach(bufnr, lang)
 end
 
 function M.detach(bufnr)
-  local buf = bufnr or api.nvim_get_current_buf()
-  if M.highlighters[buf] then
-    M.highlighters[buf]:set_query("")
-    M.highlighters[buf] = nil
+  if M.highlighters[bufnr] then
+    M.highlighters[bufnr]:set_query("")
+    M.highlighters[bufnr] = nil
   end
-  api.nvim_buf_set_option(buf, 'syntax', 'on')
+  api.nvim_buf_set_option(bufnr, 'syntax', 'on')
 end
 
 return M
diff --git a/lua/nvim-treesitter/incremental_selection.lua b/lua/nvim-treesitter/incremental_selection.lua
index dce5f112a..4ce265e26 100644
--- a/lua/nvim-treesitter/incremental_selection.lua
+++ b/lua/nvim-treesitter/incremental_selection.lua
@@ -94,8 +94,6 @@ function M.node_decremental()
 end
 
 function M.attach(bufnr)
-  local buf = bufnr or api.nvim_get_current_buf()
-
   local config = configs.get_module('incremental_selection')
   for funcname, mapping in pairs(config.keymaps) do
     local mode
@@ -105,19 +103,17 @@ function M.attach(bufnr)
       mode = 'v'
     end
     local cmd = string.format(":lua require'nvim-treesitter.incremental_selection'.%s()<CR>", funcname)
-    api.nvim_buf_set_keymap(buf, mode, mapping, cmd, { silent = true, noremap = true })
+    api.nvim_buf_set_keymap(bufnr, mode, mapping, cmd, { silent = true, noremap = true })
   end
 end
 
 function M.detach(bufnr)
-  local buf = bufnr or api.nvim_get_current_buf()
-
   local config = configs.get_module('incremental_selection')
   for f, mapping in pairs(config.keymaps) do
     if f == "init_selection" then
-      api.nvim_buf_del_keymap(buf, 'n', mapping)
+      api.nvim_buf_del_keymap(bufnr, 'n', mapping)
     else
-      api.nvim_buf_del_keymap(buf, 'v', mapping)
+      api.nvim_buf_del_keymap(bufnr, 'v', mapping)
     end
   end
 end
diff --git a/lua/nvim-treesitter/query.lua b/lua/nvim-treesitter/query.lua
index 4a2ce57d1..2dd85e2aa 100644
--- a/lua/nvim-treesitter/query.lua
+++ b/lua/nvim-treesitter/query.lua
@@ -2,20 +2,11 @@ local api = vim.api
 local ts = vim.treesitter
 local utils = require'nvim-treesitter.utils'
 local parsers = require'nvim-treesitter.parsers'
+local caching = require'nvim-treesitter.caching'
 
 local M = {}
 
-local default_dict = {
-  __index = function(table, key)
-    local exists = rawget(table, key)
-    if not exists then
-      table[key] = {}
-    end
-    return rawget(table, key)
-  end
-}
-
-local query_cache = setmetatable({}, default_dict)
+local query_cache = caching.create_buffer_cache()
 
 -- Some treesitter grammars extend others.
 -- We can use that to import the queries of the base language
@@ -55,17 +46,17 @@ local function read_query_files(filenames)
 end
 
 local function update_cached_matches(bufnr, changed_tick, query_group)
-  query_cache[query_group][bufnr] = {tick=changed_tick, cache=( M.collect_group_results(bufnr, query_group) or {} )}
+  query_cache.set(query_group, bufnr, {tick=changed_tick, cache=( M.collect_group_results(bufnr, query_group) or {} )})
 end
 
 function M.get_matches(bufnr, query_group)
   local bufnr = bufnr or api.nvim_get_current_buf()
-  local cached_local = query_cache[query_group][bufnr]
+  local cached_local = query_cache.get(query_group, bufnr)
   if not cached_local or api.nvim_buf_get_changedtick(bufnr) > cached_local.tick then
     update_cached_matches(bufnr,api.nvim_buf_get_changedtick(bufnr), query_group)
   end
 
-  return query_cache[query_group][bufnr].cache
+  return query_cache.get(query_group, bufnr).cache
 end
 
 local function filter_files(file_list)
diff --git a/lua/nvim-treesitter/refactor/highlight_current_scope.lua b/lua/nvim-treesitter/refactor/highlight_current_scope.lua
index 8b785ef8d..ef695f7b9 100644
--- a/lua/nvim-treesitter/refactor/highlight_current_scope.lua
+++ b/lua/nvim-treesitter/refactor/highlight_current_scope.lua
@@ -27,8 +27,6 @@ function M.clear_highlights(bufnr)
 end
 
 function M.attach(bufnr)
-  local bufnr = bufnr or api.nvim_get_current_buf()
-
   cmd(string.format('augroup NvimTreesitterCurrentScope_%d', bufnr))
   cmd 'au!'
   -- luacheck: push ignore 631
diff --git a/lua/nvim-treesitter/refactor/highlight_definitions.lua b/lua/nvim-treesitter/refactor/highlight_definitions.lua
index 8a75980d6..d430d38a6 100644
--- a/lua/nvim-treesitter/refactor/highlight_definitions.lua
+++ b/lua/nvim-treesitter/refactor/highlight_definitions.lua
@@ -39,8 +39,6 @@ function M.clear_usage_highlights(bufnr)
 end
 
 function M.attach(bufnr)
-  local bufnr = bufnr or api.nvim_get_current_buf()
-
   cmd(string.format('augroup NvimTreesitterUsages_%d', bufnr))
   cmd 'au!'
   -- luacheck: push ignore 631
diff --git a/lua/nvim-treesitter/refactor/navigation.lua b/lua/nvim-treesitter/refactor/navigation.lua
index 5f5eebd53..185cce89f 100644
--- a/lua/nvim-treesitter/refactor/navigation.lua
+++ b/lua/nvim-treesitter/refactor/navigation.lua
@@ -49,8 +49,6 @@ function M.list_definitions(bufnr)
 end
 
 function M.attach(bufnr)
-  local bufnr = bufnr or api.nvim_get_current_buf()
-
   local config = configs.get_module('refactor.navigation')
 
   for fn_name, mapping in pairs(config.keymaps) do
diff --git a/lua/nvim-treesitter/refactor/smart_rename.lua b/lua/nvim-treesitter/refactor/smart_rename.lua
index 877da44e1..711ab993a 100644
--- a/lua/nvim-treesitter/refactor/smart_rename.lua
+++ b/lua/nvim-treesitter/refactor/smart_rename.lua
@@ -46,7 +46,6 @@ function M.smart_rename(bufnr)
 end
 
 function M.attach(bufnr)
-  local bufnr = bufnr or api.nvim_get_current_buf()
   local config = configs.get_module('refactor.smart_rename')
 
   for fn_name, mapping in pairs(config.keymaps) do
diff --git a/lua/nvim-treesitter/textobjects/attach.lua b/lua/nvim-treesitter/textobjects/attach.lua
index 22d6b650d..c7f443759 100644
--- a/lua/nvim-treesitter/textobjects/attach.lua
+++ b/lua/nvim-treesitter/textobjects/attach.lua
@@ -6,9 +6,8 @@ local M = {}
 
 function M.make_attach(normal_mode_functions, submodule)
   return function(bufnr, lang)
-    local buf = bufnr or api.nvim_get_current_buf()
     local config = configs.get_module("textobjects."..submodule)
-    local lang = lang or parsers.get_buf_lang(buf)
+    local lang = lang or parsers.get_buf_lang(bufnr)
 
     for _, function_call in pairs(normal_mode_functions) do
       for mapping, query in pairs(config[function_call] or {}) do
@@ -19,7 +18,7 @@ function M.make_attach(normal_mode_functions, submodule)
         end
         if query then
           local cmd = ":lua require'nvim-treesitter.textobjects."..submodule.."'."..function_call.."('"..query.."')<CR>"
-          api.nvim_buf_set_keymap(buf, "n", mapping, cmd, {silent = true, noremap = true })
+          api.nvim_buf_set_keymap(bufnr, "n", mapping, cmd, {silent = true, noremap = true })
         end
       end
     end
@@ -28,7 +27,6 @@ end
 
 function M.make_detach(normal_mode_functions, submodule)
   return function(bufnr)
-    local buf = bufnr or api.nvim_get_current_buf()
     local config = configs.get_module("textobjects."..submodule)
     local lang = parsers.get_buf_lang(bufnr)
 
@@ -39,8 +37,8 @@ function M.make_detach(normal_mode_functions, submodule)
         query = nil
       end
       if query then
-        api.nvim_buf_del_keymap(buf, "o", mapping)
-        api.nvim_buf_del_keymap(buf, "v", mapping)
+        api.nvim_buf_del_keymap(bufnr, "o", mapping)
+        api.nvim_buf_del_keymap(bufnr, "v", mapping)
       end
     end
     for _, function_call in pairs(normal_mode_functions) do
@@ -51,7 +49,7 @@ function M.make_detach(normal_mode_functions, submodule)
           query = nil
         end
         if query then
-          api.nvim_buf_del_keymap(buf, "n", mapping)
+          api.nvim_buf_del_keymap(bufnr, "n", mapping)
         end
       end
     end