snacks.nvim/lua/snacks/indent.lua

---@class snacks.indent
local M = {}

M.meta = {
  desc = "Indent guides and scopes",
}

M.enabled = false

---@class snacks.indent.Config
---@field enabled? boolean
local defaults = {
  indent = {
    enabled = true, -- enable indent guides
    char = "│",
    blank = " ",
    -- blank = "∙",
    only_scope = false, -- only show indent guides of the scope
    only_current = false, -- only show indent guides in the current window
    hl = "SnacksIndent", ---@type string|string[] hl groups for indent guides
    -- can be a list of hl groups to cycle through
    -- hl = {
    --     "SnacksIndent1",
    --     "SnacksIndent2",
    --     "SnacksIndent3",
    --     "SnacksIndent4",
    --     "SnacksIndent5",
    --     "SnacksIndent6",
    --     "SnacksIndent7",
    --     "SnacksIndent8",
    -- },
  },
  -- animate scopes. Enabled by default for Neovim >= 0.10
  -- Works on older versions but has to trigger redraws during animation.
  ---@type snacks.animate.Config|{enabled?: boolean}
  animate = {
    enabled = vim.fn.has("nvim-0.10") == 1,
    easing = "linear",
    duration = {
      step = 20, -- ms per step
      total = 500, -- maximum duration
    },
  },
  ---@class snacks.indent.Scope.Config: snacks.scope.Config
  scope = {
    enabled = true, -- enable highlighting the current scope
    char = "│",
    underline = false, -- underline the start of the scope
    only_current = false, -- only show scope in the current window
    hl = "SnacksIndentScope", ---@type string|string[] hl group for scopes
  },
  chunk = {
    -- when enabled, scopes will be rendered as chunks, except for the
    -- top-level scope which will be rendered as a scope.
    enabled = false,
    -- only show chunk scopes in the current window
    only_current = false,
    hl = "SnacksIndentChunk", ---@type string|string[] hl group for chunk scopes
    char = {
      corner_top = "┌",
      corner_bottom = "└",
      -- corner_top = "╭",
      -- corner_bottom = "╰",
      horizontal = "─",
      vertical = "│",
      arrow = ">",
    },
  },
  blank = {
    char = " ",
    -- char = "·",
    hl = "SnacksIndentBlank", ---@type string|string[] hl group for blank spaces
  },
  -- filter for buffers to enable indent guides
  filter = function(buf)
    return vim.g.snacks_indent ~= false and vim.b[buf].snacks_indent ~= false and vim.bo[buf].buftype == ""
  end,
  priority = 200,
  debug = false,
}

---@class snacks.indent.Scope: snacks.scope.Scope
---@field win number
---@field step? number

local config = Snacks.config.get("scope", defaults)
local ns = vim.api.nvim_create_namespace("snacks_indent")
local cache_extmarks = {} ---@type table<string, vim.api.keyset.set_extmark|false>
local debug_timer = assert((vim.uv or vim.loop).new_timer())
local cache_underline = {} ---@type table<string, boolean>
local states = {} ---@type table<number, snacks.indent.State>
local scopes ---@type snacks.scope.Listener?
local stats = {
  indents = 0,
  extmarks = 0,
  scope = 0,
}

Snacks.util.set_hl({
  [""] = "NonText",
  Blank = "SnacksIndent",
  Scope = "Special",
  Chunk = "SnacksIndentScope",
  ["1"] = "DiagnosticInfo",
  ["2"] = "DiagnosticHint",
  ["3"] = "DiagnosticWarn",
  ["4"] = "DiagnosticError",
  ["5"] = "DiagnosticInfo",
  ["6"] = "DiagnosticHint",
  ["7"] = "DiagnosticWarn",
  ["8"] = "DiagnosticError",
}, { prefix = "SnacksIndent", default = true })

---@param level number
---@param hl string|string[]
local function get_hl(level, hl)
  return type(hl) == "string" and hl or hl[(level - 1) % #hl + 1]
end

---@param hl string
local function get_underline_hl(hl)
  local ret = "SnacksIndentUnderline_" .. hl
  if not cache_underline[hl] then
    local fg = Snacks.util.color(hl, "fg")
    vim.api.nvim_set_hl(0, ret, { sp = fg, underline = true })
    cache_underline[hl] = true
  end
  return ret
end

--- Get the virtual text for the indent guide with
--- the given indent level, left column and shiftwidth
---@param indent number
---@param state snacks.indent.State
local function get_extmark(indent, state)
  local key = indent .. ":" .. state.leftcol .. ":" .. state.shiftwidth
  if cache_extmarks[key] ~= nil then
    return cache_extmarks[key]
  end
  stats.extmarks = stats.extmarks + 1

  local sw = state.shiftwidth
  indent = math.floor(indent / sw) * sw -- align to shiftwidth
  indent = indent - state.leftcol -- adjust for visible indents
  local rem = indent % sw -- remaining spaces of the first partially visible indent
  indent = math.floor(indent / sw) -- full visible indents

  -- hide if indent is 0 and no remaining spaces
  if indent < 1 and rem == 0 then
    cache_extmarks[key] = false
    return false
  end

  local hidden = math.ceil(state.leftcol / sw) -- level of the last hidden indent
  local blank = config.indent.blank:rep(sw - vim.api.nvim_strwidth(config.indent.char))

  local text = {} ---@type string[][]
  text[1] = rem > 0 and { (config.indent.blank):rep(rem), get_hl(hidden, config.blank.hl) } or nil

  for i = 1, indent do
    text[#text + 1] = { config.indent.char, get_hl(i + hidden, config.indent.hl) }
    text[#text + 1] = { blank, get_hl(i + hidden, config.blank.hl) }
  end

  cache_extmarks[key] = {
    virt_text = text,
    virt_text_pos = "overlay",
    hl_mode = "combine",
    priority = config.priority,
    ephemeral = true,
  }
  return cache_extmarks[key]
end

local function animating(buf)
  return Snacks.animate.enabled({ buf = buf, name = "indent" })
end

---@param win number
---@param buf number
---@param top number
---@param bottom number
local function get_state(win, buf, top, bottom)
  local prev, changedtick = states[win], vim.b[buf].changedtick ---@type snacks.indent.State?, number
  if not (prev and prev.buf == buf and prev.changedtick == changedtick) then
    prev = nil
  end
  ---@class snacks.indent.State
  ---@field indents table<number, number>
  local state = {
    win = win,
    buf = buf,
    changedtick = changedtick,
    is_current = win == vim.api.nvim_get_current_win(),
    top = top,
    bottom = bottom,
    leftcol = vim.api.nvim_buf_call(buf, vim.fn.winsaveview).leftcol --[[@as number]],
    shiftwidth = vim.bo[buf].shiftwidth,
    indents = prev and prev.indents or { [0] = 0 },
  }
  state.shiftwidth = state.shiftwidth == 0 and vim.bo[buf].tabstop or state.shiftwidth
  states[win] = state
  return state
end

--- Called during every redraw cycle, so it should be fast.
--- Everything that can be cached should be cached.
---@param win number
---@param buf number
---@param top number -- 1-indexed
---@param bottom number -- 1-indexed
---@private
function M.on_win(win, buf, top, bottom)
  local state = get_state(win, buf, top, bottom)

  local scope = scopes and scopes:get(win) --[[@as snacks.indent.Scope?]]
  local indent_col = 0 -- the start column of the indent guides

  -- adjust top and bottom if only_scope is enabled
  if config.indent.only_scope then
    if not scope then
      return
    end
    indent_col = scope.indent or 0
    state.top = math.max(state.top, scope.from)
    state.bottom = math.min(state.bottom, scope.to)
  end

  local show_indent = config.indent.enabled and (not config.indent.only_current or state.is_current)
  local show_scope = config.scope.enabled and (not config.scope.only_current or state.is_current)
  local show_chunk = config.chunk.enabled and (not config.chunk.only_current or state.is_current)

  -- Calculate and render indents
  local indents = state.indents
  vim.api.nvim_buf_call(buf, function()
    for l = top, bottom do
      local indent = indents[l]
      if not indent then
        stats.indents = stats.indents + 1
        local next = vim.fn.nextnonblank(l)
        -- Indent for a blank line is the minimum of the previous and next non-blank line.
        -- If the previous and next non-blank lines have different indents, add shiftwidth.
        if next ~= l then
          local prev = vim.fn.prevnonblank(l)
          indents[prev] = indents[prev] or vim.fn.indent(prev)
          indents[next] = indents[next] or vim.fn.indent(next)
          indent = math.min(indents[prev], indents[next])
          if indents[prev] ~= indents[next] then
            indent = indent + state.shiftwidth
          end
        else
          indent = vim.fn.indent(l)
        end
        indents[l] = indent
      end
      local opts = show_indent and indent > 0 and get_extmark(indent - indent_col, state)
      if opts then
        vim.api.nvim_buf_set_extmark(buf, ns, l - 1, indent_col, opts)
      end
    end
  end)

  -- Render scope
  if scope and scope:size() > 1 then
    show_chunk = show_chunk and (scope.indent or 0) >= state.shiftwidth
    if show_chunk then
      M.render_chunk(scope, state)
    elseif show_scope then
      M.render_scope(scope, state)
    end
  end
end

--- Render the scope overlappping the given range
---@param scope snacks.indent.Scope
---@param state snacks.indent.State
---@private
function M.render_scope(scope, state)
  local indent = (scope.indent or 2)
  local hl = get_hl(scope.indent + 1, config.scope.hl)
  local to = animating(scope.buf) and scope.step or scope.to
  local col = indent - state.leftcol

  if config.scope.underline and scope.from >= state.top and scope.from <= state.bottom then
    vim.api.nvim_buf_set_extmark(scope.buf, ns, scope.from - 1, math.max(col, 0), {
      end_col = #vim.api.nvim_buf_get_lines(scope.buf, scope.from - 1, scope.from, false)[1],
      hl_group = get_underline_hl(hl),
      hl_mode = "combine",
      priority = config.priority + 1,
      strict = false,
      ephemeral = true,
    })
  end

  if col < 0 then -- scope is hidden
    return
  end

  for l = math.max(scope.from, state.top), math.min(to, state.bottom) do
    local i = state.indents[l]
    if i and i > indent then
      vim.api.nvim_buf_set_extmark(scope.buf, ns, l - 1, 0, {
        virt_text = { { config.scope.char, hl } },
        virt_text_pos = "overlay",
        virt_text_win_col = col,
        hl_mode = "combine",
        priority = config.priority + 1,
        strict = false,
        ephemeral = true,
      })
    end
  end
end

--- Render the scope overlappping the given range
---@param scope snacks.indent.Scope
---@param state snacks.indent.State
---@private
function M.render_chunk(scope, state)
  local indent = (scope.indent or 2)
  local col = indent - state.leftcol - state.shiftwidth
  if col < 0 then -- scope is hidden
    return
  end
  local to = animating(scope.buf) and scope.step or scope.to
  local hl = get_hl(scope.indent + 1, config.chunk.hl)
  local char = config.chunk.char

  ---@param l number
  ---@param line string
  local function add(l, line)
    vim.api.nvim_buf_set_extmark(scope.buf, ns, l - 1, 0, {
      virt_text = { { line, hl } },
      virt_text_pos = "overlay",
      virt_text_win_col = col,
      hl_mode = "combine",
      priority = config.priority + 2,
      strict = false,
      ephemeral = true,
    })
  end

  for l = math.max(scope.from, state.top), math.min(to, state.bottom) do
    local i = state.indents[l] - state.leftcol
    if l == scope.from then -- top line
      add(l, char.corner_top .. (char.horizontal):rep(i - col - 1))
    elseif l == scope.to then -- bottom line
      add(l, char.corner_bottom .. (char.horizontal):rep(i - col - 2) .. char.arrow)
    elseif i and i > col then -- middle line
      add(l, char.vertical)
    end
  end
end

-- Called when the scope changes
---@param win number
---@param buf number
---@param scope snacks.indent.Scope?
---@param prev snacks.indent.Scope?
---@private
function M.on_scope(win, buf, scope, prev)
  stats.scope = stats.scope + 1
  if prev then -- clear previous scope
    Snacks.util.redraw_range(win, prev.from, prev.to)
  end
  if scope then
    scope.win = win
    scope.step = scope.from
    if animating(scope.buf) then
      Snacks.animate(
        scope.from,
        scope.to,
        function(value, ctx)
          if scopes and scopes:get(win) ~= scope then
            return
          end
          scope.step = value
          Snacks.util.redraw_range(win, math.min(ctx.prev, value), math.max(ctx.prev, value))
        end,
        vim.tbl_extend("keep", {
          int = true,
          id = "indent_scope_" .. win,
          buf = buf,
        }, config.animate)
      )
    end
    Snacks.util.redraw_range(win, scope.from, animating(scope.buf) and scope.from + 1 or scope.to)
  end
end

---@private
function M.debug()
  if debug_timer:is_active() then
    debug_timer:stop()
    return
  end
  local last = {}
  debug_timer:start(50, 50, function()
    if not vim.deep_equal(stats, last) then
      last = vim.deepcopy(stats)
      Snacks.notify(vim.inspect(stats), { ft = "lua", id = "snacks_indent_debug", title = "Snacks Indent Debug" })
    end
  end)
end

--- Enable indent guides
function M.enable()
  if M.enabled then
    return
  end
  config = Snacks.config.get("indent", defaults)

  if config.debug then
    M.debug()
  end

  vim.g.snacks_animate_indent = config.animate.enabled

  M.enabled = true

  -- setup decoration provider
  vim.api.nvim_set_decoration_provider(ns, {
    on_win = function(_, win, buf, top, bottom)
      if M.enabled and config.filter(buf) then
        M.on_win(win, buf, top + 1, bottom + 1)
      end
    end,
  })

  -- Listen for scope changes
  scopes = scopes or Snacks.scope.attach(M.on_scope, config.scope)
  if not scopes.enabled then
    scopes:enable()
  end

  local group = vim.api.nvim_create_augroup("snacks_indent", { clear = true })

  vim.api.nvim_create_autocmd("ColorScheme", {
    group = group,
    callback = function()
      cache_underline = {}
    end,
  })

  -- cleanup cache
  vim.api.nvim_create_autocmd({ "WinClosed", "BufDelete", "BufWipeout" }, {
    group = group,
    callback = function()
      for win in pairs(states) do
        if not vim.api.nvim_win_is_valid(win) then
          states[win] = nil
        end
      end
    end,
  })

  -- redraw when shiftwidth changes
  vim.api.nvim_create_autocmd("OptionSet", {
    group = group,
    pattern = { "shiftwidth" },
    callback = vim.schedule_wrap(function()
      vim.cmd([[redraw!]])
    end),
  })
end

-- Disable indent guides
function M.disable()
  if not M.enabled then
    return
  end
  M.enabled = false
  if scopes then
    scopes:disable()
  end
  vim.api.nvim_del_augroup_by_name("snacks_indent")
  debug_timer:stop()
  states = {}
  stats = { indents = 0, extmarks = 0, scope = 0 }
  vim.cmd([[redraw!]])
end

return M