feat(highlights): add is predicate

3 files changed, 37 insertions, 5 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/highlight.lua b/lua/nvim-treesitter/highlight.lua
index 965eccbe3..1fbc6621d 100644
--- a/lua/nvim-treesitter/highlight.lua
+++ b/lua/nvim-treesitter/highlight.lua
@@ -35,6 +35,7 @@ hlmap["function"] = "TSFunction"
 hlmap["function.builtin"] = "TSFuncBuiltin"
 hlmap["function.macro"] = "TSFuncMacro"
 hlmap["parameter"] = "TSParameter"
+hlmap["parameter.reference"] = "TSParameterReference"
 hlmap["method"] = "TSMethod"
 hlmap["field"] = "TSField"
 hlmap["property"] = "TSProperty"
diff --git a/lua/nvim-treesitter/query_predicates.lua b/lua/nvim-treesitter/query_predicates.lua
index 67850deaf..f6a21f28c 100644
--- a/lua/nvim-treesitter/query_predicates.lua
+++ b/lua/nvim-treesitter/query_predicates.lua
@@ -4,12 +4,25 @@ local function error(str)
   vim.api.nvim_err_writeln(str)
 end
 
-query.add_predicate("nth?", function(match, pattern, bufnr, pred)
-  if #pred ~= 3 then
-    error("nth? must hav exactly two arguments")
-    return
+local function valid_args(name, pred, count, strict_count)
+  local arg_count = #pred - 1
+
+  if strict_count then
+    if arg_count ~= count then
+      error(string.format("%s must have exactly %d arguments", name, count))
+      return false
+    end
+  elseif arg_count < count then
+    error(string.format("%s must have at least %d arguments", name, count))
+    return false
   end
 
+  return true
+end
+
+query.add_predicate("nth?", function(match, pattern, bufnr, pred)
+  if not valid_args("nth?", pred, 2, true) then return end
+
   local node = match[pred[2]]
   local n = pred[3] - 1
   if node and node:parent() and node:named_child_count() > n then
@@ -20,7 +33,7 @@ query.add_predicate("nth?", function(match, pattern, bufnr, pred)
 end)
 
 query.add_predicate('has-ancestor?', function(match, pattern, bufnr, pred)
-  if #pred ~= 3 then error("has-ancestor? must have exactly two arguments!") return end
+  if not valid_args("has-ancestor?", pred, 2, true) then return end
 
   local node = match[pred[2]]
   local ancestor_type = pred[3]
@@ -35,3 +48,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 not valid_args("is?", pred, 2) then 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)