lua/nvim-treesitter/query_predicates.lua


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

local query = require"vim.treesitter.query"

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
  end

  local node = match[pred[2]]
  local n = pred[3] - 1
  if node and node:parent() and node:named_child_count() > n then
    return node:named_child(n) == node
  end

  return false
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

  local node = match[pred[2]]
  local ancestor_type = pred[3]
  if not node then return true end

  node = node:parent()
  while node do
    if node:type() == ancestor_type then
      return true
    end
    node = node:parent()
  end
  return false
end)