From ae2b88e45f2ca9958be37f0b0aae99ccbdaee5ea Mon Sep 17 00:00:00 2001
From: Steven Sojka <steelsojka@gmail.com>
Date: Sat, 15 Aug 2020 09:24:24 -0500
Subject: [PATCH] feat(highlights): add is predicate

---
 lua/nvim-treesitter.lua                  |  3 +++
 lua/nvim-treesitter/query_predicates.lua | 15 ++++++++++++
 queries/javascript/highlights.scm        | 31 +++++++++++++++---------
 queries/javascript/locals.scm            | 12 ++++-----
 4 files changed, 43 insertions(+), 18 deletions(-)

diff --git a/lua/nvim-treesitter.lua b/lua/nvim-treesitter.lua
index 0e7b85a9a..343a037bb 100644
--- a/lua/nvim-treesitter.lua
+++ b/lua/nvim-treesitter.lua
@@ -5,6 +5,9 @@ local info = require'nvim-treesitter.info'
 local configs = require'nvim-treesitter.configs'
 local parsers = require'nvim-treesitter.parsers'
 
+-- Registers all query predicates
+require"nvim-treesitter.query_predicates"
+
 local M = {}
 
 function M.setup()
diff --git a/lua/nvim-treesitter/query_predicates.lua b/lua/nvim-treesitter/query_predicates.lua
index 67850deaf..60ee391e6 100644
--- a/lua/nvim-treesitter/query_predicates.lua
+++ b/lua/nvim-treesitter/query_predicates.lua
@@ -35,3 +35,18 @@ query.add_predicate('has-ancestor?', function(match, pattern, bufnr, pred)
   end
   return false
 end)
+
+query.add_predicate('is?', function(match, pattern, bufnr, pred)
+  if #pred < 3 then error("is? must have at least two arguments!") return end
+
+  -- Avoid circular dependencies
+  local locals = require"nvim-treesitter.locals"
+  local node = match[pred[2]]
+  local types = {unpack(pred, 3)}
+
+  if not node then return true end
+
+  local _, _, kind = locals.find_definition(node, bufnr)
+
+  return vim.tbl_contains(types, kind)
+end)
diff --git a/queries/javascript/highlights.scm b/queries/javascript/highlights.scm
index fe373e617..4d1d6cb0e 100644
--- a/queries/javascript/highlights.scm
+++ b/queries/javascript/highlights.scm
@@ -1,23 +1,37 @@
 ; Types
 
 ; Javascript
+
+; Properties
+;-----------
+
+(property_identifier) @property
+
 ; Special identifiers
 ;--------------------
 
+(identifier) @variable
+
 ((identifier) @constant
- (#match? @constant "^[A-Z_][A-Z\\d_]+$"))
+ (match? @constant "^[A-Z_][A-Z\\d_]+$"))
 
 ((shorthand_property_identifier) @constant
- (#match? @constant "^[A-Z_][A-Z\\d_]+$"))
+ (match? @constant "^[A-Z_][A-Z\\d_]+$"))
 
 ((identifier) @constructor
- (#match? @constructor "^[A-Z]"))
+ (match? @constructor "^[A-Z]"))
 
 ((identifier) @variable.builtin
- (#match? @variable.builtin "^(arguments|module|console|window|document)$"))
+ (not-is? @variable.builtin import var parameter)
+ (match? @variable.builtin "^(arguments|module|console|window|document)$"))
 
 ((identifier) @function.builtin
- (#eq? @function.builtin "require"))
+ (not-is? @function.builtin import var parameter)
+ (eq? @function.builtin "require"))
+
+((identifier) @parameter
+ (is? @parameter parameter))
+
 
 ; Function and method definitions
 ;--------------------------------
@@ -78,13 +92,6 @@
   (rest_parameter
     (identifier) @parameter))
 
-(identifier) @variable
-
-; Properties
-;-----------
-
-(property_identifier) @property
-
 ; Literals
 ;---------
 
diff --git a/queries/javascript/locals.scm b/queries/javascript/locals.scm
index 74bc54050..b8d572b71 100644
--- a/queries/javascript/locals.scm
+++ b/queries/javascript/locals.scm
@@ -14,28 +14,28 @@
 ;------------
 
 (formal_parameters
-  (identifier) @definition.var)
+  (identifier) @definition.parameter)
 
 (formal_parameters
   (object_pattern
-    (identifier) @definition.var))
+    (identifier) @definition.parameter))
 
 ; function(arg = []) {
 (formal_parameters
   (assignment_pattern
-    (shorthand_property_identifier) @definition.var))
+    (shorthand_property_identifier) @definition.parameter))
 
 (formal_parameters
   (object_pattern
-    (shorthand_property_identifier) @definition.var))
+    (shorthand_property_identifier) @definition.parameter))
 
 (formal_parameters
   (array_pattern
-    (identifier) @definition.var))
+    (identifier) @definition.parameter))
 
 (formal_parameters
   (rest_parameter
-    (identifier) @definition.var))
+    (identifier) @definition.parameter))
 
 (variable_declarator
   name: (identifier) @definition.var)