From c59f2008e0d4998094123e4976210a78e8f6b4d0 Mon Sep 17 00:00:00 2001
From: Gregory Anders <greg@gpanders.com>
Date: Wed, 24 Nov 2021 19:32:26 -0700
Subject: fix(diagnostic): get line count per buffer when clamping

Fixes a bug when `get_diagnostics` is called with a nil `bufnr`.
Diagnostics should be clamped for the buffer they reside in, not the
current buffer.
---
 runtime/lua/vim/diagnostic.lua | 34 ++++++++++++++++++++++------------
 1 file changed, 22 insertions(+), 12 deletions(-)

(limited to 'runtime')

diff --git a/runtime/lua/vim/diagnostic.lua b/runtime/lua/vim/diagnostic.lua
index 74996faf3c..85b430375e 100644
--- a/runtime/lua/vim/diagnostic.lua
+++ b/runtime/lua/vim/diagnostic.lua
@@ -372,25 +372,35 @@ local function get_diagnostics(bufnr, opts, clamp)
 
   local namespace = opts.namespace
   local diagnostics = {}
-  local buf_line_count = clamp and vim.api.nvim_buf_line_count(bufnr) or math.huge
+
+  -- Memoized results of buf_line_count per bufnr
+  local buf_line_count = setmetatable({}, {
+    __index = function(t, k)
+      t[k] = vim.api.nvim_buf_line_count(k)
+      return rawget(t, k)
+    end,
+  })
 
   ---@private
-  local function add(d)
+  local function add(b, d)
     if not opts.lnum or d.lnum == opts.lnum then
-      if clamp and (d.lnum >= buf_line_count or d.end_lnum >= buf_line_count) then
-        d = vim.deepcopy(d)
-        d.lnum = math.max(math.min(d.lnum, buf_line_count - 1), 0)
-        d.end_lnum = math.max(math.min(d.end_lnum, buf_line_count - 1), 0)
+      if clamp and vim.api.nvim_buf_is_loaded(b) then
+        local line_count = buf_line_count[b] - 1
+        if (d.lnum > line_count or d.end_lnum > line_count) 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)
+        end
       end
       table.insert(diagnostics, d)
     end
   end
 
   if namespace == nil and bufnr == nil then
-    for _, t in pairs(diagnostic_cache) do
+    for b, t in pairs(diagnostic_cache) do
       for _, v in pairs(t) do
         for _, diagnostic in pairs(v) do
-          add(diagnostic)
+          add(b, diagnostic)
         end
       end
     end
@@ -398,19 +408,19 @@ local function get_diagnostics(bufnr, opts, clamp)
     bufnr = get_bufnr(bufnr)
     for iter_namespace in pairs(diagnostic_cache[bufnr]) do
       for _, diagnostic in pairs(diagnostic_cache[bufnr][iter_namespace]) do
-        add(diagnostic)
+        add(bufnr, diagnostic)
       end
     end
   elseif bufnr == nil then
-    for _, t in pairs(diagnostic_cache) do
+    for b, t in pairs(diagnostic_cache) do
       for _, diagnostic in pairs(t[namespace] or {}) do
-        add(diagnostic)
+        add(b, diagnostic)
       end
     end
   else
     bufnr = get_bufnr(bufnr)
     for _, diagnostic in pairs(diagnostic_cache[bufnr][namespace] or {}) do
-      add(diagnostic)
+      add(bufnr, diagnostic)
     end
   end
 
-- 
cgit