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.

1 files changed, 116 insertions, 0 deletions

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 [[