feature(diagnostic): add vim.diagnostic.count() (#26807)

feat(diagnostic): add `vim.diagnostic.count()`

Problem: Getting diagnostic count based on the output of
  `vim.diagnostic.get()` might become costly as number of diagnostic
  entries grows. This is because it returns a copy of diagnostic cache
  entries (so as to not allow users to change them in place).
  Getting information about diagnostic count is frequently used in
  statusline, so it is important to be as fast as reasonbly possible.

Solution: Add `vim.diagnostic.count()` which computes severity
  counts without making copies.

runtime/doc/deprecated.txt

@@ -90,7 +90,7 @@ For each of the functions below, use the corresponding function in
 - *vim.lsp.diagnostic.enable()*
 - *vim.lsp.diagnostic.get()*
 - *vim.lsp.diagnostic.get_all()*	Use |vim.diagnostic.get()| instead.
-- *vim.lsp.diagnostic.get_count()*	Use |vim.diagnostic.get()| instead.
+- *vim.lsp.diagnostic.get_count()*	Use |vim.diagnostic.count()| instead.
 - *vim.lsp.diagnostic.get_line_diagnostics()* Use |vim.diagnostic.get()| instead.
 - *vim.lsp.diagnostic.get_next()*
 - *vim.lsp.diagnostic.get_next_pos()*

runtime/doc/diagnostic.txt

@@ -481,6 +481,22 @@ config({opts}, {namespace})                          *vim.diagnostic.config()*
     Return: ~
         (table|nil) table of current diagnostic config if `opts` is omitted.
 
+count({bufnr}, {opts})                                *vim.diagnostic.count()*
+    Get current diagnostics count.
+
+    Parameters: ~
+      • {bufnr}  (integer|nil) Buffer number to get diagnostics from. Use 0
+                 for current buffer or nil for all buffers.
+      • {opts}   (table|nil) A table with the following keys:
+                 • namespace: (number) Limit diagnostics to the given
+                   namespace.
+                 • lnum: (number) Limit diagnostics to the given line number.
+                 • severity: See |diagnostic-severity|.
+
+    Return: ~
+        (table) A table with actually present severity values as keys (see
+        |diagnostic-severity|) and integer counts as values.
+
 disable({bufnr}, {namespace})                       *vim.diagnostic.disable()*
     Disable diagnostics in the given buffer.
 

runtime/doc/news.txt

@@ -275,6 +275,11 @@ The following new APIs and features were added.
 • |v_Q-default| and |v_@-default| repeat a register for each line of a visual
   selection.
 
+• |vim.diagnostic.count()| returns the number of diagnostics for a given
+  buffer and/or namespace, by severity. This is a faster alternative to
+  |vim.diagnostic.get()| when only the number of diagnostics is needed, but
+  not the diagnostics themselves.
+
 ==============================================================================
 CHANGED FEATURES                                                 *news-changed*
 

runtime/lua/vim/diagnostic.lua

@@ -363,7 +363,6 @@ local function get_diagnostics(bufnr, opts, clamp)
 
   local function add(b, d)
     if not opts.lnum or d.lnum == opts.lnum then
-      d = vim.deepcopy(d)
       if clamp and api.nvim_buf_is_loaded(b) then
         local line_count = buf_line_count[b] - 1
         if
@@ -374,6 +373,7 @@ local function get_diagnostics(bufnr, opts, clamp)
           or d.col < 0
           or d.end_col < 0
         then
+          d = vim.deepcopy(d)
           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)
@@ -756,7 +756,31 @@ function M.get(bufnr, opts)
     opts = { opts, 't', true },
   })
 
-  return get_diagnostics(bufnr, opts, false)
+  return vim.deepcopy(get_diagnostics(bufnr, opts, false))
+end
+
+--- Get current diagnostics count.
+---
+---@param bufnr integer|nil Buffer number to get diagnostics from. Use 0 for
+---                        current buffer or nil for all buffers.
+---@param opts table|nil A table with the following keys:
+---                        - namespace: (number) Limit diagnostics to the given namespace.
+---                        - lnum: (number) Limit diagnostics to the given line number.
+---                        - severity: See |diagnostic-severity|.
+---@return table A 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, 'n', true },
+    opts = { opts, 't', true },
+  })
+
+  local diagnostics = get_diagnostics(bufnr, opts, false)
+  local count = {}
+  for i = 1, #diagnostics do
+    local severity = diagnostics[i].severity
+    count[severity] = (count[severity] or 0) + 1
+  end
+  return count
 end
 
 --- Get the previous diagnostic closest to the cursor position.

test/functional/lua/diagnostic_spec.lua

@@ -930,6 +930,122 @@ end)
     end)
   end)
 
+  describe('count', function()
+    it('returns actually present severity counts', function()
+      eq(
+        exec_lua [[return {
+          [vim.diagnostic.severity.ERROR] = 4,
+          [vim.diagnostic.severity.WARN] = 3,
+          [vim.diagnostic.severity.INFO] = 2,
+          [vim.diagnostic.severity.HINT] = 1,
+        }]],
+        exec_lua [[
+          vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, {
+            make_error("Error 1", 1, 1, 1, 2),
+            make_error("Error 2", 1, 3, 1, 4),
+            make_error("Error 3", 1, 5, 1, 6),
+            make_error("Error 4", 1, 7, 1, 8),
+            make_warning("Warning 1", 2, 1, 2, 2),
+            make_warning("Warning 2", 2, 3, 2, 4),
+            make_warning("Warning 3", 2, 5, 2, 6),
+            make_info("Info 1", 3, 1, 3, 2),
+            make_info("Info 2", 3, 3, 3, 4),
+            make_hint("Hint 1", 4, 1, 4, 2),
+          })
+          return vim.diagnostic.count(diagnostic_bufnr)
+        ]]
+      )
+      eq(
+        exec_lua [[return {
+          [vim.diagnostic.severity.ERROR] = 2,
+          [vim.diagnostic.severity.INFO] = 1,
+        }]],
+        exec_lua [[
+          vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, {
+            make_error("Error 1", 1, 1, 1, 2),
+            make_error("Error 2", 1, 3, 1, 4),
+            make_info("Info 1", 3, 1, 3, 2),
+          })
+          return vim.diagnostic.count(diagnostic_bufnr)
+        ]]
+      )
+    end)
+
+    it('returns only requested diagnostics count when severity range is supplied', function()
+      eq(
+        exec_lua [[return {
+          { [vim.diagnostic.severity.ERROR] = 1, [vim.diagnostic.severity.WARN] = 1 },
+          { [vim.diagnostic.severity.WARN] = 1,  [vim.diagnostic.severity.INFO] = 1, [vim.diagnostic.severity.HINT] = 1 },
+          { [vim.diagnostic.severity.WARN] = 1,  [vim.diagnostic.severity.INFO] = 1 },
+        }]],
+        exec_lua [[
+          vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, {
+            make_error("Error 1", 1, 1, 1, 5),
+            make_warning("Warning on Server 1", 1, 1, 2, 3),
+            make_info("Ignored information", 1, 1, 2, 3),
+            make_hint("Here's a hint", 1, 1, 2, 3),
+          })
+
+          return {
+            vim.diagnostic.count(diagnostic_bufnr, { severity = {min=vim.diagnostic.severity.WARN} }),
+            vim.diagnostic.count(diagnostic_bufnr, { severity = {max=vim.diagnostic.severity.WARN} }),
+            vim.diagnostic.count(diagnostic_bufnr, {
+              severity = {
+                min=vim.diagnostic.severity.INFO,
+                max=vim.diagnostic.severity.WARN,
+              }
+            }),
+          }
+        ]]
+      )
+    end)
+
+    it('returns only requested diagnostics when severities are supplied', function()
+      eq(
+        exec_lua [[return {
+          { [vim.diagnostic.severity.WARN] = 1 },
+          { [vim.diagnostic.severity.ERROR] = 1 },
+          { [vim.diagnostic.severity.WARN] = 1, [vim.diagnostic.severity.INFO] = 1 },
+        }]],
+        exec_lua [[
+          vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, {
+            make_error("Error 1", 1, 1, 1, 5),
+            make_warning("Warning on Server 1", 1, 1, 2, 3),
+            make_info("Ignored information", 1, 1, 2, 3),
+            make_hint("Here's a hint", 1, 1, 2, 3),
+          })
+
+          return {
+            vim.diagnostic.count(diagnostic_bufnr, { severity = {vim.diagnostic.severity.WARN} }),
+            vim.diagnostic.count(diagnostic_bufnr, { severity = {vim.diagnostic.severity.ERROR} }),
+            vim.diagnostic.count(diagnostic_bufnr, {
+              severity = {
+                vim.diagnostic.severity.INFO,
+                vim.diagnostic.severity.WARN,
+              }
+            }),
+          }
+        ]]
+      )
+    end)
+
+    it('allows filtering by line', function()
+      eq(
+        exec_lua [[return { [vim.diagnostic.severity.ERROR] = 1 }]],
+        exec_lua [[
+          vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, {
+            make_error("Error 1", 1, 1, 1, 5),
+            make_warning("Warning on Server 1", 1, 1, 2, 3),
+            make_info("Ignored information", 1, 1, 2, 3),
+            make_error("Error On Other Line", 2, 1, 1, 5),
+          })
+
+          return vim.diagnostic.count(diagnostic_bufnr, {lnum = 2})
+        ]]
+      )
+    end)
+  end)
+
   describe('config()', function()
     it('works with global, namespace, and ephemeral options', function()
       eq(1, exec_lua [[