path: root/runtime/lua/vim/diagnostic/_store.lua

local api = vim.api
local nvim_on = require('vim._core.util').nvim_on

local severity_module = require('vim.diagnostic._severity')

--- @class (private) vim.diagnostic._store
local M = {}

-- bufnr -> ns -> Diagnostic[]
local diagnostic_cache = {} --- @type table<integer,table<integer,vim.Diagnostic[]?>>

local group = api.nvim_create_augroup('nvim.diagnostic.buf_wipeout')
setmetatable(diagnostic_cache, {
  --- @param t table<integer,vim.Diagnostic[]>
  --- @param bufnr integer
  __index = function(t, bufnr)
    assert(bufnr > 0, 'Invalid buffer number')
    nvim_on('BufWipeout', group, { buf = bufnr }, function()
      rawset(t, bufnr, nil)
    end)
    t[bufnr] = {}
    return t[bufnr]
  end,
})

--- @param bufnr integer
--- @param namespace integer
--- @param d vim.Diagnostic.Set
local function norm_diag(bufnr, namespace, d)
  vim.validate('diagnostic.lnum', d.lnum, 'number')
  local d1 = d --[[@as vim.Diagnostic]]
  d1.severity = d.severity and severity_module.to_severity(d.severity)
    or vim.diagnostic.severity.ERROR
  d1.end_lnum = d.end_lnum or d.lnum
  d1.col = d.col or 0
  d1.end_col = d.end_col or d.col or 0
  d1.namespace = namespace
  d1.bufnr = bufnr
end

--- Set extmarks at the computed positions of the cached diagnostics of
--- {bufnr}/{namespace}, and store their ids in diagnostic._extmark_id (used by
--- get_logical_pos to adjust positions).
---
--- @param bufnr integer
--- @param namespace integer
local function set_position_extmarks(bufnr, namespace)
  local ns = vim.diagnostic.get_namespace(namespace)

  if not ns.user_data.location_ns then
    ns.user_data.location_ns =
      api.nvim_create_namespace(string.format('nvim.%s.diagnostic', ns.name))
  end

  api.nvim_buf_clear_namespace(bufnr, ns.user_data.location_ns, 0, -1)

  local lines = api.nvim_buf_get_lines(bufnr, 0, -1, true)
  local buf_cache = rawget(diagnostic_cache, bufnr) --- @type table<integer,vim.Diagnostic[]?>?
  -- set extmarks at diagnostic locations to preserve logical positions despite text changes
  for _, diagnostic0 in ipairs(buf_cache and buf_cache[namespace] or {}) do
    --- @cast diagnostic0 vim.Diagnostic
    local last_row = #lines - 1
    local row = math.max(0, math.min(diagnostic0.lnum, last_row))
    local row_len = #lines[row + 1]
    local col = math.max(0, math.min(diagnostic0.col, row_len - 1))

    local end_row = math.max(0, math.min(diagnostic0.end_lnum or row, last_row))
    local end_row_len = #lines[end_row + 1]
    local end_col = math.max(0, math.min(diagnostic0.end_col or col, end_row_len))

    if end_row == row then
      -- avoid starting an extmark beyond end of the line
      if end_col == col then
        end_col = math.min(end_col + 1, end_row_len)
      end
    else
      -- avoid ending an extmark before start of the line
      if end_col == 0 then
        end_row = end_row - 1

        local end_line = lines[end_row + 1]

        if not end_line then
          error(
            'Failed to adjust diagnostic position to the end of a previous line. #lines in a buffer: '
              .. #lines
              .. ', lnum: '
              .. diagnostic0.lnum
              .. ', col: '
              .. diagnostic0.col
              .. ', end_lnum: '
              .. diagnostic0.end_lnum
              .. ', end_col: '
              .. diagnostic0.end_col
          )
        end

        end_col = #end_line
      end
    end

    diagnostic0._extmark_id = api.nvim_buf_set_extmark(bufnr, ns.user_data.location_ns, row, col, {
      end_row = end_row,
      end_col = end_col,
      invalidate = true,
    })
  end
end

-- Positions can only be computed (and handlers can only draw) once the buffer is loaded, so set
-- the extmarks and show diagnostics for any buffer with cached diagnostics when it is read.
nvim_on('BufRead', api.nvim_create_augroup('nvim.diagnostic.buf_load'), function(ev)
  local buf_cache = rawget(diagnostic_cache, ev.buf) --- @type table<integer,vim.Diagnostic[]?>?
  if buf_cache then
    for namespace in pairs(buf_cache) do
      set_position_extmarks(ev.buf, namespace)
    end
    vim.diagnostic.show(nil, ev.buf)
  end
end)

--- @param bufnr integer?
--- @param opts vim.diagnostic.GetOpts?
--- @param clamp boolean
--- @return vim.Diagnostic[]
function M.get_diagnostics(bufnr, opts, clamp)
  opts = opts or {}

  local namespace = opts.namespace

  if type(namespace) == 'number' then
    namespace = { namespace }
  end

  --- @cast namespace integer[]

  --- @type vim.Diagnostic[]
  local diagnostics = {}

  -- Memoized results of buf_line_count per bufnr
  --- @type table<integer,integer>
  local buf_line_count = setmetatable({}, {
    --- @param t table<integer,integer>
    --- @param k integer
    --- @return integer
    __index = function(t, k)
      t[k] = api.nvim_buf_line_count(k)
      return rawget(t, k)
    end,
  })

  local match_severity = opts.severity and severity_module.severity_predicate(opts.severity)
    or function(_)
      return true
    end

  --- @param b integer
  --- @param d vim.Diagnostic
  local match_enablement = function(d, b)
    if opts.enabled == nil then
      return true
    end

    local enabled = vim.diagnostic.is_enabled({ bufnr = b, ns_id = d.namespace })

    return (enabled and opts.enabled) or (not enabled and not opts.enabled)
  end

  --- @param b integer
  --- @param d vim.Diagnostic
  local function add(b, d)
    if
      match_severity(d)
      and match_enablement(d, b)
      and (not opts.lnum or (opts.lnum >= d.lnum and opts.lnum <= (d.end_lnum or d.lnum)))
    then
      if clamp and api.nvim_buf_is_loaded(b) then
        local line_count = buf_line_count[b] - 1
        if
          d.lnum > line_count
          or d.end_lnum > line_count
          or d.lnum < 0
          or d.end_lnum < 0
          or d.col < 0
          or d.end_col < 0
        then
          d = vim.deepcopy(d, true)
          d.lnum = math.max(math.min(d.lnum, line_count), 0)
          d.end_lnum = math.max(math.min(d.end_lnum, line_count), 0)
          d.col = math.max(d.col, 0)
          d.end_col = math.max(d.end_col, 0)
        end
      end
      table.insert(diagnostics, d)
    end
  end

  --- @param buf integer
  --- @param diags vim.Diagnostic[]
  local function add_all_diags(buf, diags)
    for _, diagnostic0 in pairs(diags) do
      add(buf, diagnostic0)
    end
  end

  if not namespace and not bufnr then
    for buf, ns_diags in pairs(diagnostic_cache) do
      for _, diags in pairs(ns_diags) do
        add_all_diags(buf, diags)
      end
    end
  elseif not namespace then
    bufnr = vim._resolve_bufnr(bufnr)
    for _, diags in pairs(diagnostic_cache[bufnr]) do
      add_all_diags(bufnr, diags)
    end
  elseif bufnr == nil then
    for b, t in pairs(diagnostic_cache) do
      for _, iter_namespace in ipairs(namespace) do
        add_all_diags(b, t[iter_namespace] or {})
      end
    end
  else
    bufnr = vim._resolve_bufnr(bufnr)
    for _, iter_namespace in ipairs(namespace) do
      add_all_diags(bufnr, diagnostic_cache[bufnr][iter_namespace] or {})
    end
  end

  return diagnostics
end

--- @return integer[]
function M.get_bufnrs()
  return vim.tbl_keys(diagnostic_cache)
end

--- @param bufnr integer
--- @return integer[]
function M.get_buf_namespaces(bufnr)
  return vim.tbl_keys(diagnostic_cache[vim._resolve_bufnr(bufnr)])
end

--- @param namespace integer
--- @param bufnr integer
function M.clear(namespace, bufnr)
  diagnostic_cache[vim._resolve_bufnr(bufnr)][namespace] = nil
end

--- @param bufnr integer
function M.drop_buf(bufnr)
  diagnostic_cache[vim._resolve_bufnr(bufnr)] = nil
end

--- Set diagnostics for the given namespace and buffer.
---
--- @param namespace integer The diagnostic namespace
--- @param bufnr integer Buffer number
--- @param diagnostics vim.Diagnostic.Set[]
function M.set(namespace, bufnr, diagnostics)
  vim.validate('namespace', namespace, 'number')
  vim.validate('bufnr', bufnr, 'number')
  vim.validate('diagnostics', diagnostics, vim.islist, 'a list of diagnostics')

  bufnr = vim._resolve_bufnr(bufnr)

  for _, diagnostic0 in ipairs(diagnostics) do
    norm_diag(bufnr, namespace, diagnostic0)
  end

  --- @cast diagnostics vim.Diagnostic[]

  if vim.tbl_isempty(diagnostics) then
    diagnostic_cache[bufnr][namespace] = nil
  else
    diagnostic_cache[bufnr][namespace] = diagnostics
  end

  if api.nvim_buf_is_loaded(bufnr) then
    set_position_extmarks(bufnr, namespace)
  end
end

--- @param bufnr integer? Buffer number to get diagnostics from. Use 0 for
---                      current buffer or nil for all buffers.
--- @param opts? vim.diagnostic.GetOpts
--- @return vim.Diagnostic[] : Fields `bufnr`, `end_lnum`, `end_col`, and `severity`
---                           are guaranteed to be present.
function M.get(bufnr, opts)
  vim.validate('bufnr', bufnr, 'number', true)
  vim.validate('opts', opts, 'table', true)

  return vim.deepcopy(M.get_diagnostics(bufnr, opts, false), true)
end

--- @param bufnr? integer Buffer number to get diagnostics from. Use 0 for
---                      current buffer or nil for all buffers.
--- @param opts? vim.diagnostic.GetOpts
--- @return table<integer, integer> : Table with actually present severity values as keys
---                (see |diagnostic-severity|) and integer counts as values.
function M.count(bufnr, opts)
  vim.validate('bufnr', bufnr, 'number', true)
  vim.validate('opts', opts, 'table', true)

  local diagnostics = M.get_diagnostics(bufnr, opts, false)
  local count = {} --- @type table<integer,integer>
  for _, d in ipairs(diagnostics) do
    count[d.severity] = (count[d.severity] or 0) + 1
  end
  return count
end

return M