Merge pull request #295 from steelsojka/feature/set-scopes-for-definitions

feat(defintions): allow setting of scope

3 files changed, 70 insertions, 3 deletions

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 62a38b71c..0e72dcb66 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -177,3 +177,30 @@ Mainly for markup languages.
 @scope
 @reference
 ```
+
+#### Definition Scope
+
+You can set the scope of a definition by setting the `scope` property on the definition.
+
+For example, a javascript function declaration creates a scope. The function name is captured as the definition.
+This means that the function definition would only be available WITHIN the scope of the function, which is not the case.
+The definition can be used in the scope the function was defined in.
+
+```javascript
+function doSomething() {}
+
+doSomething(); // Should point to the declaration as the definition
+```
+
+```scheme
+(function_declaration
+  ((identifier) @definition.var)
+   (set! "definition.var.scope" "parent"))
+```
+
+Possible scope values are:
+
+- `parent`: The definition is valid in the containing scope and one more scope above that scope
+- `global`: The definition is valid in the root scope
+- `local`: The definition is valid in the containing scope. This is the default behavior
+
diff --git a/lua/nvim-treesitter/locals.lua b/lua/nvim-treesitter/locals.lua
index 18d355c35..14ec47abb 100644
--- a/lua/nvim-treesitter/locals.lua
+++ b/lua/nvim-treesitter/locals.lua
@@ -111,8 +111,8 @@ end
 function M.get_local_nodes(local_def)
   local result = {}
 
-  M.recurse_local_nodes(local_def, function(_, node, kind)
-    table.insert(result, { node = node, kind = kind })
+  M.recurse_local_nodes(local_def, function(def, node, kind)
+    table.insert(result, vim.tbl_extend("keep", { kind = kind }, def))
   end)
 
   return result
@@ -161,7 +161,9 @@ M.get_definitions_lookup_table = ts_utils.memoize_by_buf_tick(function(bufnr)
 
   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 scopes = M.get_definition_scopes(node_entry.node, bufnr, node_entry.scope)
+      -- Always use the highest valid scope
+      local scope = scopes[#scopes]
       local node_text = ts_utils.get_node_text(node_entry.node, bufnr)[1]
       local id = M.get_definition_id(scope, node_text)
 
@@ -172,6 +174,40 @@ M.get_definitions_lookup_table = ts_utils.memoize_by_buf_tick(function(bufnr)
   return result
 end)
 
+--- Gets all the scopes of a definition based on the scope type
+-- Scope types can be
+--
+-- "parent": Uses the parent of the containing scope, basically, skipping a scope
+-- "global": Uses the top most scope
+-- "local": Uses the containg scope of the definition. This is the default
+--
+-- @param node: the definition node
+-- @param bufnr: the buffer
+-- @param scope_type: the scope type
+function M.get_definition_scopes(node, bufnr, scope_type)
+  local scopes = {}
+  local scope_count = 1
+
+  -- Definition is valid for the containing scope
+  -- and the containing scope of that scope
+  if scope_type == 'parent' then
+    scope_count = 2
+  -- Definition is valid in all parent scopes
+  elseif scope_type == 'global' then
+    scope_count = nil
+  end
+
+  local i = 0
+  for scope in M.iter_scope_tree(node, bufnr) do
+    table.insert(scopes, scope)
+    i = i + 1
+
+    if scope_count and i >= scope_count then break end
+  end
+
+  return scopes
+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]
diff --git a/queries/javascript/locals.scm b/queries/javascript/locals.scm
index 74bc54050..ee4915929 100644
--- a/queries/javascript/locals.scm
+++ b/queries/javascript/locals.scm
@@ -43,6 +43,10 @@
 (import_specifier
   (identifier) @definition.import)
 
+(function_declaration
+  ((identifier) @definition.var)
+   (#set! definition.var.scope parent))
+
 ; References
 ;------------