u.nvim/lua/u/range.lua

local Pos = require 'u.pos'

local orig_on_yank = (vim.hl or vim.highlight).on_yank
local on_yank_enabled = true;
((vim.hl or vim.highlight) --[[@as any]]).on_yank = function(opts)
  if not on_yank_enabled then return end
  return orig_on_yank(opts)
end

---@class u.Range
---@field start u.Pos
---@field stop u.Pos|nil
---@field mode 'v'|'V'
local Range = {}

---@param start u.Pos
---@param stop u.Pos|nil
---@param mode? 'v'|'V'
---@return u.Range
function Range.new(start, stop, mode)
  if stop ~= nil and stop < start then
    start, stop = stop, start
  end

  local r = { start = start, stop = stop, mode = mode or 'v' }
  local function str()
    ---@param p u.Pos
    local function posstr(p)
      if p == nil then
        return 'nil'
      elseif p.off ~= 0 then
        return string.format('Pos(%d:%d){off=%d}', p.lnum, p.col, p.off)
      else
        return string.format('Pos(%d:%d)', p.lnum, p.col)
      end
    end

    local _1 = posstr(r.start)
    local _2 = posstr(r.stop)
    return string.format('Range{buf=%d, mode=%s, start=%s, stop=%s}', r.start.buf, r.mode, _1, _2)
  end
  setmetatable(r, { __index = Range, __tostring = str })
  return r
end

function Range.is(x)
  local mt = getmetatable(x)
  return mt and mt.__index == Range
end

---@param lpos string
---@param rpos string
---@return u.Range
function Range.from_marks(lpos, rpos)
  local start = Pos.from_pos(lpos)
  local stop = Pos.from_pos(rpos)

  ---@type 'v'|'V'
  local mode
  if stop:is_col_max() then
    mode = 'V'
  else
    mode = 'v'
  end

  return Range.new(start, stop, mode)
end

---@param buf? number
function Range.from_buf_text(buf)
  if buf == nil or buf == 0 then buf = vim.api.nvim_get_current_buf() end
  local num_lines = vim.api.nvim_buf_line_count(buf)

  local start = Pos.new(buf, 1, 1)
  local stop = Pos.new(buf, num_lines, Pos.MAX_COL)
  return Range.new(start, stop, 'V')
end

-- TODO: make 1-based
---@param buf? number
---@param line number 0-based line index
function Range.from_line(buf, line) return Range.from_lines(buf, line, line) end

-- TODO: make 1-based
---@param buf? number
---@param start_line number 0-based line index
---@param stop_line number 0-based line index
function Range.from_lines(buf, start_line, stop_line)
  if buf == nil or buf == 0 then buf = vim.api.nvim_get_current_buf() end
  if stop_line < 0 then
    local num_lines = vim.api.nvim_buf_line_count(buf)
    stop_line = num_lines + stop_line
  end
  return Range.new(Pos.new(buf, start_line + 1, 1), Pos.new(buf, stop_line + 1, Pos.MAX_COL), 'V')
end

---@param text_obj string
---@param opts? { buf?: number; contains_cursor?: boolean; pos?: u.Pos, user_defined?: boolean }
---@return u.Range|nil
function Range.from_text_object(text_obj, opts)
  opts = opts or {}
  if opts.buf == nil then opts.buf = vim.api.nvim_get_current_buf() end
  if opts.contains_cursor == nil then opts.contains_cursor = false end
  if opts.user_defined == nil then opts.user_defined = false end

  ---@type "a" | "i"
  local selection_type = text_obj:sub(1, 1)
  local obj_type = text_obj:sub(#text_obj, #text_obj)
  local is_quote = vim.tbl_contains({ "'", '"', '`' }, obj_type)
  local cursor = Pos.from_pos '.'

  --- @type u.Pos
  local start
  --- @type u.Pos
  local stop

  vim.api.nvim_buf_call(opts.buf, function()
    local original_state = {
      winview = vim.fn.winsaveview(),
      regquote = vim.fn.getreg '"',
      posdot = vim.fn.getpos '.',
      poslb = vim.fn.getpos "'[",
      posrb = vim.fn.getpos "']",
    }

    if opts.pos ~= nil then opts.pos:save_to_pos '.' end

    Pos.invalid():save_to_pos "'["
    Pos.invalid():save_to_pos "']"

    local prev_on_yank_enabled = on_yank_enabled
    on_yank_enabled = false
    vim.cmd {
      cmd = 'normal',
      bang = not opts.user_defined,
      args = { '""y' .. text_obj },
      mods = { silent = true },
    }
    on_yank_enabled = prev_on_yank_enabled

    start = Pos.from_pos "'["
    stop = Pos.from_pos "']"

    -- Restore original state:
    vim.fn.winrestview(original_state.winview)
    vim.fn.setreg('"', original_state.regquote)
    vim.fn.setpos('.', original_state.posdot)
    vim.fn.setpos("'[", original_state.poslb)
    vim.fn.setpos("']", original_state.posrb)

    if
      -- I have no idea why, but when yanking `i"`, the stop-mark is
      -- placed on the ending quote. For other text-objects, the stop-
      -- mark is placed before the closing character.
      (is_quote and selection_type == 'i' and stop:char() == obj_type)
      -- *Sigh*, this also sometimes happens for `it` as well.
      or (text_obj == 'it' and stop:char() == '<')
    then
      stop = stop:next(-1) or stop
    end
  end)

  if start == stop and start:is_invalid() then return nil end
  if opts.contains_cursor and not Range.new(start, stop):contains(cursor) then return nil end

  if is_quote and selection_type == 'a' then
    start = start:find_next(1, obj_type) or start
    stop = stop:find_next(-1, obj_type) or stop
  end

  return Range.new(start, stop)
end

--- Get range information from the currently selected visual text.
--- Note: from within a command mapping or an opfunc, use other specialized
--- utilities, such as:
--- * Range.from_cmd_args
--- * Range.from_op_func
function Range.from_vtext()
  local r = Range.from_marks('v', '.')
  if vim.fn.mode() == 'V' then r = r:to_linewise() end
  return r
end

--- Get range information from the current text range being operated on
--- as defined by an operator-pending function. Infers line-wise vs. char-wise
--- based on the type, as given by the operator-pending function.
---@param type 'line'|'char'|'block'
function Range.from_op_func(type)
  if type == 'block' then error 'block motions not supported' end

  local range = Range.from_marks("'[", "']")
  if type == 'line' then range = range:to_linewise() end
  return range
end

--- Get range information from command arguments.
---@param args unknown
---@return u.Range|nil
function Range.from_cmd_args(args)
  ---@type 'v'|'V'
  local mode
  ---@type nil|u.Pos
  local start
  local stop
  if args.range == 0 then
    return nil
  else
    start = Pos.from_pos "'<"
    stop = Pos.from_pos "'>"
    mode = stop:is_col_max() and 'V' or 'v'
  end
  return Range.new(start, stop, mode)
end

---
function Range.find_nearest_brackets()
  return Range.smallest {
    Range.from_text_object('a<', { contains_cursor = true }),
    Range.from_text_object('a[', { contains_cursor = true }),
    Range.from_text_object('a(', { contains_cursor = true }),
    Range.from_text_object('a{', { contains_cursor = true }),
  }
end

function Range.find_nearest_quotes()
  return Range.smallest {
    Range.from_text_object([[a']], { contains_cursor = true }),
    Range.from_text_object([[a"]], { contains_cursor = true }),
    Range.from_text_object([[a`]], { contains_cursor = true }),
  }
end

---@param ranges (u.Range|nil)[]
function Range.smallest(ranges)
  ---@type u.Range[]
  ranges = vim.iter(ranges):filter(function(r) return r ~= nil and not r:is_empty() end):totable()
  if #ranges == 0 then return nil end

  -- find smallest match
  local smallest = ranges[1]
  for _, r in ipairs(ranges) do
    local start, stop = r.start, r.stop
    if start > smallest.start and stop < smallest.stop then smallest = r end
  end
  return smallest
end

function Range:clone() return Range.new(self.start:clone(), self.stop ~= nil and self.stop:clone() or nil, self.mode) end
function Range:line_count()
  if self:is_empty() then return 0 end
  return self.stop.lnum - self.start.lnum + 1
end

function Range:to_linewise()
  local r = self:clone()

  r.mode = 'V'
  r.start.col = 1
  if r.stop ~= nil then r.stop.col = Pos.MAX_COL end

  return r
end

function Range:is_empty() return self.stop == nil end

function Range:trim_start()
  if self:is_empty() then return end

  local r = self:clone()
  while r.start:char():match '%s' do
    local next = r.start:next(1)
    if next == nil then break end
    r.start = next
  end
  return r
end

function Range:trim_stop()
  if self:is_empty() then return end

  local r = self:clone()
  while r.stop:char():match '%s' do
    local next = r.stop:next(-1)
    if next == nil then break end
    r.stop = next
  end
  return r
end

---@param p u.Pos
function Range:contains(p) return not self:is_empty() and p >= self.start and p <= self.stop end

---@return string[]
function Range:lines()
  if self:is_empty() then return {} end
  return vim.fn.getregion(self.start:as_vim(), self.stop:as_vim(), { type = self.mode })
end

---@return string
function Range:text() return vim.fn.join(self:lines(), '\n') end

---@param i number 1-based
---@param j? number 1-based
function Range:sub(i, j) return self:text():sub(i, j) end

---@param l number
---@return { line: string; idx0: { start: number; stop: number; }; lnum: number; range: fun():u.Range; text: fun():string }|nil
function Range:line(l)
  if l < 0 then l = self:line_count() + l + 1 end
  if l > self:line_count() then return end

  local line_indices = vim.fn.getregionpos(self.start:as_vim(), self.stop:as_vim(), { type = self.mode })
  local line_bounds = line_indices[l]

  local start = Pos.new(unpack(line_bounds[1]))
  local stop = Pos.new(unpack(line_bounds[2]))
  return Range.new(start, stop)
end

---@param replacement nil|string|string[]
function Range:replace(replacement)
  if replacement == nil then replacement = {} end
  if type(replacement) == 'string' then replacement = vim.fn.split(replacement, '\n') end

  local buf = self.start.buf
  local replace_type = (self:is_empty() and 'insert') or (self.mode == 'v' and 'region') or 'lines'

  local function update_stop_non_linewise()
    local new_last_line_num = self.start.lnum + #replacement - 1
    local new_last_col = #(replacement[#replacement] or '')
    if new_last_line_num == self.start.lnum then new_last_col = new_last_col + self.start.col - 1 end
    self.stop = Pos.new(buf, new_last_line_num, new_last_col)
  end
  local function update_stop_linewise()
    if #replacement == 0 then
      self.stop = nil
    else
      local new_last_line_num = self.start.lnum - 1 + #replacement - 1
      self.stop = Pos.new(buf, new_last_line_num + 1, Pos.MAX_COL, self.stop.off)
    end
    self.mode = 'v'
  end

  if replace_type == 'insert' then
    -- To insert text at a given `(row, column)` location, use `start_row =
    -- end_row = row` and `start_col = end_col = col`.
    vim.api.nvim_buf_set_text(
      buf,
      self.start.lnum - 1,
      self.start.col - 1,
      self.start.lnum - 1,
      self.start.col - 1,
      replacement
    )
    update_stop_non_linewise()
  elseif replace_type == 'region' then
    -- Fixup the bounds:
    local max_col = #self.stop:line()

    -- Indexing is zero-based. Row indices are end-inclusive, and column indices
    -- are end-exclusive.
    vim.api.nvim_buf_set_text(
      buf,
      self.start.lnum - 1,
      self.start.col - 1,
      self.stop.lnum - 1,
      math.min(self.stop.col, max_col),
      replacement
    )
    update_stop_non_linewise()
  elseif replace_type == 'lines' then
    -- Indexing is zero-based, end-exclusive.
    vim.api.nvim_buf_set_lines(buf, self.start.lnum - 1, self.stop.lnum, true, replacement)
    update_stop_linewise()
  else
    error 'unreachable'
  end
end

---@param amount number
function Range:shrink(amount)
  local start = self.start
  local stop = self.stop
  if stop == nil then return self:clone() end

  for _ = 1, amount do
    local next_start = start:next(1)
    local next_stop = stop:next(-1)
    if next_start == nil or next_stop == nil then return end
    start = next_start
    stop = next_stop
    if next_start > next_stop then break end
  end
  if start > stop then stop = nil end
  return Range.new(start, stop, self.mode)
end

---@param amount number
function Range:must_shrink(amount)
  local shrunk = self:shrink(amount)
  if shrunk == nil or shrunk:is_empty() then error 'error in Range:must_shrink: Range:shrink() returned nil' end
  return shrunk
end

---@param left string
---@param right string
function Range:save_to_pos(left, right)
  if self:is_empty() then
    self.start:save_to_pos(left)
    self.start:save_to_pos(right)
  else
    self.start:save_to_pos(left)
    self.stop:save_to_pos(right)
  end
end

---@param left string
---@param right string
function Range:save_to_marks(left, right)
  if self:is_empty() then
    self.start:save_to_mark(left)
    self.start:save_to_mark(right)
  else
    self.start:save_to_mark(left)
    self.stop:save_to_mark(right)
  end
end

function Range:set_visual_selection()
  if self:is_empty() then return end

  if vim.api.nvim_get_current_buf() ~= self.start.buf then vim.api.nvim_set_current_buf(self.start.buf) end

  local curr_mode = vim.fn.mode()
  if curr_mode ~= self.mode then vim.fn.feedkeys(self.mode, 'x') end
  self.start:save_to_pos '.'
  vim.fn.feedkeys('o', 'x')
  self.stop:save_to_pos '.'
end

---@param group string
---@param opts? { timeout?: number, priority?: number, on_macro?: boolean }
function Range:highlight(group, opts)
  if self:is_empty() then return end

  opts = opts or { on_macro = false }
  if opts.on_macro == nil then opts.on_macro = false end

  local in_macro = vim.fn.reg_executing() ~= ''
  if not opts.on_macro and in_macro then return { clear = function() end } end

  local ns = vim.api.nvim_create_namespace ''

  local winview = vim.fn.winsaveview();
  (vim.hl or vim.highlight).range(
    self.start.buf,
    ns,
    group,
    { self.start.lnum - 1, self.start.col - 1 },
    { self.stop.lnum - 1, self.stop.col - 1 },
    {
      inclusive = true,
      priority = opts.priority,
      timeout = opts.timeout,
      regtype = self.mode,
    }
  )
  if not in_macro then vim.fn.winrestview(winview) end
  vim.cmd.redraw()

  return {
    ns = ns,
    clear = function()
      vim.api.nvim_buf_clear_namespace(self.start.buf, ns, self.start.lnum - 1, self.stop.lnum)
      vim.cmd.redraw()
    end,
  }
end

return Range