fix(definitions): optimize and fix definition highlighting

1 files changed, 76 insertions, 41 deletions

diff --git a/lua/nvim-treesitter/locals.lua b/lua/nvim-treesitter/locals.lua
index f43a8e24..18d355c3 100644
--- a/lua/nvim-treesitter/locals.lua
+++ b/lua/nvim-treesitter/locals.lua
@@ -24,6 +24,16 @@ function M.get_locals(bufnr)
   return queries.get_matches(bufnr, 'locals')
 end
 
+--- Creates unique id for a node based on text and range
+-- @param scope: the scope node of the definition
+-- @param bufnr: the buffer
+-- @param node_text: the node text to use
+-- @returns a string id
+function M.get_definition_id(scope, node_text)
+  -- Add a vaild starting character in case node text doesn't start with a valid one.
+  return table.concat({ 'k', node_text, scope:range() }, '_')
+end
+
 function M.get_definitions(bufnr)
   local locals = M.get_locals(bufnr)
 
@@ -69,56 +79,30 @@ end
 --- Gets a table with all the scopes containing a node
 -- The order is from most specific to least (bottom up)
 function M.get_scope_tree(node, bufnr)
-  local current_scope = M.containing_scope(node, bufnr, false) or parsers.get_tree_root(bufnr)
   local scopes = {}
 
-  while current_scope do
-    table.insert(scopes, current_scope)
-    current_scope = current_scope:parent()
-      and (M.containing_scope(current_scope:parent(), bufnr, false) or parsers.get_tree_root(bufnr))
-      or nil
+  for scope in M.iter_scope_tree(node, bufnr) do
+    table.insert(scopes, scope)
   end
 
   return scopes
 end
 
--- Finds the definition node and it's scope node of a node
--- @param node starting node
--- @param bufnr buffer
--- @returns the definition node and the definition nodes scope node
-function M.find_definition(node, bufnr)
-  local bufnr = bufnr or api.nvim_get_current_buf()
-  local node_text = ts_utils.get_node_text(node, bufnr)[1]
-  local scope_tree = M.get_scope_tree(node, bufnr)
-  local match
-  local last_scope_index
+--- Iterates over a nodes scopes moving from the bottom up
+function M.iter_scope_tree(node, bufnr)
+  local last_node = node
 
-  -- Loop over every definition
-  for _, definition in ipairs(M.get_definitions(bufnr)) do
-    for _, node_entry in ipairs(M.get_local_nodes(definition)) do
-      local def_scope = M.containing_scope(node_entry.node, bufnr, false) or parsers.get_tree_root(bufnr)
-
-      -- Only match definitions that match the text of the node
-      -- Look for the most specific definition in the tree
-      -- The lower the index, the more specific the definition is
-      if ts_utils.get_node_text(node_entry.node, bufnr)[1] == node_text then
-	for i, scope_node in ipairs(scope_tree) do
-	  -- If we already found a close definition in scope, just skip checking
-	  if last_scope_index and i >= last_scope_index then break end
-	  if scope_node == def_scope then
-	    last_scope_index = i
-	    match = node_entry
-	  end
-	end
-      end
+  return function()
+    if not last_node then
+      return
     end
-  end
 
-  if match and last_scope_index then
-    return match.node, scope_tree[last_scope_index], match.kind
-  end
+    local scope = M.containing_scope(last_node, bufnr, false) or parsers.get_tree_root(bufnr)
+
+    last_node = scope:parent()
 
-  return node, parsers.get_parser(bufnr).tree:root(), nil
+    return scope
+  end
 end
 
 -- Gets a table of all nodes and their 'kinds' from a locals list
@@ -158,7 +142,54 @@ function M.recurse_local_nodes(local_def, accumulator, full_match, last_match)
   end
 end
 
--- Finds usages of a node in a given scope
+--- Get a single dimension table to look definition nodes.
+-- Keys are generated by using the range of the containing scope and the text of the definition node.
+-- This makes looking up a definition for a given scope a simple key lookup.
+--
+-- This is memoized by buffer tick. If the function is called in succession
+-- without the buffer tick changing, then the previous result will be used
+-- since the syntax tree hasn't changed.
+--
+-- Usage lookups require finding the definition of the node, so `find_definition`
+-- is called very frequently, which is why this lookup must be fast as possible.
+--
+-- @param bufnr: the buffer
+-- @returns a table for looking up definitions
+M.get_definitions_lookup_table = ts_utils.memoize_by_buf_tick(function(bufnr)
+  local definitions = M.get_definitions(bufnr)
+  local result = {}
+
+  for _, definition in ipairs(definitions) do
+    for _, node_entry in ipairs(M.get_local_nodes(definition)) do
+      local scope = M.containing_scope(node_entry.node, bufnr, false) or parsers.get_tree_root(bufnr)
+      local node_text = ts_utils.get_node_text(node_entry.node, bufnr)[1]
+      local id = M.get_definition_id(scope, node_text)
+
+      result[id] = node_entry
+    end
+  end
+
+  return result
+end)
+
+function M.find_definition(node, bufnr)
+  local def_lookup = M.get_definitions_lookup_table(bufnr)
+  local node_text = ts_utils.get_node_text(node, bufnr)[1]
+
+  for scope in M.iter_scope_tree(node, bufnr) do
+    local id = M.get_definition_id(scope, node_text)
+
+    if def_lookup[id] then
+      local entry = def_lookup[id]
+
+      return entry.node, scope, entry.kind
+    end
+  end
+
+  return node, parsers.get_tree_root(bufnr), nil
+end
+
+-- 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
@@ -176,7 +207,11 @@ function M.find_usages(node, scope_node, bufnr)
       and match.reference.node
       and ts_utils.get_node_text(match.reference.node, bufnr)[1] == node_text
     then
-      table.insert(usages, match.reference.node)
+      local def_node, _, kind = M.find_definition(match.reference.node, bufnr)
+
+      if kind == nil or def_node == node then
+        table.insert(usages, match.reference.node)
+      end
     end
   end