Add virtual range for creating node-like objects from multiple nodes

3 files changed, 151 insertions, 0 deletions

diff --git a/lua/nvim-treesitter/query.lua b/lua/nvim-treesitter/query.lua
index 09fdbbdde..4cb5993b4 100644
--- a/lua/nvim-treesitter/query.lua
+++ b/lua/nvim-treesitter/query.lua
@@ -1,5 +1,6 @@
 local api = vim.api
 local tsq = require'vim.treesitter.query'
+local tsrange = require'nvim-treesitter.tsrange'
 local utils = require'nvim-treesitter.utils'
 local parsers = require'nvim-treesitter.parsers'
 local caching = require'nvim-treesitter.caching'
@@ -98,6 +99,10 @@ function M.iter_prepared_matches(query, qnode, bufnr, start_row, end_row)
           if pred[1] == "set!" and type(pred[2]) == "string" then
             insert_to_path(prepared_match, split(pred[2]), pred[3])
           end
+          if pred[1] == "make-range!" and type(pred[2]) == "string" and #pred == 4 then
+            insert_to_path(prepared_match, split(pred[2]..'.node'),
+                           tsrange.TSRange.from_nodes(bufnr, match[pred[3]], match[pred[4]]))
+          end
         end
       end
 
diff --git a/lua/nvim-treesitter/query_predicates.lua b/lua/nvim-treesitter/query_predicates.lua
index 0bcb905ec..0c0728997 100644
--- a/lua/nvim-treesitter/query_predicates.lua
+++ b/lua/nvim-treesitter/query_predicates.lua
@@ -66,3 +66,4 @@ end)
 
 -- Just avoid some anoying warnings for this predicate
 query.add_predicate('set!', function() return true end)
+query.add_predicate('make-range!', function() return true end)
diff --git a/lua/nvim-treesitter/tsrange.lua b/lua/nvim-treesitter/tsrange.lua
new file mode 100644
index 000000000..e6aa73f5d
--- /dev/null
+++ b/lua/nvim-treesitter/tsrange.lua
@@ -0,0 +1,145 @@
+local M = {}
+local TSRange = {}
+TSRange.__index = TSRange
+
+local api = vim.api
+local parsers = require'nvim-treesitter.parsers'
+
+local function get_byte_offset(buf, row, col)
+  return api.nvim_buf_get_offset(buf, row)
+         + vim.fn.byteidx(api.nvim_buf_get_lines(buf, row, row + 1, false), col)
+end
+
+function TSRange.new(buf, start_row, start_col, end_row, end_col)
+  return setmetatable(
+    {
+      start_pos = {start_row, start_col, get_byte_offset(buf, start_row, start_col)},
+      end_pos = {end_row, end_col, get_byte_offset(buf, end_row, end_col)},
+      buf = buf,
+      [1] = start_row,
+      [2] = start_col,
+      [3] = end_row,
+      [4] = end_col,
+    },
+    TSRange)
+end
+
+function TSRange.from_nodes(buf, start_node, end_node)
+  TSRange.__index = TSRange
+  local start_pos = {start_node:start()}
+  local end_pos = {end_node:end_()}
+  return setmetatable(
+    {
+      start_pos = {start_pos[1], start_pos[2], start_pos[3]},
+      end_pos = {end_pos[1], end_pos[2], end_pos[3]},
+      buf = buf,
+      [1] = start_pos[1],
+      [2] = start_pos[2],
+      [3] = end_pos[1],
+      [4] = end_pos[2],
+    },
+    TSRange)
+end
+
+function TSRange.from_table(buf, range)
+  return setmetatable(
+    {
+      start_pos = {range[1], range[2], get_byte_offset(buf, range[1], range[2])},
+      end_pos = {range[3], range[4], get_byte_offset(buf, range[3], range[4])},
+      buf = buf,
+      [1] = range[1],
+      [2] = range[2],
+      [3] = range[3],
+      [4] = range[4],
+    },
+    TSRange)
+end
+
+function TSRange:parent(range)
+  local parser = parsers.get_parser(self.buf, parsers.get_buf_lang(range))
+  local root = parser.tree:root()
+  return root:named_descendant_for_range(self.start_pos[1], self.start_pos[2], self.end_pos[1], self.end_pos[2])
+end
+
+function TSRange:field()
+end
+
+function TSRange:child_count()
+  return #self:collect_children()
+end
+
+function TSRange:named_child_count()
+  return #self:collect_children(function(c) return c:named() end)
+end
+
+function TSRange:iter_children()
+  local raw_iterator = self:parent().iter_children()
+  return function()
+    while true do
+      local node = raw_iterator()
+      if not node then return end
+      local _, _, start_byte = node:start()
+      local _, _, end_byte = node:end_()
+      if start_byte >= self.start_pos[3] and end_byte <= self.end_pos[3] then
+        return node
+      end
+    end
+  end
+end
+
+function TSRange:collect_children(filter_fun)
+  local children = {}
+  for _, c in self:iter_children() do
+    if not filter_fun or filter_fun(c) then
+      table.insert(children, c)
+    end
+  end
+  return children
+end
+
+function TSRange:child(index)
+  return self:collect_children()[index + 1]
+end
+
+function TSRange:named_child(index)
+  return self:collect_children(function(c) return c.named() end)[index + 1]
+end
+
+function TSRange:start()
+  return unpack(self.start_pos)
+end
+
+function TSRange:end_()
+  return unpack(self.end_pos)
+end
+
+function TSRange:range()
+  return self.start_pos[1], self.start_pos[2], self.end_pos[1], self.end_pos[2]
+end
+
+function TSRange:type()
+  return 'nvim-treesitter-range'
+end
+
+function TSRange:symbol()
+  return -1
+end
+
+function TSRange:named()
+  return false
+end
+
+function TSRange:missing()
+  return false
+end
+
+function TSRange:has_error()
+  return #self:collect_children(function(c) return c:has_error() end) > 0 and true or false
+end
+
+function TSRange:sexpr()
+  return table.concat(vim.tbl_map(function(c) return c:sexpr() end, self:collect_children()), ' ')
+end
+
+M.TSRange = TSRange
+return M