feat: extend vim.Pos, vim.Range #36397

Problem:
Using nested `vim.Pos` objects to represent each `vim.Range` object
requires 3 tables for each `vim.Range`, which may be undesirable in
performance critical code. Using key-value tables performs worse than
using array-like tables (lists).

Solution:
Use array-like indices for the internal fields of both `vim.Pos` and
`vim.Range` objects. Use a metatable to allow users to access them like
if they were key-value tables.

---

Problem:
The `vim.Pos` conversion interface for `extmark` indexing does not take
into account the difference in how a position on top of a newline is
represented in `vim.Pos` and `extmark`.
- `vim.Pos`: for a newline at the end of row `n`, `row` takes the value
  `n + 1` and `col` takes the value `0`.
- `extmark`: for a newline at the end of for `n`, `row` takes the value
  `n` and `col` takes the value `#row_text`.

Solution:
Handle this in the `extmark` interface.

---

Problem:
Not all `to_xxx` interfaces have wrapping objects like `to_lsp`.

Solution:
Return unwrapped values in `to_xxx` interfaces where it makes sense.
Accept unwrapped values in "from" interfaces where it makes sense.

---

Problem:
`start` and `end` positions have different semantics, so they can't be
compared. `vim.Range` relies on comparing the `end` and `start` of two
ranges to decide which one is greater, which doesn't work as expected
because this of the different semantics.

For example, for the ranges:

    local a = {
      start = { row = 0, col = 22, },
      end_ = { row = 0, col = 24, },
    }
    local b = {
      start = { row = 0, col = 17, },
      end_ = { row = 0, col = 22, },
    }

in this code:

    local foo, bar = "foo",  "bar"
    --               |---||-|
    --                 b  a

The range `b` is smaller than the range `a`, but the current
implementation compares `b._end` (`col = 22`) and `a.start` (`col = 22`)
and concludes that, since `b.col` is not smaller than `a.col`, `b`
should be greater than `a`.

Solution:
- Use a `to_inclusive_pos` to normalize end positions inside of
  `vim.Range` whenever a comparison between a start and an end position
  is necessary.

1 files changed, 239 insertions, 28 deletions

diff --git a/runtime/lua/vim/range.lua b/runtime/lua/vim/range.lua
index f0cd01145e..78e201e3fd 100644
--- a/runtime/lua/vim/range.lua
+++ b/runtime/lua/vim/range.lua
@@ -8,6 +8,7 @@
 --- objects.
 
 local validate = vim.validate
+local api = vim.api
 
 --- Represents a well-defined range.
 ---
@@ -26,7 +27,8 @@ local validate = vim.validate
 --- local range2 = vim.range(3, 5, 4, 0)
 ---
 --- -- Because `vim.Range` is end exclusive, `range1` and `range2` both represent
---- -- a range starting at the row 3, column 5 and ending at where the row 3 ends.
+--- -- a range starting at the row 3, column 5 and ending at where the row 3 ends
+--- -- (including the newline at the end of line 3).
 ---
 --- -- Operators are overloaded for comparing two `vim.Pos` objects.
 --- if range1 == range2 then
@@ -39,71 +41,161 @@ local validate = vim.validate
 --- need to have the same optional fields.
 ---
 ---@class vim.Range
----@field start vim.Pos Start position.
----@field end_ vim.Pos End position, exclusive.
+---@field start_row integer 0-based byte index.
+---@field start_col integer 0-based byte index.
+---@field end_row integer 0-based byte index.
+---@field end_col integer 0-based byte index.
+---
+--- Optional buffer handle.
+---
+--- When specified, it indicates that this range belongs to a specific buffer.
+--- This field is required when performing range conversions.
+---@field buf? integer
+---@field private [1] integer underlying representation of start_row
+---@field private [2] integer underlying representation of start_col
+---@field private [3] integer underlying representation of end_row
+---@field private [4] integer underlying representation of end_col
 local Range = {}
-Range.__index = Range
+
+---@private
+---@param pos vim.Range
+---@param key any
+function Range.__index(pos, key)
+  if key == 'start_row' then
+    return pos[1]
+  elseif key == 'start_col' then
+    return pos[2]
+  elseif key == 'end_row' then
+    return pos[3]
+  elseif key == 'end_col' then
+    return pos[4]
+  elseif key == 'buf' then
+    return pos[5]
+  end
+
+  return Range[key]
+end
 
 ---@package
 ---@overload fun(self: vim.Range, start: vim.Pos, end_: vim.Pos): vim.Range
 ---@overload fun(self: vim.Range, start_row: integer, start_col: integer, end_row: integer, end_col: integer, opts?: vim.Pos.Optional): vim.Range
 function Range.new(...)
-  ---@type vim.Pos, vim.Pos, vim.Pos.Optional
-  local start, end_
+  ---@type integer, integer, integer, integer, integer|nil
+  local start_row, start_col, end_row, end_col, buf
 
   local nargs = select('#', ...)
   if nargs == 2 then
     ---@type vim.Pos, vim.Pos
-    start, end_ = ...
+    local start, end_ = ...
     validate('start', start, 'table')
     validate('end_', end_, 'table')
 
     if start.buf ~= end_.buf then
       error('start and end positions must belong to the same buffer')
     end
+    start_row, start_col, end_row, end_col, buf =
+      start.row, start.col, end_.row, end_.col, start.buf
   elseif nargs == 4 or nargs == 5 then
-    ---@type integer, integer, integer, integer, vim.Pos.Optional
-    local start_row, start_col, end_row, end_col, opts = ...
-    start, end_ = vim.pos(start_row, start_col, opts), vim.pos(end_row, end_col, opts)
+    local opts
+    ---@type integer, integer, integer, integer, vim.Pos.Optional|nil
+    start_row, start_col, end_row, end_col, opts = ...
+    buf = opts and opts.buf
   else
     error('invalid parameters')
   end
 
   ---@type vim.Range
   local self = setmetatable({
-    start = start,
-    end_ = end_,
+    start_row,
+    start_col,
+    end_row,
+    end_col,
+    buf,
   }, Range)
 
   return self
 end
 
+--- TODO(ofseed): Make it work for unloaded buffers. Check get_line() in vim.lsp.util.
+---@param buf integer
+---@param row integer
+local function get_line(buf, row)
+  return api.nvim_buf_get_lines(buf, row, row + 1, true)[1]
+end
+
+---@param p1_row integer Row of first position to compare.
+---@param p1_col integer Col of first position to compare.
+---@param p2_row integer Row of second position to compare.
+---@param p2_col integer Col of second position to compare.
+---@return integer
+--- 1: a > b
+--- 0: a == b
+--- -1: a < b
+local function cmp_pos(p1_row, p1_col, p2_row, p2_col)
+  if p1_row == p2_row then
+    if p1_col > p2_col then
+      return 1
+    elseif p1_col < p2_col then
+      return -1
+    else
+      return 0
+    end
+  elseif p1_row > p2_row then
+    return 1
+  end
+
+  return -1
+end
+
+---@param row integer
+---@param col integer
+---@param buf integer
+---@return integer, integer
+local function to_inclusive_pos(row, col, buf)
+  if col > 0 then
+    col = col - 1
+  elseif col == 0 and row > 0 then
+    row = row - 1
+    col = #get_line(buf, row)
+  end
+
+  return row, col
+end
+
 ---@private
 ---@param r1 vim.Range
 ---@param r2 vim.Range
 function Range.__lt(r1, r2)
-  return r1.end_ < r2.start
+  local r1_inclusive_end_row, r1_inclusive_end_col =
+    to_inclusive_pos(r1.end_row, r1.end_col, r1.buf)
+  return cmp_pos(r1_inclusive_end_row, r1_inclusive_end_col, r2.start_row, r2.start_col) == -1
 end
 
 ---@private
 ---@param r1 vim.Range
 ---@param r2 vim.Range
 function Range.__le(r1, r2)
-  return r1.end_ <= r2.start
+  local r1_inclusive_end_row, r1_inclusive_end_col =
+    to_inclusive_pos(r1.end_row, r1.end_col, r1.buf)
+  return cmp_pos(r1_inclusive_end_row, r1_inclusive_end_col, r2.start_row, r2.start_col) ~= 1
 end
 
 ---@private
 ---@param r1 vim.Range
 ---@param r2 vim.Range
 function Range.__eq(r1, r2)
-  return r1.start == r2.start and r1.end_ == r2.end_
+  return cmp_pos(r1.start_row, r1.start_col, r2.start_row, r2.start_col) == 0
+    and cmp_pos(r1.end_row, r1.end_col, r2.end_row, r2.end_col) == 0
 end
 
 --- Checks whether the given range is empty; i.e., start >= end.
 ---
 ---@return boolean `true` if the given range is empty
 function Range:is_empty()
-  return self.start >= self.end_
+  local inclusive_end_row, inclusive_end_col =
+    to_inclusive_pos(self.end_row, self.end_col, self.buf)
+
+  return cmp_pos(self.start_row, self.start_col, inclusive_end_row, inclusive_end_col) ~= -1
 end
 
 --- Checks whether {outer} range contains {inner} range or position.
@@ -112,13 +204,30 @@ end
 ---@param inner vim.Range|vim.Pos
 ---@return boolean `true` if {outer} range fully contains {inner} range or position.
 function Range.has(outer, inner)
-  if inner.start then
-    -- inner is a range
-    return outer.start <= inner.start and outer.end_ >= inner.end_
-  else
-    -- inner is a position
-    return outer.start <= inner and outer.end_ >= inner
+  if getmetatable(inner) == vim.pos then
+    ---@cast inner -vim.Range
+    return cmp_pos(outer.start_row, outer.start_col, inner.row, inner.col) ~= 1
+      and cmp_pos(outer.end_row, outer.end_col, inner.row, inner.col) ~= -1
   end
+  ---@cast inner -vim.Pos
+
+  local outer_inclusive_end_row, outer_inclusive_end_col =
+    to_inclusive_pos(outer.end_row, outer.end_col, outer.buf)
+  local inner_inclusive_end_row, inner_inclusive_end_col =
+    to_inclusive_pos(inner.end_row, inner.end_col, inner.buf)
+
+  return cmp_pos(outer.start_row, outer.start_col, inner.start_row, inner.start_col) ~= 1
+    and cmp_pos(outer.end_row, outer.end_col, inner.end_row, inner.end_col) ~= -1
+    -- accounts for empty ranges at the start/end of `outer` that per Neovim API and LSP logic
+    -- insert the text outside `outer`
+    and cmp_pos(outer.start_row, outer.start_col, inner_inclusive_end_row, inner_inclusive_end_col) == -1
+    and cmp_pos(
+        outer_inclusive_end_row,
+        outer_inclusive_end_col,
+        inner.start_row,
+        inner.start_col
+      )
+      == 1
 end
 
 --- Computes the common range shared by the given ranges.
@@ -128,12 +237,21 @@ end
 ---@return vim.Range? range that is present inside both `r1` and `r2`.
 ---                   `nil` if such range does not exist.
 function Range.intersect(r1, r2)
-  if r1.end_ <= r2.start or r1.start >= r2.end_ then
+  local r1_inclusive_end_row, r1_inclusive_end_col =
+    to_inclusive_pos(r1.end_row, r1.end_col, r1.buf)
+  local r2_inclusive_end_row, r2_inclusive_end_col =
+    to_inclusive_pos(r2.end_row, r2.end_col, r2.buf)
+
+  if
+    cmp_pos(r1_inclusive_end_row, r1_inclusive_end_col, r2.start_row, r2.start_col) ~= 1
+    or cmp_pos(r1.start_row, r1.start_col, r2_inclusive_end_row, r2_inclusive_end_col) ~= -1
+  then
     return nil
   end
-  local rs = r1.start <= r2.start and r2 or r1
-  local re = r1.end_ >= r2.end_ and r2 or r1
-  return Range.new(rs.start, re.end_)
+
+  local rs = cmp_pos(r1.start_row, r1.start_col, r2.start_row, r2.start_col) ~= 1 and r2 or r1
+  local re = cmp_pos(r1.end_row, r1.end_col, r2.end_row, r2.end_col) ~= -1 and r2 or r1
+  return Range.new(rs.start_row, rs.start_col, re.end_row, re.end_col)
 end
 
 --- Converts |vim.Range| to `lsp.Range`.
@@ -156,8 +274,10 @@ function Range.to_lsp(range, position_encoding)
 
   ---@type lsp.Range
   return {
-    ['start'] = range.start:to_lsp(position_encoding),
-    ['end'] = range.end_:to_lsp(position_encoding),
+    ['start'] = vim
+      .pos(range.start_row, range.start_col, { buf = range.buf })
+      :to_lsp(position_encoding),
+    ['end'] = vim.pos(range.end_row, range.end_col, { buf = range.buf }):to_lsp(position_encoding),
   }
 end
 
@@ -190,6 +310,97 @@ function Range.lsp(buf, range, position_encoding)
   return Range.new(start, end_)
 end
 
+--- Converts |vim.Range| to extmark range (see |api-indexing|).
+---
+--- Example:
+--- ```lua
+--- -- `buf` is required for conversion to extmark range.
+--- local buf = vim.api.nvim_get_current_buf()
+--- local range = vim.range(3, 5, 4, 0, { buf = buf })
+---
+--- -- Convert to extmark range, you can call it in a method style.
+--- local extmark_range = range:to_extmark()
+--- ```
+---@param range vim.Range
+function Range.to_extmark(range)
+  validate('range', range, 'table')
+
+  local srow, scol = vim.pos(range.start_row, range.start_col, { buf = range.buf }):to_extmark()
+  local erow, ecol = vim.pos(range.end_row, range.end_col, { buf = range.buf }):to_extmark()
+  return srow, scol, erow, ecol
+end
+
+--- Creates a new |vim.Range| from extmark range (see |api-indexing|).
+---
+--- Example:
+--- ```lua
+--- local buf = vim.api.nvim_get_current_buf()
+---
+--- local range = vim.range.extmark(3, 5, 4, 0, { buf = buf })
+--- ```
+---@param start_row integer
+---@param start_col integer
+---@param end_row integer
+---@param end_col integer
+---@param opts vim.Pos.Optional|nil
+function Range.extmark(start_row, start_col, end_row, end_col, opts)
+  validate('range', start_row, 'number')
+  validate('range', start_col, 'number')
+  validate('range', end_row, 'number')
+  validate('range', end_col, 'number')
+
+  local start = vim.pos.extmark(start_row, start_col, opts)
+  local end_ = vim.pos.extmark(end_row, end_col, opts)
+
+  return Range.new(start, end_)
+end
+
+--- Converts |vim.Range| to mark-like range (see |api-indexing|).
+---
+--- Example:
+--- ```lua
+--- -- `buf` is required for conversion to extmark range.
+--- local buf = vim.api.nvim_get_current_buf()
+--- local range = vim.range(3, 5, 4, 0, { buf = buf })
+---
+--- -- Convert to cursor range, you can call it in a method style.
+--- local cursor_range = range:to_cursor()
+--- ```
+---@param range vim.Range
+function Range.to_cursor(range)
+  validate('range', range, 'table')
+
+  local srow, scol = vim.pos(range.start_row, range.start_col, { buf = range.buf }):to_cursor()
+  local erow, ecol = vim.pos(range.end_row, range.end_col, { buf = range.buf }):to_cursor()
+  return srow, scol, erow, ecol
+end
+
+--- Creates a new |vim.Range| from mark-like range (see |api-indexing|).
+---
+--- Example:
+--- ```lua
+--- local buf = vim.api.nvim_get_current_buf()
+--- local start = vim.api.nvim_win_get_cursor(0)
+--- -- move the cursor
+--- local end_ = vim.api.nvim_win_get_cursor(0)
+---
+--- local range = vim.range.cursor(start, end_, { buf = buf })
+--- ```
+---@param buf integer
+---@param start_pos [integer, integer]
+---@param end_pos [integer, integer]
+---@param opts vim.Pos.Optional|nil
+function Range.cursor(buf, start_pos, end_pos, opts)
+  validate('buf', buf, 'number')
+  validate('range', start_pos, 'table')
+  validate('range', end_pos, 'table')
+
+  local start = vim.pos.cursor(start_pos, opts)
+  local end_ = vim.pos.cursor(end_pos, opts)
+
+  return Range.new(start, end_)
+end
+
 -- Overload `Range.new` to allow calling this module as a function.
 setmetatable(Range, {
   __call = function(_, ...)