Virt-Indent.nvim/lua/vindent/indent.lua

---@class VirtualIndent
---@field private _ns_id number extmarks namespace id
---@field private _bufnr integer Buffer VirtualIndent is attached to
---@field private _attached boolean Whether or not VirtualIndent is attached for its buffer
---@field private _bufnrs table<integer, VirtualIndent> Buffers with VirtualIndent attached
---@field private _watcher_running boolean Whether or not VirtualIndent is reacting to `vim.b.org_indent_mode`
---@field private _extmark_ids table<integer, integer[]> Mapping of line number to active extmarks
---@field private _tree_utils table<string, function> Treesitter utilities for the given filetype
---@field private _fallback_pattern string Pattern to search for if treesitter parser fails
local VirtualIndent = {
  _ns_id = vim.api.nvim_create_namespace("VirtIndent"),
  _bufnrs = {},
}
VirtualIndent.__index = VirtualIndent

--- Creates a new instance of VirtualIndent for a given buffer or returns the existing instance if
--- one exists
---@param bufnr? integer Buffer to use for VirtualIndent when attached
---@return VirtualIndent?
function VirtualIndent:new(bufnr)
  -- TODO: Improve organization of this, ideally should be a separate module that returns these.
  local ft_settings = {
    markdown = {
      utils = require("vindent.treesitter.markdown"),
      fallback_pattern = "^%#+",
    },
    org = {
      utils = require("vindent.treesitter.org"),
      fallback_pattern = "^%*+",
    },
  }

  local filetype = vim.api.nvim_get_option_value("filetype", { buf = bufnr }) or ""
  local ft_setting = ft_settings[filetype]
  if not ft_setting then
    return
  end

  bufnr = bufnr or vim.api.nvim_get_current_buf()
  if self._bufnrs[bufnr] then
    return self._bufnrs[bufnr]
  end

  local this = setmetatable({
    _bufnr = bufnr,
    _watcher_running = false,
    _extmark_ids = {},
    _attached = false,
    _fallback_pattern = ft_setting.fallback_pattern,
    _tree_utils = ft_setting.utils,
  }, self)
  self._bufnrs[bufnr] = this
  return this
end

function VirtualIndent:_delete_old_extmarks(start_line, end_line)
  for line = start_line, end_line do
    local line_extmark_ids = self._extmark_ids[line] or {}
    for idx, extmark_id in ipairs(line_extmark_ids) do
      vim.api.nvim_buf_del_extmark(self._bufnr, self._ns_id, extmark_id)
      line_extmark_ids[idx] = nil
    end
  end
end

function VirtualIndent:_get_indent_size(line, ts_has_errors, ts_fallback_pat)
  -- If tree has errors, we can't rely on treesitter to get the correct indentation
  -- Fallback to searching closest headline by checking each previous line
  if ts_has_errors then
    local linenr = line
    while linenr > 0 do
      local _, level = vim.fn.getline(linenr):find(ts_fallback_pat)
      if level then
        -- If the current line is a headline we should return no virtual indentation, otherwise
        -- return virtual indentation
        return (linenr == line and 0 or level + 1)
      end
      linenr = linenr - 1
    end
  end

  local headline = self._tree_utils.closest_headline_node({ line + 1, 1 })

  if headline then
    local headline_line = headline:start()

    if headline_line ~= line then
      return self._tree_utils.headline_level(headline) + 1
    end
  end

  return 0
end

---@param start_line number start line number to set the indentation, 0-based inclusive
---@param end_line number end line number to set the indentation, 0-based inclusive
---@param ignore_ts? boolean whether or not to skip the treesitter start & end lookup
function VirtualIndent:set_indent(start_line, end_line, ignore_ts)
  ignore_ts = ignore_ts or false
  local headline = self._tree_utils.closest_headline_node({ start_line + 1, 1 })
  if headline and not ignore_ts then
    local parent = headline:parent()
    if parent then
      start_line = math.min(parent:start(), start_line)
      end_line = math.max(parent:end_(), end_line)
    end
  end
  if start_line > 0 then
    start_line = start_line - 1
  end

  local node_at_cursor = vim.treesitter.get_node()
  local ts_has_errors = false
  if node_at_cursor then
    ts_has_errors = node_at_cursor:tree():root():has_error()
  end

  self:_delete_old_extmarks(start_line, end_line)
  for line = start_line, end_line do
    local indent = self:_get_indent_size(line, ts_has_errors, self._fallback_pattern)

    if indent > 0 then
      -- NOTE: `ephemeral = true` is not implemented for `inline` virt_text_pos :(
      local new_extmark_id = vim.api.nvim_buf_set_extmark(self._bufnr, self._ns_id, line, 0, {
        -- HACK: The 'space' character below is not a space, it is actually a "Braille Pattern Blank"
        -- character, U+2800. This avoids issues with how `indentexpr` is calculated by not using
        -- spaces (which the `indentexpr` is looking for).
        virt_text = { { string.rep("⠀", indent), "VirtIndent" } },
        virt_text_pos = "inline",
        right_gravity = false,
        priority = 110,
      })
      if not self._extmark_ids[line] then
        self._extmark_ids[line] = { new_extmark_id }
      else
        table.insert(self._extmark_ids[line], new_extmark_id)
      end
    end
  end
end

--- Enables virtual indentation in registered buffer
function VirtualIndent:attach()
  if self._attached then
    return
  end
  self:set_indent(0, vim.api.nvim_buf_line_count(self._bufnr) - 1, true)

  vim.api.nvim_buf_attach(self._bufnr, false, {
    on_lines = function(_, _, _, start_line, _, end_line)
      if not self._attached then
        return true
      end

      vim.schedule(function()
        self:set_indent(start_line, end_line)
      end)
    end,
    on_reload = function()
      self:set_indent(0, vim.api.nvim_buf_line_count(self._bufnr) - 1, true)
    end,
    on_detach = function(_, bufnr)
      self:detach()
      self._bufnrs[bufnr] = nil
    end,
  })
  self._attached = true
end

function VirtualIndent:detach()
  if not self._attached then
    return
  end
  self:_delete_old_extmarks(0, vim.api.nvim_buf_line_count(self._bufnr) - 1)
  self._attached = false
end

return VirtualIndent