treesitter: add and test vim-match? predicate

runtime/lua/vim/treesitter/query.lua

@@ -24,7 +24,11 @@ function M.parse_query(lang, query)
 end
 
 -- TODO(vigoux): support multiline nodes too
-local function get_node_text(node, bufnr)
+
+--- 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
@@ -34,11 +38,11 @@ local function get_node_text(node, bufnr)
 end
 
 -- Predicate handler receive the following arguments
--- (match, pattern, bufnr, regexes, index, predicate)
+-- (match, pattern, bufnr, predicate)
 local predicate_handlers = {
   ["eq?"] = function(match, _, bufnr, predicate)
       local node = match[predicate[2]]
-      local node_text = get_node_text(node, bufnr)
+      local node_text = M.get_node_text(node, bufnr)
 
       local str
       if type(predicate[3]) == "string" then
@@ -46,7 +50,7 @@ local predicate_handlers = {
         str = predicate[3]
       else
         -- (#eq? @aa @bb)
-        str = get_node_text(match[predicate[3]], bufnr)
+        str = M.get_node_text(match[predicate[3]], bufnr)
       end
 
       if node_text ~= str or str == nil then
@@ -63,12 +67,42 @@ local predicate_handlers = {
         return false
       end
 
-      return string.find(get_node_text(node, bufnr), regex)
+      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 = get_node_text(node, bufnr)
+    local node_text = M.get_node_text(node, bufnr)
 
     for i=3,#predicate do
       if string.find(node_text, predicate[i], 1, true) then

test/functional/lua/treesitter_spec.lua

@@ -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
 
@@ -246,6 +281,7 @@ static int nlua_schedule(lua_State *const lstate)
 ; 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))