refactor(treesitter): add Range type aliase for Range4|Range6

3 files changed, 24 insertions, 17 deletions

diff --git a/runtime/lua/vim/treesitter/_range.lua b/runtime/lua/vim/treesitter/_range.lua
index 0017a567ec..f4db5016ac 100644
--- a/runtime/lua/vim/treesitter/_range.lua
+++ b/runtime/lua/vim/treesitter/_range.lua
@@ -16,6 +16,8 @@ local M = {}
 ---@field [5] integer end column
 ---@field [6] integer end bytes
 
+---@alias Range Range4|Range6
+
 ---@private
 ---@param a_row integer
 ---@param a_col integer
@@ -85,8 +87,8 @@ function M.validate(r)
 end
 
 ---@private
----@param r1 Range4|Range6
----@param r2 Range4|Range6
+---@param r1 Range
+---@param r2 Range
 ---@return boolean
 function M.intercepts(r1, r2)
   local srow_1, scol_1, erow_1, ecol_1 = M.unpack4(r1)
@@ -106,7 +108,7 @@ function M.intercepts(r1, r2)
 end
 
 ---@private
----@param r Range4|Range6
+---@param r Range
 ---@return integer, integer, integer, integer
 function M.unpack4(r)
   local off_1 = #r == 6 and 1 or 0
@@ -121,8 +123,8 @@ function M.unpack6(r)
 end
 
 ---@private
----@param r1 Range4|Range6
----@param r2 Range4|Range6
+---@param r1 Range
+---@param r2 Range
 ---@return boolean whether r1 contains r2
 function M.contains(r1, r2)
   local srow_1, scol_1, erow_1, ecol_1 = M.unpack4(r1)
@@ -143,7 +145,7 @@ end
 
 ---@private
 ---@param source integer|string
----@param range Range4|Range6
+---@param range Range
 ---@return Range6
 function M.add_bytes(source, range)
   if type(range) == 'table' and #range == 6 then
diff --git a/runtime/lua/vim/treesitter/languagetree.lua b/runtime/lua/vim/treesitter/languagetree.lua
index 0bb0601241..6869fae92c 100644
--- a/runtime/lua/vim/treesitter/languagetree.lua
+++ b/runtime/lua/vim/treesitter/languagetree.lua
@@ -484,7 +484,6 @@ function LanguageTree:included_regions()
 end
 
 ---@private
---- TODO(lewis6991): cleanup of the node_range interface
 ---@param node TSNode
 ---@param source string|integer
 ---@param metadata TSMetadata
@@ -530,7 +529,7 @@ end
 ---@param pattern integer
 ---@param lang string
 ---@param combined boolean
----@param ranges Range4[]
+---@param ranges Range6[]
 local function add_injection(t, tree_index, pattern, lang, combined, ranges)
   assert(type(lang) == 'string')
 
@@ -558,13 +557,11 @@ end
 --- https://tree-sitter.github.io/tree-sitter/syntax-highlighting#language-injection
 ---@param match table<integer,TSNode>
 ---@param metadata TSMetadata
----@return string, boolean, Range4[]
+---@return string?, boolean, Range6[]
 function LanguageTree:_get_injection(match, metadata)
-  local ranges = {} ---@type Range4[]
+  local ranges = {} ---@type Range6[]
   local combined = metadata['injection.combined'] ~= nil
-  local lang = metadata['injection.language']
-  assert(type(lang) == 'string')
-
+  local lang = metadata['injection.language'] --[[@as string?]]
   local include_children = metadata['injection.include-children'] ~= nil
 
   for id, node in pairs(match) do
@@ -584,7 +581,7 @@ end
 ---@private
 ---@param match table<integer,TSNode>
 ---@param metadata TSMetadata
----@return string, boolean, Range4[]
+---@return string, boolean, Range6[]
 function LanguageTree:_get_injection_deprecated(match, metadata)
   local lang = nil ---@type string
   local ranges = {} ---@type Range6[]
@@ -910,7 +907,7 @@ end
 
 ---@private
 ---@param tree TSTree
----@param range Range4
+---@param range Range
 ---@return boolean
 local function tree_contains(tree, range)
   return Range.contains({ tree:root():range() }, range)
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index 70af4f7bce..f4e038b2d8 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -59,10 +59,18 @@ end
 
 ---@private
 ---@param buf integer
----@param range Range6
+---@param range Range
 ---@returns string
 local function buf_range_get_text(buf, range)
   local start_row, start_col, end_row, end_col = Range.unpack4(range)
+  if end_col == 0 then
+    if start_row == end_row then
+      start_col = -1
+      start_row = start_row - 1
+    end
+    end_col = -1
+    end_row = end_row - 1
+  end
   local lines = a.nvim_buf_get_text(buf, start_row, start_col, end_row, end_col, {})
   return table.concat(lines, '\n')
 end
@@ -396,7 +404,7 @@ local predicate_handlers = {
 predicate_handlers['vim-match?'] = predicate_handlers['match?']
 
 ---@class TSMetadata
----@field range Range4|Range6
+---@field range Range
 ---@field [integer] TSMetadata
 ---@field [string] integer|string