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.

5 files changed, 164 insertions, 3 deletions

diff --git a/runtime/doc/deprecated.txt b/runtime/doc/deprecated.txt
index 9d32f64aef..0f2dbaa77c 100644
--- a/runtime/doc/deprecated.txt
+++ b/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()*
diff --git a/runtime/doc/diagnostic.txt b/runtime/doc/diagnostic.txt
index fa5ef22e37..106e130a41 100644
--- a/runtime/doc/diagnostic.txt
+++ b/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.
 
diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 112a76c013..8efad2454b 100644
--- a/runtime/doc/news.txt
+++ b/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*
 
diff --git a/runtime/lua/vim/diagnostic.lua b/runtime/lua/vim/diagnostic.lua
index ad40723737..a447463dff 100644
--- a/runtime/lua/vim/diagnostic.lua
+++ b/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.
diff --git a/test/functional/lua/diagnostic_spec.lua b/test/functional/lua/diagnostic_spec.lua
index 688118a085..d59ee02f01 100644
--- a/test/functional/lua/diagnostic_spec.lua
+++ b/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 [[