refactor(refactor): use higher local apis and some cleanup

README.md

@@ -102,9 +102,23 @@ require'nvim-treesitter.configs'.setup {
           init_selection = 'gnn',         -- maps in normal mode to init the node/scope selection
           node_incremental = "grn",       -- increment to the upper named parent
           scope_incremental = "grc",      -- increment to the upper scope (as defined in locals.scm)
-          node_decremental = "grm",      -- decrement to the previous node
+          node_decremental = "grm",       -- decrement to the previous node
         }
     },
+    refactor = {
+      highlight_defintions = {
+        enable = true
+      },
+      smart_rename = {
+        enable = true,
+        smart_rename = "grr"              -- mapping to rename reference under cursor
+      },
+      navigation = {
+        enable = true,
+        goto_definition = "gnd",          -- mapping to go to definition of symbol under cursor
+        list_definitions = "gnD"          -- mapping to list all definitions in current file
+      }
+    },
     ensure_installed = 'all' -- one of 'all', 'language', or a list of languages
 }
 EOF
@@ -134,6 +148,16 @@ Some of these features are :
 You can find the roadmap [here](https://github.com/nvim-treesitter/nvim-treesitter/projects/1).
 The roadmap and all features of this plugin are open to change, and any suggestion will be highly appreciated!
 
+## Available Modules
+
+- `highlight`: Consistent syntax highlighting.
+- `incremental_selection`: Syntax based selection.
+- `refactor.highlight_definitions`: Syntax based definition and usage highlighting.
+- `refactor.smart_rename`: Syntax based definition and usage renaming.
+- `refactor.navigation`: Syntax based definition listing and navigation.
+  * List all definitions
+  * Go to definition
+
 ## Utils
 
 you can get some utility functions with

lua/nvim-treesitter/configs.lua

@@ -237,21 +237,19 @@ end
 function M.setup(user_data)
   if not user_data then return end
 
-  for mod, data in pairs(user_data) do
-    if config.modules[mod] then
-      M.setup_module(config.modules[mod], data)
-    elseif mod == 'ensure_installed' then
-      config.ensure_installed = data
-      require'nvim-treesitter.install'.ensure_installed(data)
-    end
-  end
+  M.setup_module(config.modules, user_data)
 end
 
---- Sets up a single module or all submodules of a group
+-- Sets up a single module or all submodules of a group.
+-- Note, this method is recursive.
 -- @param mod the module or group of modules
 -- @param data user defined configuration for the module
-function M.setup_module(mod, data)
-  if M.is_module(mod) then
+-- @param mod_name name of the module if it exists
+function M.setup_module(mod, data, mod_name)
+  if mod_name == 'ensure_installed' then
+    config.ensure_installed = data
+    require'nvim-treesitter.install'.ensure_installed(data)
+  elseif M.is_module(mod) then
     if type(data.enable) == 'boolean' then
       mod.enable = data.enable
     end
@@ -267,7 +265,7 @@ function M.setup_module(mod, data)
     end
   elseif type(data) == 'table' and type(mod) == 'table' then
     for key, value in pairs(data) do 
-      M.setup_module(mod[key], value)
+      M.setup_module(mod[key], value, key)
     end
   end
 end

lua/nvim-treesitter/locals.lua

@@ -11,7 +11,7 @@ local M = {
   locals = {}
 }
 
-function M.collect_locals(bufnr, root)
+function M.collect_locals(bufnr)
   local lang = parsers.ft_to_lang(api.nvim_buf_get_option(bufnr, "ft"))
   if not lang then return end
 
@@ -21,7 +21,7 @@ function M.collect_locals(bufnr, root)
   local parser = parsers.get_parser(bufnr, lang)
   if not parser then return end
 
-  local root = root or parser:parse():root()
+  local root = parser:parse():root()
   local start_row, _, end_row, _ = root:range()
 
   local locals = {}

lua/nvim-treesitter/refactor/highlight_definitions.lua

@@ -15,8 +15,11 @@ function M.highlight_usages(bufnr)
   M.clear_usage_highlights(bufnr)
 
   local node_at_point = ts_utils.get_node_at_cursor()
+  local references = locals.get_references(bufnr)
 
-  if not node_at_point then return end
+  if not node_at_point or not vim.tbl_contains(references, node_at_point) then 
+    return
+  end
 
   local def_node, scope = ts_utils.find_definition(node_at_point, bufnr)
   local usages = ts_utils.find_usages(node_at_point, scope)
@@ -28,25 +31,23 @@ function M.highlight_usages(bufnr)
       api.nvim_buf_add_highlight(
         bufnr, 
         usage_namespace, 
-        'Visual', 
+        'TSDefinitionUsage', 
         start_row,
         start_col,
         end_col)
     end
   end
 
-  if def_node then
+  if def_node ~= node_at_point then
     local start_row, start_col, _, end_col = def_node:range()
 
-    if def_node ~= node_at_point then
-      api.nvim_buf_add_highlight(
-        bufnr, 
-        usage_namespace, 
-        'Search', 
-        start_row,
-        start_col,
-        end_col)
-    end
+    api.nvim_buf_add_highlight(
+      bufnr, 
+      usage_namespace, 
+      'TSDefinition', 
+      start_row,
+      start_col,
+      end_col)
   end
 end
 

lua/nvim-treesitter/refactor/navigation.lua

@@ -26,12 +26,6 @@ function M.goto_definition(bufnr)
   if not node_at_point then return end
 
   local definition, _ = ts_utils.find_definition(node_at_point, bufnr)
-
-  if not definition then
-    print('No definition found')
-    return
-  end
-
   local start_row, start_col, _ = definition:start()
 
   api.nvim_win_set_cursor(0, { start_row + 1, start_col })

lua/nvim-treesitter/refactor/smart_rename.lua

@@ -3,6 +3,7 @@
 
 local ts_utils = require'nvim-treesitter.ts_utils'
 local configs = require'nvim-treesitter.configs'
+local utils = require'nvim-treesitter.utils'
 local api = vim.api
 
 local M = {}
@@ -12,7 +13,7 @@ function M.smart_rename(bufnr)
   local node_at_point = ts_utils.get_node_at_cursor()
 
   if not node_at_point then
-    print('No node to rename!')
+    utils.print_warning("No node to rename!")
     return
   end
 
@@ -29,23 +30,16 @@ function M.smart_rename(bufnr)
     table.insert(nodes_to_rename, node_at_point)
   end
 
-  if definition and not vim.tbl_contains(nodes_to_rename, definition) then
+  if not vim.tbl_contains(nodes_to_rename, definition) then
     table.insert(nodes_to_rename, definition)
   end
 
-  if #nodes_to_rename < 1 then
-    print('No nodes to rename!')
-    return
-  end
-
   for _, node in ipairs(nodes_to_rename) do
     local start_row, start_col, end_row, end_col = node:range() 
-
     local line = api.nvim_buf_get_lines(bufnr, start_row, start_row + 1, false)[1]
 
     if line then
       local new_line = line:sub(1, start_col) .. new_name .. line:sub(end_col + 1, -1)
-
       api.nvim_buf_set_lines(bufnr, start_row, start_row + 1, false, { new_line })
     end
   end
@@ -53,12 +47,10 @@ 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
     local cmd = string.format([[:lua require'nvim-treesitter.refactor.smart_rename'.%s(%d)<CR>]], fn_name, bufnr)
-
     api.nvim_buf_set_keymap(bufnr, 'n', mapping, cmd, { silent = true })
   end
 end

lua/nvim-treesitter/ts_utils.lua

@@ -223,29 +223,33 @@ function M.find_definition(node, bufnr)
   local node_text = M.get_node_text(node)[1]
   local current_scope = M.containing_scope(node)
   local _, _, node_start = node:start()
+  local matching_def_nodes = {}
 
   -- If a scope wasn't found then use the root node
   if current_scope == node then
     current_scope = parsers.get_parser(bufnr).tree:root()
   end
-  
-  while current_scope ~= nil and current_scope ~= node do
-    for _, def in ipairs(locals.collect_locals(bufnr, current_scope)) do
-      if def.definition then
-        for _, def_node in ipairs(M.get_local_nodes(def.definition)) do
-          local _, _, def_start = def_node:start()
 
-          if M.get_node_text(def_node)[1] == node_text and def_start < node_start then
-            return def_node, current_scope
-          end
-        end
+  -- Get all definitions that match the node text
+  for _, def in ipairs(locals.get_definitions(bufnr)) do
+    for _, def_node in ipairs(M.get_local_nodes(def)) do
+      if M.get_node_text(def_node)[1] == node_text then
+        table.insert(matching_def_nodes, def_node)       
       end
     end
+  end
 
+  -- Continue up each scope until we find the scope that contains the definition
+  while current_scope do
+    for _, def_node in ipairs(matching_def_nodes) do
+      if M.is_parent(current_scope, def_node) then
+        return def_node, current_scope
+      end
+    end
     current_scope = M.containing_scope(current_scope:parent())
   end
 
-  return nil, nil
+  return node, parsers.get_parser(bufnr).tree:root()
 end
 
 -- Gets all nodes from a local list result.
@@ -285,26 +289,41 @@ function M.recurse_local_nodes(local_def, accumulator, full_match, last_match)
   end
 end
 
--- Finds usages of a node in a particula scope
+-- Finds usages of a node in a given scope
 -- @param node the node to find usages for
 -- @param scope_node the node to look within
 -- @returns a list of nodes
-function M.find_usages(node, scope_node)
-  local usages = {}
+function M.find_usages(node, scope_node, bufnr)
+  local bufnr = bufnr or api.nvim_get_current_buf()
   local node_text = M.get_node_text(node)[1]
 
   if not node_text or #node_text < 1 then return {} end
 
-  for _, def in ipairs(locals.collect_locals(bufnr, scope_node)) do
-    if def.reference 
-      and def.reference.node 
-      and M.get_node_text(def.reference.node)[1] == node_text then
-      
-      table.insert(usages, def.reference.node)
+  local scope_node = scope_node or parsers.get_parser(bufnr).tree:root()
+  local references = locals.get_references(bufnr)
+  local usages = {}
+
+  M.recurse_tree(scope_node, function(iter_node, _, next)
+    if vim.tbl_contains(references, iter_node) and M.get_node_text(iter_node)[1] == node_text then
+      table.insert(usages, iter_node)
     end
-  end
+    next()
+  end)
 
   return usages
 end
 
+-- Recurses all child nodes of a tree.
+-- The callback is provided the child node, parent_node, and a callback to recurse into
+-- the child node. This allows for the ability to short circuit the recursion
+-- if we found what we are looking for, we can then stop the recursion or skip a node
+-- if need be.
+-- @param tree the node root
+-- @param cb the callback for each node
+function M.recurse_tree(tree, cb)
+  for _, child in ipairs(M.get_named_children(tree)) do
+    cb(child, tree, function(next_node) M.recurse_tree(next_node or child, cb) end)
+  end
+end
+
 return M

lua/nvim-treesitter/utils.lua

@@ -45,7 +45,7 @@ function M.get_cache_dir()
   return nil, 'Invalid cache rights, $XDG_CACHE_HOME or /tmp should be read/write'
 end
 
---- Gets a property at path
+-- Gets a property at path
 -- @param tbl the table to access
 -- @param path the '.' seperated path
 -- @returns the value at path or nil
@@ -62,4 +62,10 @@ function M.get_at_path(tbl, path)
   return result
 end
 
+-- Prints a warning message
+-- @param text the text message
+function M.print_warning(text)
+  api.nvim_command(string.format([[echohl WarningMsg | echo "%s" | echohl None]], text))
+end
+
 return M

plugin/nvim-treesitter.vim

@@ -57,3 +57,6 @@ highlight default link TSTypeBuiltin Type
 highlight default link TSStructure Structure
 highlight default link TSInclude Include
 
+highlight default link TSDefinitionUsage Visual
+highlight default link TSDefinition Search
+

queries/javascript/locals.scm

@@ -28,9 +28,6 @@
 (variable_declarator
   name: (identifier) @definition)
 
-(import_specifier
-  (identifier) @definition)
-
 ; References
 ;------------