u.nvim/lua/u/pos.lua

local MAX_COL = vim.v.maxcol

---@param buf number
---@param lnum number 1-based
local function line_text(buf, lnum) return vim.api.nvim_buf_get_lines(buf, lnum - 1, lnum, false)[1] end

---@class u.Pos
---@field buf number buffer number
---@field lnum number 1-based line index
---@field col number 1-based column index
---@field off number
local Pos = {}
Pos.MAX_COL = MAX_COL

---@param buf? number
---@param lnum number 1-based
---@param col number 1-based
---@param off? number
---@return u.Pos
function Pos.new(buf, lnum, col, off)
  if buf == nil or buf == 0 then buf = vim.api.nvim_get_current_buf() end
  if off == nil then off = 0 end
  local pos = {
    buf = buf,
    lnum = lnum,
    col = col,
    off = off,
  }

  local function str()
    if pos.off ~= 0 then
      return string.format('Pos(%d:%d){buf=%d, off=%d}', pos.lnum, pos.col, pos.buf, pos.off)
    else
      return string.format('Pos(%d:%d){buf=%d}', pos.lnum, pos.col, pos.buf)
    end
  end
  setmetatable(pos, {
    __index = Pos,
    __tostring = str,
    __lt = Pos.__lt,
    __le = Pos.__le,
    __eq = Pos.__eq,
  })
  return pos
end

function Pos.invalid() return Pos.new(0, 0, 0, 0) end

function Pos.is(x)
  if not type(x) == 'table' then return false end
  local mt = getmetatable(x)
  return mt and mt.__index == Pos
end

function Pos.__lt(a, b) return a.lnum < b.lnum or (a.lnum == b.lnum and a.col < b.col) end
function Pos.__le(a, b) return a < b or a == b end
function Pos.__eq(a, b) return Pos.is(a) and Pos.is(b) and a.buf == b.buf and a.lnum == b.lnum and a.col == b.col end
function Pos.__add(x, y)
  if type(x) == 'number' then
    x, y = y, x
  end
  if not Pos.is(x) or type(y) ~= 'number' then return nil end
  return x:next(y)
end
function Pos.__sub(x, y)
  if type(x) == 'number' then
    x, y = y, x
  end
  if not Pos.is(x) or type(y) ~= 'number' then return nil end
  return x:next(-y)
end

---@param name string
---@return u.Pos
function Pos.from_pos(name)
  local p = vim.fn.getpos(name)
  return Pos.new(p[1], p[2], p[3], p[4])
end

function Pos:is_invalid() return self.buf == 0 and self.lnum == 0 and self.col == 0 and self.off == 0 end

function Pos:clone() return Pos.new(self.buf, self.lnum, self.col, self.off) end

---@return boolean
function Pos:is_col_max() return self.col == MAX_COL end

--- Normalize the position to a real position (take into account vim.v.maxcol).
function Pos:as_real()
  local maxlen = #line_text(self.buf, self.lnum)
  local col = self.col
  if col > maxlen then
    -- We could use utilities in this file to get the given line, but
    -- since this is a low-level function, we are going to optimize and
    -- use the API directly:
    col = maxlen
  end
  return Pos.new(self.buf, self.lnum, col, self.off)
end

function Pos:as_vim() return { self.buf, self.lnum, self.col, self.off } end

---@param pos string
function Pos:save_to_pos(pos) vim.fn.setpos(pos, { self.buf, self.lnum, self.col, self.off }) end

---@param mark string
function Pos:save_to_mark(mark)
  local p = self:as_real()
  vim.api.nvim_buf_set_mark(p.buf, mark, p.lnum, p.col, {})
end

---@return string
function Pos:char()
  local line = line_text(self.buf, self.lnum)
  if line == nil then return '' end
  return line:sub(self.col, self.col)
end

function Pos:line() return line_text(self.buf, self.lnum) end

---@param dir? -1|1
---@param must? boolean
---@return u.Pos|nil
function Pos:next(dir, must)
  if must == nil then must = false end

  if dir == nil or dir == 1 then
    -- Next:
    local num_lines = vim.api.nvim_buf_line_count(self.buf)
    local last_line = line_text(self.buf, num_lines)
    if self.lnum == num_lines and self.col == #last_line then
      if must then error 'error in Pos:next(): Pos:next() returned nil' end
      return nil
    end

    local col = self.col + 1
    local line = self.lnum
    local line_max_col = #line_text(self.buf, self.lnum)
    if col > line_max_col then
      col = 1
      line = line + 1
    end
    return Pos.new(self.buf, line, col, self.off)
  else
    -- Previous:
    if self.col == 1 and self.lnum == 1 then
      if must then error 'error in Pos:next(): Pos:next() returned nil' end
      return nil
    end

    local col = self.col - 1
    local line = self.lnum
    local prev_line_max_col = #(line_text(self.buf, self.lnum - 1) or '')
    if col < 1 then
      col = math.max(prev_line_max_col, 1)
      line = line - 1
    end
    return Pos.new(self.buf, line, col, self.off)
  end
end

---@param dir? -1|1
function Pos:must_next(dir)
  local next = self:next(dir, true)
  if next == nil then error 'unreachable' end
  return next
end

---@param dir -1|1
---@param predicate fun(p: u.Pos): boolean
---@param test_current? boolean
function Pos:next_while(dir, predicate, test_current)
  if test_current and not predicate(self) then return end
  local curr = self
  while true do
    local next = curr:next(dir)
    if next == nil or not predicate(next) then break end
    curr = next
  end
  return curr
end

---@param dir -1|1
---@param predicate string|fun(p: u.Pos): boolean
function Pos:find_next(dir, predicate)
  if type(predicate) == 'string' then
    local s = predicate
    predicate = function(p) return s == p:char() end
  end

  ---@type u.Pos|nil
  local curr = self
  while curr ~= nil do
    if predicate(curr) then return curr end
    curr = curr:next(dir)
  end
  return curr
end

--- Finds the matching bracket/paren for the current position.
---@param max_chars? number|nil
---@param invocations? u.Pos[]
---@return u.Pos|nil
function Pos:find_match(max_chars, invocations)
  if invocations == nil then invocations = {} end
  if vim.tbl_contains(invocations, function(p) return self == p end, { predicate = true }) then return nil end
  table.insert(invocations, self)

  local openers = { '{', '[', '(', '<' }
  local closers = { '}', ']', ')', '>' }
  local c = self:char()
  local is_opener = vim.tbl_contains(openers, c)
  local is_closer = vim.tbl_contains(closers, c)
  if not is_opener and not is_closer then return nil end

  local i, _ = vim.iter(is_opener and openers or closers):enumerate():find(function(_, c2) return c == c2 end)
  -- Store the character we will be looking for:
  local c_match = (is_opener and closers or openers)[i]

  ---@type u.Pos|nil
  local cur = self
  --- `adv` is a helper that moves the current position backward or forward,
  --- depending on whether we are looking for an opener or a closer. It returns
  --- nil if 1) the watch-dog `max_chars` falls bellow 0, or 2) if we have gone
  --- beyond the beginning/end of the file.
  ---@return u.Pos|nil
  local function adv()
    if cur == nil then return nil end

    if max_chars ~= nil then
      max_chars = max_chars - 1
      if max_chars < 0 then return nil end
    end

    return cur:next(is_opener and 1 or -1)
  end

  -- scan until we find `c_match`:
  cur = adv()
  while cur ~= nil and cur:char() ~= c_match do
    cur = adv()
    if cur == nil then break end

    local c2 = cur:char()
    if c2 == c_match then break end

    if vim.tbl_contains(openers, c2) or vim.tbl_contains(closers, c2) then
      cur = cur:find_match(max_chars, invocations)
      cur = adv() -- move past the match
    end
  end

  return cur
end

return Pos