2 files changed, 27 insertions, 31 deletions
diff --git a/lua/nvim-treesitter/locals.lua b/lua/nvim-treesitter/locals.lua
index b52d42b3f..93bcd90fd 100644
--- a/lua/nvim-treesitter/locals.lua
+++ b/lua/nvim-treesitter/locals.lua
@@ -23,6 +23,24 @@ setmetatable(query_cache, default_dict)
 local M = {}
 
 function M.collect_locals(bufnr, query_kind)
+  local locals = {}
+
+  for prepared_match in M.iter_locals(bufnr, nil, query_kind) do
+    table.insert(locals, prepared_match)
+  end
+
+  return locals
+end
+
+local function update_cached_locals(bufnr, changed_tick, query_kind)
+  query_cache[query_kind][bufnr] = {tick=changed_tick, cache=( M.collect_locals(bufnr, query_kind) or {} )}
+end
+
+-- Iterates matches from a locals query file.
+-- @param bufnr the buffer
+-- @param root the root node
+-- @param query_kind the query file to use
+function M.iter_locals(bufnr, root, query_kind)
   query_kind = query_kind or 'locals'
 
   local lang = parsers.ft_to_lang(api.nvim_buf_get_option(bufnr, "ft"))
@@ -34,20 +52,10 @@ function M.collect_locals(bufnr, query_kind)
   local parser = parsers.get_parser(bufnr, lang)
   if not parser then return end
 
-  local root = parser:parse():root()
+  local root = root or parser:parse():root()
   local start_row, _, end_row, _ = root:range()
 
-  local locals = {}
-
-  for prepared_match in queries.iter_prepared_matches(query, root, bufnr, start_row, end_row) do
-    table.insert(locals, prepared_match)
-  end
-
-  return locals
-end
-
-local function update_cached_locals(bufnr, changed_tick, query_kind)
-  query_cache[query_kind][bufnr] = {tick=changed_tick, cache=( M.collect_locals(bufnr, query_kind) or {} )}
+  return queries.iter_prepared_matches(query, root, bufnr, start_row, end_row)
 end
 
 function M.get_locals(bufnr, query_kind)
diff --git a/lua/nvim-treesitter/ts_utils.lua b/lua/nvim-treesitter/ts_utils.lua
index 7d5e97cdd..463490f5e 100644
--- a/lua/nvim-treesitter/ts_utils.lua
+++ b/lua/nvim-treesitter/ts_utils.lua
@@ -299,30 +299,18 @@ function M.find_usages(node, scope_node, bufnr)
   if not node_text or #node_text < 1 then return {} end
 
   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)
+  for match in locals.iter_locals(bufnr, scope_node) do
+    if match.reference
+      and match.reference.node
+      and M.get_node_text(match.reference.node)[1] == node_text
+    then
+      table.insert(usages, match.reference.node)
     end
-    next()
-  end)
+  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