Merge pull request #12739 from vigoux/ts-refactor-predicates

treesitter: refactor

6 files changed, 474 insertions, 214 deletions

diff --git a/runtime/doc/lua.txt b/runtime/doc/lua.txt
index 60c7a60d25..aa9addece8 100644
--- a/runtime/doc/lua.txt
+++ b/runtime/doc/lua.txt
@@ -574,6 +574,14 @@ retained for the lifetime of a buffer but this is subject to change. A plugin
 should keep a reference to the parser object as long as it wants incremental
 updates.
 
+Parser files						*treesitter-parsers*
+
+Parsers are the heart of tree-sitter. They are libraries that tree-sitter will
+search for in the `parsers` runtime directory.
+
+For a parser to be available for a given language, there must be a file named
+`{lang}.so` within the parser directory.
+
 Parser methods						*lua-treesitter-parser*
 
 tsparser:parse()					*tsparser:parse()*
@@ -593,9 +601,9 @@ shouldn't be done directly in the change callback anyway as they will be very
 frequent. Rather a plugin that does any kind of analysis on a tree should use
 a timer to throttle too frequent updates.
 
-tsparser:set_included_ranges(ranges)			*tsparser:set_included_ranges()*
+tsparser:set_included_ranges({ranges})			*tsparser:set_included_ranges()*
 	Changes the ranges the parser should consider. This is used for
-	language injection.  `ranges` should be of the form (all zero-based): >
+	language injection.  {ranges} should be of the form (all zero-based): >
 	{
 		{start_node, end_node},
 		...
@@ -617,15 +625,15 @@ tsnode:parent()						*tsnode:parent()*
 tsnode:child_count()					*tsnode:child_count()*
 	Get the node's number of children.
 
-tsnode:child(N)						*tsnode:child()*
-	Get the node's child at the given index, where zero represents the
+tsnode:child({index})						*tsnode:child()*
+	Get the node's child at the given {index}, where zero represents the
 	first child.
 
 tsnode:named_child_count()			*tsnode:named_child_count()*
 	Get the node's number of named children.
 
-tsnode:named_child(N)					*tsnode:named_child()*
-	Get the node's named child at the given index, where zero represents
+tsnode:named_child({index})					*tsnode:named_child()*
+	Get the node's named child at the given {index}, where zero represents
 	the first named child.
 
 tsnode:start()						*tsnode:start()*
@@ -661,12 +669,12 @@ tsnode:has_error()					*tsnode:has_error()*
 tsnode:sexpr()						*tsnode:sexpr()*
 	Get an S-expression representing the node as a string.
 
-tsnode:descendant_for_range(start_row, start_col, end_row, end_col)
+tsnode:descendant_for_range({start_row}, {start_col}, {end_row}, {end_col})
 						*tsnode:descendant_for_range()*
 	Get the smallest node within this node that spans the given range of
 	(row, column) positions
 
-tsnode:named_descendant_for_range(start_row, start_col, end_row, end_col)
+tsnode:named_descendant_for_range({start_row}, {start_col}, {end_row}, {end_col})
 					*tsnode:named_descendant_for_range()*
 	Get the smallest named node within this node that spans the given
 	range of (row, column) positions
@@ -677,17 +685,17 @@ Tree-sitter queries are supported, with some limitations. Currently, the only
 supported match predicate is `eq?` (both comparing a capture against a string
 and two captures against each other).
 
-vim.treesitter.parse_query(lang, query)
-						*vim.treesitter.parse_query(()*
-	Parse the query as a string. (If the query is in a file, the caller
+vim.treesitter.parse_query({lang}, {query})
+						*vim.treesitter.parse_query()*
+	Parse {query} as a string. (If the query is in a file, the caller
         should read the contents into a string before calling).
 
-query:iter_captures(node, bufnr, start_row, end_row)
+query:iter_captures({node}, {bufnr}, {start_row}, {end_row})
 							*query:iter_captures()*
-	Iterate over all captures from all matches inside a `node`.
-	`bufnr` is needed if the query contains predicates, then the caller
+	Iterate over all captures from all matches inside {node}.
+	{bufnr} is needed if the query contains predicates, then the caller
 	must ensure to use a freshly parsed tree consistent with the current
-	text of the buffer. `start_row` and `end_row` can be used to limit
+	text of the buffer. {start_row} and {end_row} can be used to limit
 	matches inside a row range (this is typically used with root node
 	as the node, i e to get syntax highlight matches in the current
 	viewport)
@@ -704,7 +712,7 @@ query:iter_captures(node, bufnr, start_row, end_row)
 	  ... use the info here ...
 	end
 <
-query:iter_matches(node, bufnr, start_row, end_row)
+query:iter_matches({node}, {bufnr}, {start_row}, {end_row})
 							*query:iter_matches()*
 	Iterate over all matches within a node. The arguments are the same as
 	for |query:iter_captures()| but the iterated values are different:
@@ -721,8 +729,52 @@ query:iter_matches(node, bufnr, start_row, end_row)
 	    ... use the info here ...
 	  end
 	end
->
-Treesitter syntax highlighting (WIP)		    *lua-treesitter-highlight*
+
+Treesitter Query Predicates				*lua-treesitter-predicates*
+
+When writing queries for treesitter, one might use `predicates`, that is,
+special scheme nodes that are evaluted to verify things on a captured node for
+example, the |eq?| predicate : >
+	((identifier) @foo (#eq? @foo "foo"))
+
+This will only match identifier corresponding to the `"foo"` text.
+Here is a list of built-in predicates :
+
+	`eq?`						*ts-predicate-eq?*
+		This predicate will check text correspondance between nodes or
+		strings : >
+			((identifier) @foo (#eq? @foo "foo"))
+			((node1) @left (node2) @right (#eq? @left @right))
+<
+	`match?`					*ts-predicate-match?*
+		This will match if the provived lua regex matches the text
+		corresponding to a node : >
+			((idenfitier) @constant (#match? @constant "^[A-Z_]+$"))
+<		Note: the `^` and `$` anchors will respectively match the
+			start and end of the node's text.
+
+	`vim-match?`					*ts-predicate-vim-match?*
+		This will match the same way than |match?| but using vim
+		regexes.
+	
+	`contains?`					*ts-predicate-contains?*
+		Will check if any of the following arguments appears in the
+		text corresponding to the node : >
+			((identifier) @foo (#contains? @foo "foo"))
+			((identifier) @foo-bar (#contains @foo-bar "foo" "bar"))
+<
+							*lua-treesitter-not-predicate*
+Each predicate has a `not-` prefixed predicate that is just the negation of
+the predicate.
+
+					*vim.treesitter.query.add_predicate()*
+vim.treesitter.query.add_predicate({name}, {handler})
+
+This adds a predicate with the name {name} to be used in queries.
+{handler} should be a function whose signature will be : >
+	handler(match, pattern, bufnr, predicate)
+
+Treesitter syntax highlighting (WIP)			*lua-treesitter-highlight*
 
 NOTE: This is a partially implemented feature, and not usable as a default
 solution yet. What is documented here is a temporary interface indented
diff --git a/runtime/lua/vim/treesitter.lua b/runtime/lua/vim/treesitter.lua
index 927456708c..550dee1e3f 100644
--- a/runtime/lua/vim/treesitter.lua
+++ b/runtime/lua/vim/treesitter.lua
@@ -1,4 +1,6 @@
 local a = vim.api
+local query = require'vim.treesitter.query'
+local language = require'vim.treesitter.language'
 
 -- TODO(bfredl): currently we retain parsers for the lifetime of the buffer.
 -- Consider use weak references to release parser if all plugins are done with
@@ -8,6 +10,12 @@ local parsers = {}
 local Parser = {}
 Parser.__index = Parser
 
+--- Parses the buffer if needed and returns a tree.
+--
+-- Calling this will call the on_changedtree callbacks if the tree has changed.
+--
+-- @returns An up to date tree
+-- @returns If the tree changed with this call, the changed ranges
 function Parser:parse()
   if self.valid then
     return self.tree
@@ -38,48 +46,39 @@ function Parser:_on_lines(bufnr, changed_tick, start_row, old_stop_row, stop_row
   end
 end
 
+--- Sets the included ranges for the current parser
+--
+-- @param ranges A table of nodes that will be used as the ranges the parser should include.
 function Parser:set_included_ranges(ranges)
   self._parser:set_included_ranges(ranges)
   -- The buffer will need to be parsed again later
   self.valid = false
 end
 
-local M = {
-  parse_query = vim._ts_parse_query,
-}
+local M = vim.tbl_extend("error", query, language)
 
 setmetatable(M, {
   __index = function (t, k)
       if k == "TSHighlighter" then
-        t[k] = require'vim.tshighlighter'
+        a.nvim_err_writeln("vim.TSHighlighter is deprecated, please use vim.treesitter.highlighter")
+        t[k] = require'vim.treesitter.highlighter'
+        return t[k]
+      elseif k == "highlighter" then
+        t[k] = require'vim.treesitter.highlighter'
         return t[k]
       end
    end
  })
 
-function M.require_language(lang, path)
-  if vim._ts_has_language(lang) then
-    return true
-  end
-  if path == nil then
-    local fname = 'parser/' .. lang .. '.*'
-    local paths = a.nvim_get_runtime_file(fname, false)
-    if #paths == 0 then
-      -- TODO(bfredl): help tag?
-      error("no parser for '"..lang.."' language")
-    end
-    path = paths[1]
-  end
-  vim._ts_add_language(path, lang)
-end
-
-function M.inspect_language(lang)
-  M.require_language(lang)
-  return vim._ts_inspect_language(lang)
-end
-
-function M.create_parser(bufnr, lang, id)
-  M.require_language(lang)
+--- Creates a new parser.
+--
+-- It is not recommended to use this, use vim.treesitter.get_parser() instead.
+--
+-- @param bufnr The buffer the parser will be tied to
+-- @param lang The language of the parser.
+-- @param id The id the parser will have
+function M._create_parser(bufnr, lang, id)
+  language.require_language(lang)
   if bufnr == 0 then
     bufnr = a.nvim_get_current_buf()
   end
@@ -91,8 +90,8 @@ function M.create_parser(bufnr, lang, id)
   self.changedtree_cbs = {}
   self.lines_cbs = {}
   self:parse()
-    -- TODO(bfredl): use weakref to self, so that the parser is free'd is no plugin is
-    -- using it.
+  -- TODO(bfredl): use weakref to self, so that the parser is free'd is no plugin is
+  -- using it.
   local function lines_cb(_, ...)
     return self:_on_lines(...)
   end
@@ -108,17 +107,31 @@ function M.create_parser(bufnr, lang, id)
   return self
 end
 
-function M.get_parser(bufnr, ft, buf_attach_cbs)
+--- Gets the parser for this bufnr / ft combination.
+--
+-- If needed this will create the parser.
+-- Unconditionnally attach the provided callback
+--
+-- @param bufnr The buffer the parser should be tied to
+-- @param ft The filetype of this parser
+-- @param buf_attach_cbs An `nvim_buf_attach`-like table argument with the following keys :
+--  `on_lines` : see `nvim_buf_attach`, but this will be called _after_ the parsers callback.
+--  `on_changedtree` : a callback that will be called everytime the tree has syntactical changes.
+--      it will only be passed one argument, that is a table of the ranges (as node ranges) that
+--      changed.
+--
+-- @returns The parser
+function M.get_parser(bufnr, lang, buf_attach_cbs)
   if bufnr == nil or bufnr == 0 then
     bufnr = a.nvim_get_current_buf()
   end
-  if ft == nil then
-    ft = a.nvim_buf_get_option(bufnr, "filetype")
+  if lang == nil then
+    lang = a.nvim_buf_get_option(bufnr, "filetype")
   end
-  local id = tostring(bufnr)..'_'..ft
+  local id = tostring(bufnr)..'_'..lang
 
   if parsers[id] == nil then
-    parsers[id] = M.create_parser(bufnr, ft, id)
+    parsers[id] = M._create_parser(bufnr, lang, id)
   end
 
   if buf_attach_cbs and buf_attach_cbs.on_changedtree then
@@ -132,129 +145,4 @@ function M.get_parser(bufnr, ft, buf_attach_cbs)
   return parsers[id]
 end
 
--- query: pattern matching on trees
--- predicate matching is implemented in lua
-local Query = {}
-Query.__index = Query
-
-local magic_prefixes = {['\\v']=true, ['\\m']=true, ['\\M']=true, ['\\V']=true}
-local function check_magic(str)
-  if string.len(str) < 2 or magic_prefixes[string.sub(str,1,2)] then
-    return str
-  end
-  return '\\v'..str
-end
-
-function M.parse_query(lang, query)
-  M.require_language(lang)
-  local self = setmetatable({}, Query)
-  self.query = vim._ts_parse_query(lang, vim.fn.escape(query,'\\'))
-  self.info = self.query:inspect()
-  self.captures = self.info.captures
-  self.regexes = {}
-  for id,preds in pairs(self.info.patterns) do
-    local regexes = {}
-    for i, pred in ipairs(preds) do
-      if (pred[1] == "match?" and type(pred[2]) == "number"
-          and type(pred[3]) == "string") then
-        regexes[i] = vim.regex(check_magic(pred[3]))
-      end
-    end
-    if next(regexes) then
-      self.regexes[id] = regexes
-    end
-  end
-  return self
-end
-
-local function get_node_text(node, bufnr)
-  local start_row, start_col, end_row, end_col = node:range()
-  if start_row ~= end_row then
-    return nil
-  end
-  local line = a.nvim_buf_get_lines(bufnr, start_row, start_row+1, true)[1]
-  return string.sub(line, start_col+1, end_col)
-end
-
-function Query:match_preds(match, pattern, bufnr)
-  local preds = self.info.patterns[pattern]
-  if not preds then
-    return true
-  end
-  local regexes = self.regexes[pattern]
-  for i, pred in pairs(preds) do
-    -- Here we only want to return if a predicate DOES NOT match, and
-    -- continue on the other case. This way unknown predicates will not be considered,
-    -- which allows some testing and easier user extensibility (#12173).
-    -- Also, tree-sitter strips the leading # from predicates for us.
-    if pred[1] == "eq?" then
-      local node = match[pred[2]]
-      local node_text = get_node_text(node, bufnr)
-
-      local str
-      if type(pred[3]) == "string" then
-        -- (#eq? @aa "foo")
-        str = pred[3]
-      else
-        -- (#eq? @aa @bb)
-        str = get_node_text(match[pred[3]], bufnr)
-      end
-
-      if node_text ~= str or str == nil then
-        return false
-      end
-    elseif pred[1] == "match?" then
-      if not regexes or not regexes[i] then
-        return false
-      end
-      local node = match[pred[2]]
-      local start_row, start_col, end_row, end_col = node:range()
-      if start_row ~= end_row then
-        return false
-      end
-      if not regexes[i]:match_line(bufnr, start_row, start_col, end_col) then
-        return false
-      end
-    end
-  end
-  return true
-end
-
-function Query:iter_captures(node, bufnr, start, stop)
-  if bufnr == 0 then
-    bufnr = vim.api.nvim_get_current_buf()
-  end
-  local raw_iter = node:_rawquery(self.query,true,start,stop)
-  local function iter()
-    local capture, captured_node, match = raw_iter()
-    if match ~= nil then
-      local active = self:match_preds(match, match.pattern, bufnr)
-      match.active = active
-      if not active then
-        return iter() -- tail call: try next match
-      end
-    end
-    return capture, captured_node
-  end
-  return iter
-end
-
-function Query:iter_matches(node, bufnr, start, stop)
-  if bufnr == 0 then
-    bufnr = vim.api.nvim_get_current_buf()
-  end
-  local raw_iter = node:_rawquery(self.query,false,start,stop)
-  local function iter()
-    local pattern, match = raw_iter()
-    if match ~= nil then
-      local active = self:match_preds(match, pattern, bufnr)
-      if not active then
-        return iter() -- tail call: try next match
-      end
-    end
-    return pattern, match
-  end
-  return iter
-end
-
 return M
diff --git a/runtime/lua/vim/tshighlighter.lua b/runtime/lua/vim/treesitter/highlighter.lua
index 6465751ae8..681d2c6324 100644
--- a/runtime/lua/vim/tshighlighter.lua
+++ b/runtime/lua/vim/treesitter/highlighter.lua
@@ -7,20 +7,50 @@ local ts_hs_ns = a.nvim_create_namespace("treesitter_hl")
 
 -- These are conventions defined by tree-sitter, though it
 -- needs to be user extensible also.
--- TODO(bfredl): this is very much incomplete, we will need to
--- go through a few tree-sitter provided queries and decide
--- on translations that makes the most sense.
 TSHighlighter.hl_map = {
-    keyword="Keyword",
-    string="String",
-    type="Type",
-    comment="Comment",
-    constant="Constant",
-    operator="Operator",
-    number="Number",
-    label="Label",
-    ["function"]="Function",
-    ["function.special"]="Function",
+    ["error"] = "Error",
+
+-- Miscs
+    ["comment"] = "Comment",
+    ["punctuation.delimiter"] = "Delimiter",
+    ["punctuation.bracket"] = "Delimiter",
+    ["punctuation.special"] = "Delimiter",
+
+-- Constants
+    ["constant"] = "Constant",
+    ["constant.builtin"] = "Special",
+    ["constant.macro"] = "Define",
+    ["string"] = "String",
+    ["string.regex"] = "String",
+    ["string.escape"] = "SpecialChar",
+    ["character"] = "Character",
+    ["number"] = "Number",
+    ["boolean"] = "Boolean",
+    ["float"] = "Float",
+
+-- Functions
+    ["function"] = "Function",
+    ["function.special"] = "Function",
+    ["function.builtin"] = "Special",
+    ["function.macro"] = "Macro",
+    ["parameter"] = "Identifier",
+    ["method"] = "Function",
+    ["field"] = "Identifier",
+    ["property"] = "Identifier",
+    ["constructor"] = "Special",
+
+-- Keywords
+    ["conditional"] = "Conditional",
+    ["repeat"] = "Repeat",
+    ["label"] = "Label",
+    ["operator"] = "Operator",
+    ["keyword"] = "Keyword",
+    ["exception"] = "Exception",
+
+    ["type"] = "Type",
+    ["type.builtin"] = "Type",
+    ["structure"] = "Structure",
+    ["include"] = "Include",
 }
 
 function TSHighlighter.new(query, bufnr, ft)
@@ -75,7 +105,15 @@ end
 function TSHighlighter:set_query(query)
   if type(query) == "string" then
     query = vim.treesitter.parse_query(self.parser.lang, query)
+  elseif query == nil then
+    query = vim.treesitter.get_query(self.parser.lang, 'highlights')
+
+    if query == nil then
+      a.nvim_err_writeln("No highlights.scm query found for " .. self.parser.lang)
+      query = vim.treesitter.parse_query(self.parser.lang, "")
+    end
   end
+
   self.query = query
 
   self.hl_cache = setmetatable({}, {
diff --git a/runtime/lua/vim/treesitter/language.lua b/runtime/lua/vim/treesitter/language.lua
new file mode 100644
index 0000000000..a7e36a0b89
--- /dev/null
+++ b/runtime/lua/vim/treesitter/language.lua
@@ -0,0 +1,37 @@
+local a = vim.api
+
+local M = {}
+
+--- Asserts that the provided language is installed, and optionnaly provide a path for the parser
+--
+-- Parsers are searched in the `parser` runtime directory.
+--
+-- @param lang The language the parser should parse
+-- @param path Optionnal path the parser is located at
+function M.require_language(lang, path)
+  if vim._ts_has_language(lang) then
+    return true
+  end
+  if path == nil then
+    local fname = 'parser/' .. lang .. '.*'
+    local paths = a.nvim_get_runtime_file(fname, false)
+    if #paths == 0 then
+      -- TODO(bfredl): help tag?
+      error("no parser for '"..lang.."' language, see :help treesitter-parsers")
+    end
+    path = paths[1]
+  end
+  vim._ts_add_language(path, lang)
+end
+
+--- Inspects the provided language.
+--
+-- Inspecting provides some useful informations on the language like node names, ...
+--
+-- @param lang The language.
+function M.inspect_language(lang)
+  M.require_language(lang)
+  return vim._ts_inspect_language(lang)
+end
+
+return M
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
new file mode 100644
index 0000000000..17f61b24f1
--- /dev/null
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -0,0 +1,210 @@
+local a = vim.api
+local language = require'vim.treesitter.language'
+
+-- query: pattern matching on trees
+-- predicate matching is implemented in lua
+local Query = {}
+Query.__index = Query
+
+local M = {}
+
+--- Parses a query.
+--
+-- @param language The language
+-- @param query A string containing the query (s-expr syntax)
+--
+-- @returns The query
+function M.parse_query(lang, query)
+  language.require_language(lang)
+  local self = setmetatable({}, Query)
+  self.query = vim._ts_parse_query(lang, vim.fn.escape(query,'\\'))
+  self.info = self.query:inspect()
+  self.captures = self.info.captures
+  return self
+end
+
+-- TODO(vigoux): support multiline nodes too
+
+--- Gets the text corresponding to a given node
+-- @param node the node
+-- @param bufnr the buffer from which the node in extracted.
+function M.get_node_text(node, bufnr)
+  local start_row, start_col, end_row, end_col = node:range()
+  if start_row ~= end_row then
+    return nil
+  end
+  local line = a.nvim_buf_get_lines(bufnr, start_row, start_row+1, true)[1]
+  return string.sub(line, start_col+1, end_col)
+end
+
+-- Predicate handler receive the following arguments
+-- (match, pattern, bufnr, predicate)
+local predicate_handlers = {
+  ["eq?"] = function(match, _, bufnr, predicate)
+      local node = match[predicate[2]]
+      local node_text = M.get_node_text(node, bufnr)
+
+      local str
+      if type(predicate[3]) == "string" then
+        -- (#eq? @aa "foo")
+        str = predicate[3]
+      else
+        -- (#eq? @aa @bb)
+        str = M.get_node_text(match[predicate[3]], bufnr)
+      end
+
+      if node_text ~= str or str == nil then
+        return false
+      end
+
+      return true
+  end,
+
+  ["match?"] = function(match, _, bufnr, predicate)
+      local node = match[predicate[2]]
+      local regex = predicate[3]
+      local start_row, _, end_row, _ = node:range()
+      if start_row ~= end_row then
+        return false
+      end
+
+      return string.find(M.get_node_text(node, bufnr), regex)
+  end,
+
+  ["vim-match?"] = (function()
+    local magic_prefixes = {['\\v']=true, ['\\m']=true, ['\\M']=true, ['\\V']=true}
+    local function check_magic(str)
+      if string.len(str) < 2 or magic_prefixes[string.sub(str,1,2)] then
+        return str
+      end
+      return '\\v'..str
+    end
+
+    local compiled_vim_regexes = setmetatable({}, {
+      __index = function(t, pattern)
+        local res = vim.regex(check_magic(pattern))
+        rawset(t, pattern, res)
+        return res
+      end
+    })
+
+    return function(match, _, bufnr, pred)
+      local node = match[pred[2]]
+      local start_row, start_col, end_row, end_col = node:range()
+      if start_row ~= end_row then
+        return false
+      end
+
+      local regex = compiled_vim_regexes[pred[3]]
+      return regex:match_line(bufnr, start_row, start_col, end_col)
+    end
+  end)(),
+
+  ["contains?"] = function(match, _, bufnr, predicate)
+    local node = match[predicate[2]]
+    local node_text = M.get_node_text(node, bufnr)
+
+    for i=3,#predicate do
+      if string.find(node_text, predicate[i], 1, true) then
+        return true
+      end
+    end
+
+    return false
+  end
+}
+
+--- Adds a new predicates to be used in queries
+--
+-- @param name the name of the predicate, without leading #
+-- @param handler the handler function to be used
+--    signature will be (match, pattern, bufnr, predicate)
+function M.add_predicate(name, handler, force)
+  if predicate_handlers[name] and not force then
+    a.nvim_err_writeln(string.format("Overriding %s", name))
+  end
+
+  predicate_handlers[name] = handler
+end
+
+function Query:match_preds(match, pattern, bufnr)
+  local preds = self.info.patterns[pattern]
+  if not preds then
+    return true
+  end
+  for _, pred in pairs(preds) do
+    -- Here we only want to return if a predicate DOES NOT match, and
+    -- continue on the other case. This way unknown predicates will not be considered,
+    -- which allows some testing and easier user extensibility (#12173).
+    -- Also, tree-sitter strips the leading # from predicates for us.
+    if string.sub(pred[1], 1, 4) == "not-" then
+      local pred_name = string.sub(pred[1], 5)
+      if predicate_handlers[pred_name] and
+        predicate_handlers[pred_name](match, pattern, bufnr, pred) then
+        return false
+      end
+
+    elseif predicate_handlers[pred[1]] and
+      not predicate_handlers[pred[1]](match, pattern, bufnr, pred) then
+      return false
+    end
+  end
+  return true
+end
+
+--- Iterates of the captures of self on a given range.
+--
+-- @param node The node under witch the search will occur
+-- @param buffer The source buffer to search
+-- @param start The starting line of the search
+-- @param stop The stoping line of the search (end-exclusive)
+--
+-- @returns The matching capture id
+-- @returns The captured node
+function Query:iter_captures(node, bufnr, start, stop)
+  if bufnr == 0 then
+    bufnr = vim.api.nvim_get_current_buf()
+  end
+  local raw_iter = node:_rawquery(self.query, true, start, stop)
+  local function iter()
+    local capture, captured_node, match = raw_iter()
+    if match ~= nil then
+      local active = self:match_preds(match, match.pattern, bufnr)
+      match.active = active
+      if not active then
+        return iter() -- tail call: try next match
+      end
+    end
+    return capture, captured_node
+  end
+  return iter
+end
+
+--- Iterates of the matches of self on a given range.
+--
+-- @param node The node under witch the search will occur
+-- @param buffer The source buffer to search
+-- @param start The starting line of the search
+-- @param stop The stoping line of the search (end-exclusive)
+--
+-- @returns The matching pattern id
+-- @returns The matching match
+function Query:iter_matches(node, bufnr, start, stop)
+  if bufnr == 0 then
+    bufnr = vim.api.nvim_get_current_buf()
+  end
+  local raw_iter = node:_rawquery(self.query, false, start, stop)
+  local function iter()
+    local pattern, match = raw_iter()
+    if match ~= nil then
+      local active = self:match_preds(match, pattern, bufnr)
+      if not active then
+        return iter() -- tail call: try next match
+      end
+    end
+    return pattern, match
+  end
+  return iter
+end
+
+return M
diff --git a/test/functional/lua/treesitter_spec.lua b/test/functional/lua/treesitter_spec.lua
index aa3d55b06d..b0ac9e079a 100644
--- a/test/functional/lua/treesitter_spec.lua
+++ b/test/functional/lua/treesitter_spec.lua
@@ -15,14 +15,14 @@ before_each(clear)
 describe('treesitter API', function()
   -- error tests not requiring a parser library
   it('handles missing language', function()
-    eq("Error executing lua: .../treesitter.lua: no parser for 'borklang' language",
-       pcall_err(exec_lua, "parser = vim.treesitter.create_parser(0, 'borklang')"))
+    eq("Error executing lua: .../language.lua: no parser for 'borklang' language, see :help treesitter-parsers",
+       pcall_err(exec_lua, "parser = vim.treesitter.get_parser(0, 'borklang')"))
 
     -- actual message depends on platform
     matches("Error executing lua: Failed to load parser: uv_dlopen: .+",
        pcall_err(exec_lua, "parser = vim.treesitter.require_language('borklang', 'borkbork.so')"))
 
-    eq("Error executing lua: .../treesitter.lua: no parser for 'borklang' language",
+    eq("Error executing lua: .../language.lua: no parser for 'borklang' language, see :help treesitter-parsers",
        pcall_err(exec_lua, "parser = vim.treesitter.inspect_language('borklang')"))
   end)
 
@@ -198,6 +198,41 @@ void ui_refresh(void)
     }, res)
   end)
 
+  it('allows to add predicates', function()
+    insert([[
+    int main(void) {
+      return 0;
+    }
+    ]])
+
+    local custom_query = "((identifier) @main (#is-main? @main))"
+
+    local res = exec_lua([[
+    local query = require"vim.treesitter.query"
+
+    local function is_main(match, pattern, bufnr, predicate)
+      local node = match[ predicate[2] ]
+
+      return query.get_node_text(node, bufnr)
+    end
+
+    local parser = vim.treesitter.get_parser(0, "c")
+
+    query.add_predicate("is-main?", is_main)
+
+    local query = query.parse_query("c", ...)
+
+    local nodes = {}
+    for _, node in query:iter_captures(parser:parse():root(), 0, 0, 19) do
+      table.insert(nodes, {node:range()})
+    end
+
+    return nodes
+    ]], custom_query)
+
+    eq({{0, 4, 0, 8}}, res)
+  end)
+
   it('supports highlighting', function()
     if not check_parser() then return end
 
@@ -243,10 +278,10 @@ static int nlua_schedule(lua_State *const lstate)
 (primitive_type) @type
 (sized_type_specifier) @type
 
-; defaults to very magic syntax, for best compatibility
-((identifier) @Identifier (#match? @Identifier "^l(u)a_"))
-; still support \M etc prefixes
-((identifier) @Constant (#match? @Constant "\M^\[A-Z_]\+$"))
+; Use lua regexes
+((identifier) @Identifier (#contains? @Identifier "lua_"))
+((identifier) @Constant (#match? @Constant "^[A-Z_]+$"))
+((identifier) @Normal (#vim-match? @Constant "^lstate$"))
 
 ((binary_expression left: (identifier) @WarningMsg.left right: (identifier) @WarningMsg.right) (#eq? @WarningMsg.left @WarningMsg.right))
 
@@ -292,13 +327,13 @@ static int nlua_schedule(lua_State *const lstate)
     ]]}
 
     exec_lua([[
-      local TSHighlighter = vim.treesitter.TSHighlighter
+      local highlighter = vim.treesitter.highlighter
       local query = ...
-      test_hl = TSHighlighter.new(query, 0, "c")
+      test_hl = highlighter.new(query, 0, "c")
     ]], hl_query)
     screen:expect{grid=[[
       {2:/// Schedule Lua callback on main loop's event queue}             |
-      {3:static} {3:int} nlua_schedule({3:lua_State} *{3:const} lstate)                |
+      {3:static} {3:int} {11:nlua_schedule}({3:lua_State} *{3:const} lstate)                |
       {                                                                |
         {4:if} ({11:lua_type}(lstate, {5:1}) != {5:LUA_TFUNCTION}                       |
             || {6:lstate} != {6:lstate}) {                                     |
@@ -306,9 +341,9 @@ static int nlua_schedule(lua_State *const lstate)
           {4:return} {11:lua_error}(lstate);                                    |
         }                                                              |
                                                                        |
-        {7:LuaRef} cb = nlua_ref(lstate, {5:1});                               |
+        {7:LuaRef} cb = {11:nlua_ref}(lstate, {5:1});                               |
                                                                        |
-        multiqueue_put(main_loop.events, nlua_schedule_event,          |
+        multiqueue_put(main_loop.events, {11:nlua_schedule_event},          |
                        {5:1}, ({3:void} *)({3:ptrdiff_t})cb);                      |
         {4:return} {5:0};                                                      |
       ^}                                                                |
@@ -320,7 +355,7 @@ static int nlua_schedule(lua_State *const lstate)
     feed('7Go*/<esc>')
     screen:expect{grid=[[
       {2:/// Schedule Lua callback on main loop's event queue}             |
-      {3:static} {3:int} nlua_schedule({3:lua_State} *{3:const} lstate)                |
+      {3:static} {3:int} {11:nlua_schedule}({3:lua_State} *{3:const} lstate)                |
       {                                                                |
         {4:if} ({11:lua_type}(lstate, {5:1}) != {5:LUA_TFUNCTION}                       |
             || {6:lstate} != {6:lstate}) {                                     |
@@ -329,9 +364,9 @@ static int nlua_schedule(lua_State *const lstate)
       {8:*^/}                                                               |
         }                                                              |
                                                                        |
-        {7:LuaRef} cb = nlua_ref(lstate, {5:1});                               |
+        {7:LuaRef} cb = {11:nlua_ref}(lstate, {5:1});                               |
                                                                        |
-        multiqueue_put(main_loop.events, nlua_schedule_event,          |
+        multiqueue_put(main_loop.events, {11:nlua_schedule_event},          |
                        {5:1}, ({3:void} *)({3:ptrdiff_t})cb);                      |
         {4:return} {5:0};                                                      |
       }                                                                |
@@ -342,7 +377,7 @@ static int nlua_schedule(lua_State *const lstate)
     feed('3Go/*<esc>')
     screen:expect{grid=[[
       {2:/// Schedule Lua callback on main loop's event queue}             |
-      {3:static} {3:int} nlua_schedule({3:lua_State} *{3:const} lstate)                |
+      {3:static} {3:int} {11:nlua_schedule}({3:lua_State} *{3:const} lstate)                |
       {                                                                |
       {2:/^*}                                                               |
       {2:  if (lua_type(lstate, 1) != LUA_TFUNCTION}                       |
@@ -352,9 +387,9 @@ static int nlua_schedule(lua_State *const lstate)
       {2:*/}                                                               |
         }                                                              |
                                                                        |
-        {7:LuaRef} cb = nlua_ref(lstate, {5:1});                               |
+        {7:LuaRef} cb = {11:nlua_ref}(lstate, {5:1});                               |
                                                                        |
-        multiqueue_put(main_loop.events, nlua_schedule_event,          |
+        multiqueue_put(main_loop.events, {11:nlua_schedule_event},          |
                        {5:1}, ({3:void} *)({3:ptrdiff_t})cb);                      |
         {4:return} {5:0};                                                      |
       {8:}}                                                                |
@@ -365,7 +400,7 @@ static int nlua_schedule(lua_State *const lstate)
     feed("~")
     screen:expect{grid=[[
       {2:/// Schedule Lua callback on main loop's event queu^E}             |
-      {3:static} {3:int} nlua_schedule({3:lua_State} *{3:const} lstate)                |
+      {3:static} {3:int} {11:nlua_schedule}({3:lua_State} *{3:const} lstate)                |
       {                                                                |
       {2:/*}                                                               |
       {2:  if (lua_type(lstate, 1) != LUA_TFUNCTION}                       |
@@ -375,9 +410,9 @@ static int nlua_schedule(lua_State *const lstate)
       {2:*/}                                                               |
         }                                                              |
                                                                        |
-        {7:LuaRef} cb = nlua_ref(lstate, {5:1});                               |
+        {7:LuaRef} cb = {11:nlua_ref}(lstate, {5:1});                               |
                                                                        |
-        multiqueue_put(main_loop.events, nlua_schedule_event,          |
+        multiqueue_put(main_loop.events, {11:nlua_schedule_event},          |
                        {5:1}, ({3:void} *)({3:ptrdiff_t})cb);                      |
         {4:return} {5:0};                                                      |
       {8:}}                                                                |
@@ -388,7 +423,7 @@ static int nlua_schedule(lua_State *const lstate)
     feed("re")
     screen:expect{grid=[[
       {2:/// Schedule Lua callback on main loop's event queu^e}             |
-      {3:static} {3:int} nlua_schedule({3:lua_State} *{3:const} lstate)                |
+      {3:static} {3:int} {11:nlua_schedule}({3:lua_State} *{3:const} lstate)                |
       {                                                                |
       {2:/*}                                                               |
       {2:  if (lua_type(lstate, 1) != LUA_TFUNCTION}                       |
@@ -398,9 +433,9 @@ static int nlua_schedule(lua_State *const lstate)
       {2:*/}                                                               |
         }                                                              |
                                                                        |
-        {7:LuaRef} cb = nlua_ref(lstate, {5:1});                               |
+        {7:LuaRef} cb = {11:nlua_ref}(lstate, {5:1});                               |
                                                                        |
-        multiqueue_put(main_loop.events, nlua_schedule_event,          |
+        multiqueue_put(main_loop.events, {11:nlua_schedule_event},          |
                        {5:1}, ({3:void} *)({3:ptrdiff_t})cb);                      |
         {4:return} {5:0};                                                      |
       {8:}}                                                                |