nvim-treesitter/lua/nvim-treesitter/tsrange.lua

local M = {}
local TSRange = {}
TSRange.__index = TSRange

local api = vim.api
local parsers = require'nvim-treesitter.parsers'
local ts_utils = require'nvim-treesitter.ts_utils'

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 and {start_node:start()} or {end_node:start()}
  local end_pos = end_node and {end_node:end_()} or {start_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 = ts_utils.get_root_for_position(range[1], range[2], parser)

  return root
    and root:named_descendant_for_range(self.start_pos[1], self.start_pos[2], self.end_pos[1], self.end_pos[2])
    or nil
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