LSP: Expose diagnostics grouped by bufnr (#11932)

Expose `vim.lsp.buf.diagnostics_by_buf`

This makes it easier to customize the diagnostics behavior. For example
to defer the update they can override the
`textDocument/publishDiagnostics` callback to only call
`buf_diagnostics_save_positions` and then defer the other actions to a
autocmd event.

runtime/doc/lsp.txt

@@ -915,6 +915,21 @@ apply_text_edits({text_edits}, {bufnr})
 apply_workspace_edit({workspace_edit})
                 TODO: Documentation
 
+                                         *vim.lsp.util.diagnostics_by_buf*
+diagnostics_by_buf
+                A table containing diagnostics grouped by buf.
+
+                {<bufnr>: {diagnostics}}
+
+		{diagnostics} is an array of diagnostics.
+
+                By default this is populated by the
+                `textDocument/publishDiagnostics` callback via
+                |vim.lsp.util.buf_diagnostics_save_positions|.
+
+                It contains entries for active buffers. Once a buffer is
+                detached the entries for it are discarded.
+
 buf_clear_diagnostics({bufnr})          *vim.lsp.util.buf_clear_diagnostics()*
                 TODO: Documentation
  
@@ -945,9 +960,14 @@ buf_diagnostics_count({kind})
 buf_clear_references({bufnr})            *vim.lsp.util.buf_clear_references()*
                 TODO: Documentation
 
-                               *vim.lsp.util.buf_diagnostics_save_positions()*
+                                        *vim.lsp.util.buf_diagnostics_save()*
 buf_diagnostics_save_positions({bufnr}, {diagnostics})
-                TODO: Documentation
+                Stores the diagnostics into |vim.lsp.util.diagnostics_by_buf|
+
+                Parameters: ~
+                    {bufr}         bufnr for which the diagnostics are for.
+                    {diagnostics}  Diagnostics[] received from the
+                                   langauge server.
 
                                     *vim.lsp.util.buf_diagnostics_underline()*
 buf_diagnostics_underline({bufnr}, {diagnostics})

runtime/lua/vim/lsp/util.lua

@@ -6,6 +6,31 @@ local list_extend = vim.list_extend
 
 local M = {}
 
+--- Diagnostics received from the server via `textDocument/publishDiagnostics`
+-- by buffer.
+--
+--  {<bufnr>: {diagnostics}}
+--
+-- This contains only entries for active buffers. Entries for detached buffers
+-- are discarded.
+--
+-- If you override the `textDocument/publishDiagnostic` callback,
+-- this will be empty unless you call `buf_diagnostics_save_positions`.
+--
+--
+-- Diagnostic is:
+--
+-- {
+--    range: Range
+--    message: string
+--    severity?: DiagnosticSeverity
+--    code?: number | string
+--    source?: string
+--    tags?: DiagnosticTag[]
+--    relatedInformation?: DiagnosticRelatedInformation[]
+-- }
+M.diagnostics_by_buf = {}
+
 local split = vim.split
 local function split_lines(value)
   return split(value, '\n', true)
@@ -635,8 +660,6 @@ local function highlight_range(bufnr, ns, hiname, start, finish)
 end
 
 do
-  local all_buffer_diagnostics = {}
-
   local diagnostic_ns = api.nvim_create_namespace("vim_lsp_diagnostics")
   local reference_ns = api.nvim_create_namespace("vim_lsp_references")
   local sign_ns = 'vim_lsp_signs'
@@ -692,13 +715,12 @@ do
     -- if #marks == 0 then
     --   return
     -- end
-    -- local buffer_diagnostics = all_buffer_diagnostics[bufnr]
     local lines = {"Diagnostics:"}
     local highlights = {{0, "Bold"}}
 
-    local buffer_diagnostics = all_buffer_diagnostics[bufnr]
+    local buffer_diagnostics = M.diagnostics_by_buf[bufnr]
     if not buffer_diagnostics then return end
-    local line_diagnostics = buffer_diagnostics[line]
+    local line_diagnostics = M.diagnostics_group_by_line(buffer_diagnostics[line])
     if not line_diagnostics then return end
 
     for i, diagnostic in ipairs(line_diagnostics) do
@@ -726,6 +748,8 @@ do
     return popup_bufnr, winnr
   end
 
+  --- Saves the diagnostics (Diagnostic[]) into diagnostics_by_buf
+  --
   function M.buf_diagnostics_save_positions(bufnr, diagnostics)
     validate {
       bufnr = {bufnr, 'n', true};
@@ -734,28 +758,15 @@ do
     if not diagnostics then return end
     bufnr = bufnr == 0 and api.nvim_get_current_buf() or bufnr
 
-    if not all_buffer_diagnostics[bufnr] then
+    if not M.diagnostics_by_buf[bufnr] then
       -- Clean up our data when the buffer unloads.
       api.nvim_buf_attach(bufnr, false, {
         on_detach = function(b)
-          all_buffer_diagnostics[b] = nil
+          M.diagnostics_by_buf[b] = nil
         end
       })
     end
-    all_buffer_diagnostics[bufnr] = {}
-    local buffer_diagnostics = all_buffer_diagnostics[bufnr]
-
-    for _, diagnostic in ipairs(diagnostics) do
-      local start = diagnostic.range.start
-      -- local mark_id = api.nvim_buf_set_extmark(bufnr, diagnostic_ns, 0, start.line, 0, {})
-      -- buffer_diagnostics[mark_id] = diagnostic
-      local line_diagnostics = buffer_diagnostics[start.line]
-      if not line_diagnostics then
-        line_diagnostics = {}
-        buffer_diagnostics[start.line] = line_diagnostics
-      end
-      table.insert(line_diagnostics, diagnostic)
-    end
+    M.diagnostics_by_buf[bufnr] = diagnostics
   end
 
   function M.buf_diagnostics_underline(bufnr, diagnostics)
@@ -799,15 +810,26 @@ do
     end
   end
 
-  function M.buf_diagnostics_virtual_text(bufnr, diagnostics)
-    local buffer_line_diagnostics = all_buffer_diagnostics[bufnr]
-    if not buffer_line_diagnostics then
-      M.buf_diagnostics_save_positions(bufnr, diagnostics)
+  function M.diagnostics_group_by_line(diagnostics)
+    if not diagnostics then return end
+    local diagnostics_by_line = {}
+    for _, diagnostic in ipairs(diagnostics) do
+      local start = diagnostic.range.start
+      local line_diagnostics = diagnostics_by_line[start.line]
+      if not line_diagnostics then
+        line_diagnostics = {}
+        diagnostics_by_line[start.line] = line_diagnostics
+      end
+      table.insert(line_diagnostics, diagnostic)
     end
-    buffer_line_diagnostics = all_buffer_diagnostics[bufnr]
-    if not buffer_line_diagnostics then
+    return diagnostics_by_line
+  end
+
+  function M.buf_diagnostics_virtual_text(bufnr, diagnostics)
+    if not diagnostics then
       return
     end
+    local buffer_line_diagnostics = M.diagnostics_group_by_line(diagnostics)
     for line, line_diags in pairs(buffer_line_diagnostics) do
       local virt_texts = {}
       for i = 1, #line_diags - 1 do
@@ -821,12 +843,12 @@ do
   end
   function M.buf_diagnostics_count(kind)
     local bufnr = vim.api.nvim_get_current_buf()
-    local buffer_line_diagnostics = all_buffer_diagnostics[bufnr]
-    if not buffer_line_diagnostics then return end
+    local diagnostics = M.diagnostics_by_buf[bufnr]
+    if not diagnostics then return end
     local count = 0
-    for _, line_diags in pairs(buffer_line_diagnostics) do
-      for _, diag in ipairs(line_diags) do
-        if protocol.DiagnosticSeverity[kind] == diag.severity then count = count + 1 end
+    for _, diagnostic in pairs(diagnostics) do
+      if protocol.DiagnosticSeverity[kind] == diagnostic.severity then
+        count = count + 1
       end
     end
     return count

test/functional/plugin/lsp_spec.lua

@@ -845,4 +845,20 @@ describe('LSP', function()
       eq({}, exec_lua([[return vim.lsp.util.text_document_completion_list_to_complete_items(...)]], {}, prefix))
     end)
   end)
+  describe('buf_diagnostics_save_positions', function()
+    it('stores the diagnostics in diagnostics_by_buf', function ()
+      local diagnostics = {
+        { range = {}; message = "diag1" },
+        { range = {}; message = "diag2" },
+      }
+      exec_lua([[
+        vim.lsp.util.buf_diagnostics_save_positions(...)]], 0, diagnostics)
+      eq(1, exec_lua [[ return #vim.lsp.util.diagnostics_by_buf ]])
+      eq(diagnostics, exec_lua [[
+        for _, diagnostics in pairs(vim.lsp.util.diagnostics_by_buf) do
+          return diagnostics
+        end
+      ]])
+    end)
+  end)
 end)