Repository: neovim/neovim
File: runtime/lua/vim/lsp/diagnostic.lua (Lua)

Code before:
    return
  end

  if result == nil then
    return
  end

  local client_id = ctx.client_id
  local bufnr = assert(ctx.bufnr)
  local bufstate = bufstates[bufnr]
  bufstate.client_result_id[client_id] = result.resultId

  if result.kind == 'unchanged' then
    return
  end

  handle_diagnostics(ctx.params.textDocument.uri, client_id, result.items, true)

  for uri, related_result in pairs(result.relatedDocuments or {}) do
    if related_result.kind == 'full' then
      handle_diagnostics(uri, client_id, related_result.items, true)
    end

    local related_bufnr = vim.uri_to_bufnr(uri)
    local related_bufstate = bufstates[related_bufnr]
      or { pull_kind = 'document', client_result_id = {} }
    bufstates[related_bufnr] = related_bufstate

    related_bufstate.client_result_id[client_id] = related_result.resultId
  end
end

--- Clear push diagnostics and diagnostic cache.
---
--- Diagnostic producers should prefer |vim.diagnostic.reset()|. However,
--- this method signature is still used internally in some parts of the LSP
--- implementation so it's simply marked @private rather than @deprecated.
---
---@param client_id integer
---@param buffer_client_map table<integer, table<integer, table>> map of buffers to active clients
---@private
function M.reset(client_id, buffer_client_map)
  buffer_client_map = vim.deepcopy(buffer_client_map)
  vim.schedule(function()
    for bufnr, client_ids in pairs(buffer_client_map) do
      if client_ids[client_id] then
        local namespace = M.get_namespace(client_id, false)
        vim.diagnostic.reset(namespace, bufnr)
      end
    end
  end)

Review comment (suggestion):
Hmmm it's a bit interesting that VS Code doesn't support this yet: https://github.com/microsoft/vscode-languageserver-node/blob/df56e720c01c6e2d7873733807418f6ce33187ad/client/src/common/diagnostic.ts#L1067

What do you think of having a partial approach here of only updating `related_bufstate` if `bufstates` contains `related_bufnr`? Once we add something to `bufstates` we'll start pulling for diagnostics, and I don't want for the state to blow up from a bunch of related documents that haven't even been opened.

Code after:
    end
    return
  end

  if result == nil then
    return
  end

  local client_id = ctx.client_id
  local bufnr = assert(ctx.bufnr)
  local bufstate = bufstates[bufnr]
  bufstate.client_result_id[client_id] = result.resultId

  if result.kind == 'unchanged' then
    return
  end

  handle_diagnostics(ctx.params.textDocument.uri, client_id, result.items, true)

  for uri, related_result in pairs(result.relatedDocuments or {}) do
    if related_result.kind == 'full' then
      handle_diagnostics(uri, client_id, related_result.items, true)
    end

    local related_bufnr = vim.uri_to_bufnr(uri)
    local related_bufstate = bufstates[related_bufnr]
      -- Create a new bufstate if it doesn't exist for the related document. This will not enable
      -- diagnostic pulling by itself, but will allow previous result IDs to be passed correctly the
      -- next time this buffer's diagnostics are pulled.
      or { pull_kind = 'document', client_result_id = {} }
    bufstates[related_bufnr] = related_bufstate

    related_bufstate.client_result_id[client_id] = related_result.resultId
  end
end

--- Clear push diagnostics and diagnostic cache.
---
--- Diagnostic producers should prefer |vim.diagnostic.reset()|. However,
--- this method signature is still used internally in some parts of the LSP
--- implementation so it's simply marked @private rather than @deprecated.
---
---@param client_id integer
---@param buffer_client_map table<integer, table<integer, table>> map of buffers to active clients
---@private
function M.reset(client_id, buffer_client_map)
  buffer_client_map = vim.deepcopy(buffer_client_map)
  vim.schedule(function()
    for bufnr, client_ids in pairs(buffer_client_map) do
      if client_ids[client_id] then
        local namespace = M.get_namespace(client_id, false)