From 2abc799ffd8b3b250f0ffce2f74a930ad3d19cf3 Mon Sep 17 00:00:00 2001
From: Gregory Anders <greg@gpanders.com>
Date: Fri, 19 Nov 2021 11:20:04 -0700
Subject: [PATCH 1/3] fix(diagnostic): deepcopy diagnostics before clamping
 line numbers

The current 'clamp_line_numbers' implementation modifies diagnostics in
place, which can have adverse downstream side effects. Before clamping
line numbers, make a copy of the diagnostic. This commit also merges the
'clamp_line_numbers' method into a new 'get_diagnostics' local function
which also implements the more general "get" method. The public
'vim.diagnostic.get()' API now just uses this function (without
clamping). This has the added benefit that other internal API functions
that need to use get() no longer have to go through vim.validate.

Finally, reorganize the source code a bit by grouping all of the data
structures together near the top of the file.
---
 runtime/lua/vim/diagnostic.lua          | 179 +++++++++++-------------
 runtime/lua/vim/lsp/diagnostic.lua      |   2 -
 test/functional/lua/diagnostic_spec.lua |  10 +-
 3 files changed, 89 insertions(+), 102 deletions(-)

diff --git a/runtime/lua/vim/diagnostic.lua b/runtime/lua/vim/diagnostic.lua
index 0bc7a305f0..0ffc263ec1 100644
--- a/runtime/lua/vim/diagnostic.lua
+++ b/runtime/lua/vim/diagnostic.lua
@@ -36,7 +36,37 @@ M.handlers = setmetatable({}, {
   end,
 })
 
--- Local functions {{{
+-- Metatable that automatically creates an empty table when assigning to a missing key
+local bufnr_and_namespace_cacher_mt = {
+  __index = function(t, bufnr)
+    if not bufnr or bufnr == 0 then
+      bufnr = vim.api.nvim_get_current_buf()
+    end
+
+    if rawget(t, bufnr) == nil then
+      rawset(t, bufnr, {})
+    end
+
+    return rawget(t, bufnr)
+  end,
+
+  __newindex = function(t, bufnr, v)
+    if not bufnr or bufnr == 0 then
+      bufnr = vim.api.nvim_get_current_buf()
+    end
+
+    rawset(t, bufnr, v)
+  end,
+}
+
+local diagnostic_cleanup = setmetatable({}, bufnr_and_namespace_cacher_mt)
+local diagnostic_cache = setmetatable({}, bufnr_and_namespace_cacher_mt)
+local diagnostic_cache_extmarks = setmetatable({}, bufnr_and_namespace_cacher_mt)
+local diagnostic_attached_buffers = {}
+local diagnostic_disabled = {}
+local bufs_waiting_to_update = setmetatable({}, bufnr_and_namespace_cacher_mt)
+
+local all_namespaces = {}
 
 ---@private
 local function to_severity(severity)
@@ -106,8 +136,6 @@ local function reformat_diagnostics(format, diagnostics)
   return formatted
 end
 
-local all_namespaces = {}
-
 ---@private
 local function enabled_value(option, namespace)
   local ns = namespace and M.get_namespace(namespace) or {}
@@ -213,36 +241,6 @@ local function get_bufnr(bufnr)
   return bufnr
 end
 
--- Metatable that automatically creates an empty table when assigning to a missing key
-local bufnr_and_namespace_cacher_mt = {
-  __index = function(t, bufnr)
-    if not bufnr or bufnr == 0 then
-      bufnr = vim.api.nvim_get_current_buf()
-    end
-
-    if rawget(t, bufnr) == nil then
-      rawset(t, bufnr, {})
-    end
-
-    return rawget(t, bufnr)
-  end,
-
-  __newindex = function(t, bufnr, v)
-    if not bufnr or bufnr == 0 then
-      bufnr = vim.api.nvim_get_current_buf()
-    end
-
-    rawset(t, bufnr, v)
-  end,
-}
-
-local diagnostic_cleanup = setmetatable({}, bufnr_and_namespace_cacher_mt)
-local diagnostic_cache = setmetatable({}, bufnr_and_namespace_cacher_mt)
-local diagnostic_cache_extmarks = setmetatable({}, bufnr_and_namespace_cacher_mt)
-local diagnostic_attached_buffers = {}
-local diagnostic_disabled = {}
-local bufs_waiting_to_update = setmetatable({}, bufnr_and_namespace_cacher_mt)
-
 ---@private
 local function is_disabled(namespace, bufnr)
   local ns = M.get_namespace(namespace)
@@ -399,17 +397,56 @@ local function set_list(loclist, opts)
 end
 
 ---@private
---- To (slightly) improve performance, modifies diagnostics in place.
-local function clamp_line_numbers(bufnr, diagnostics)
-  local buf_line_count = vim.api.nvim_buf_line_count(bufnr)
-  if buf_line_count == 0 then
-    return
+local function get_diagnostics(bufnr, opts, clamp)
+  opts = opts or {}
+
+  local namespace = opts.namespace
+  local diagnostics = {}
+  local buf_line_count = clamp and vim.api.nvim_buf_line_count(bufnr) or math.huge
+
+  ---@private
+  local function add(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)
+      end
+      table.insert(diagnostics, d)
+    end
   end
 
-  for _, diagnostic in ipairs(diagnostics) do
-    diagnostic.lnum = math.max(math.min(diagnostic.lnum, buf_line_count - 1), 0)
-    diagnostic.end_lnum = math.max(math.min(diagnostic.end_lnum, buf_line_count - 1), 0)
+  if namespace == nil and bufnr == nil then
+    for _, t in pairs(diagnostic_cache) do
+      for _, v in pairs(t) do
+        for _, diagnostic in pairs(v) do
+          add(diagnostic)
+        end
+      end
+    end
+  elseif namespace == nil then
+    for iter_namespace in pairs(diagnostic_cache[bufnr]) do
+      for _, diagnostic in pairs(diagnostic_cache[bufnr][iter_namespace]) do
+        add(diagnostic)
+      end
+    end
+  elseif bufnr == nil then
+    for _, t in pairs(diagnostic_cache) do
+      for _, diagnostic in pairs(t[namespace] or {}) do
+        add(diagnostic)
+      end
+    end
+  else
+    for _, diagnostic in pairs(diagnostic_cache[bufnr][namespace] or {}) do
+      add(diagnostic)
+    end
   end
+
+  if opts.severity then
+    diagnostics = filter_by_severity(opts.severity, diagnostics)
+  end
+
+  return diagnostics
 end
 
 ---@private
@@ -418,8 +455,7 @@ local function next_diagnostic(position, search_forward, bufnr, opts, namespace)
   bufnr = get_bufnr(bufnr)
   local wrap = vim.F.if_nil(opts.wrap, true)
   local line_count = vim.api.nvim_buf_line_count(bufnr)
-  local diagnostics = M.get(bufnr, vim.tbl_extend("keep", opts, {namespace = namespace}))
-  clamp_line_numbers(bufnr, diagnostics)
+  local diagnostics = get_diagnostics(bufnr, vim.tbl_extend("keep", opts, {namespace = namespace}), true)
   local line_diagnostics = diagnostic_lines(diagnostics)
   for i = 0, line_count do
     local offset = i * (search_forward and 1 or -1)
@@ -481,10 +517,6 @@ local function diagnostic_move_pos(opts, pos)
   end
 end
 
--- }}}
-
--- Public API {{{
-
 --- Configure diagnostic options globally or for a specific diagnostic
 --- namespace.
 ---
@@ -689,49 +721,7 @@ function M.get(bufnr, opts)
     opts = { opts, 't', true },
   }
 
-  opts = opts or {}
-
-  local namespace = opts.namespace
-  local diagnostics = {}
-
-  ---@private
-  local function add(d)
-    if not opts.lnum or d.lnum == opts.lnum then
-      table.insert(diagnostics, d)
-    end
-  end
-
-  if namespace == nil and bufnr == nil then
-    for _, t in pairs(diagnostic_cache) do
-      for _, v in pairs(t) do
-        for _, diagnostic in pairs(v) do
-          add(diagnostic)
-        end
-      end
-    end
-  elseif namespace == nil then
-    for iter_namespace in pairs(diagnostic_cache[bufnr]) do
-      for _, diagnostic in pairs(diagnostic_cache[bufnr][iter_namespace]) do
-        add(diagnostic)
-      end
-    end
-  elseif bufnr == nil then
-    for _, t in pairs(diagnostic_cache) do
-      for _, diagnostic in pairs(t[namespace] or {}) do
-        add(diagnostic)
-      end
-    end
-  else
-    for _, diagnostic in pairs(diagnostic_cache[bufnr][namespace] or {}) do
-      add(diagnostic)
-    end
-  end
-
-  if opts.severity then
-    diagnostics = filter_by_severity(opts.severity, diagnostics)
-  end
-
-  return diagnostics
+  return get_diagnostics(bufnr, opts, false)
 end
 
 --- Get the previous diagnostic closest to the cursor position.
@@ -1115,7 +1105,7 @@ function M.show(namespace, bufnr, diagnostics, opts)
 
   M.hide(namespace, bufnr)
 
-  diagnostics = diagnostics or M.get(bufnr, {namespace=namespace})
+  diagnostics = diagnostics or get_diagnostics(bufnr, {namespace=namespace}, true)
 
   if not diagnostics or vim.tbl_isempty(diagnostics) then
     return
@@ -1141,8 +1131,6 @@ function M.show(namespace, bufnr, diagnostics, opts)
     end
   end
 
-  clamp_line_numbers(bufnr, diagnostics)
-
   for handler_name, handler in pairs(M.handlers) do
     if handler.show and opts[handler_name] then
       handler.show(namespace, bufnr, diagnostics, opts)
@@ -1213,8 +1201,7 @@ function M.open_float(bufnr, opts)
     opts = get_resolved_options({ float = float_opts }, nil, bufnr).float
   end
 
-  local diagnostics = M.get(bufnr, opts)
-  clamp_line_numbers(bufnr, diagnostics)
+  local diagnostics = get_diagnostics(bufnr, opts, true)
 
   if scope == "line" then
     diagnostics = vim.tbl_filter(function(d)
@@ -1563,6 +1550,4 @@ function M.fromqflist(list)
   return diagnostics
 end
 
--- }}}
-
 return M
diff --git a/runtime/lua/vim/lsp/diagnostic.lua b/runtime/lua/vim/lsp/diagnostic.lua
index 1e6f83c1ba..76a4dc30b7 100644
--- a/runtime/lua/vim/lsp/diagnostic.lua
+++ b/runtime/lua/vim/lsp/diagnostic.lua
@@ -717,5 +717,3 @@ end
 -- }}}
 
 return M
-
--- vim: fdm=marker
diff --git a/test/functional/lua/diagnostic_spec.lua b/test/functional/lua/diagnostic_spec.lua
index 6414483c0d..00e261c01d 100644
--- a/test/functional/lua/diagnostic_spec.lua
+++ b/test/functional/lua/diagnostic_spec.lua
@@ -110,17 +110,21 @@ describe('vim.diagnostic', function()
 
   it('retrieves diagnostics from all buffers and namespaces', function()
     local result = exec_lua [[
-      vim.diagnostic.set(diagnostic_ns, 1, {
+      local other_bufnr = vim.api.nvim_create_buf(true, false)
+      local lines = {"1st line of text", "2nd line of text", "wow", "cool", "more", "lines"}
+      vim.api.nvim_buf_set_lines(other_bufnr, 0, 1, false, lines)
+
+      vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, {
         make_error('Diagnostic #1', 1, 1, 1, 1),
         make_error('Diagnostic #2', 2, 1, 2, 1),
       })
-      vim.diagnostic.set(other_ns, 2, {
+      vim.diagnostic.set(other_ns, other_bufnr, {
         make_error('Diagnostic #3', 3, 1, 3, 1),
       })
       return vim.diagnostic.get()
     ]]
     eq(3, #result)
-    eq(2, exec_lua([[return #vim.tbl_filter(function(d) return d.bufnr == 1 end, ...)]], result))
+    eq(2, exec_lua([[return #vim.tbl_filter(function(d) return d.bufnr == diagnostic_bufnr end, ...)]], result))
     eq('Diagnostic #1', result[1].message)
   end)
 

From 34bb5fa5a942708c32b0e68eca02261d04e51ca6 Mon Sep 17 00:00:00 2001
From: Gregory Anders <greg@gpanders.com>
Date: Fri, 19 Nov 2021 11:31:33 -0700
Subject: [PATCH 2/3] fix(diagnostic): fix navigation with diagnostics placed
 past end of line

---
 runtime/lua/vim/diagnostic.lua          |  5 +++--
 test/functional/lua/diagnostic_spec.lua | 13 +++++++++++++
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/runtime/lua/vim/diagnostic.lua b/runtime/lua/vim/diagnostic.lua
index 0ffc263ec1..2a26a36d10 100644
--- a/runtime/lua/vim/diagnostic.lua
+++ b/runtime/lua/vim/diagnostic.lua
@@ -467,13 +467,14 @@ local function next_diagnostic(position, search_forward, bufnr, opts, namespace)
       lnum = (lnum + line_count) % line_count
     end
     if line_diagnostics[lnum] and not vim.tbl_isempty(line_diagnostics[lnum]) then
+      local line_length = #vim.api.nvim_buf_get_lines(bufnr,  lnum, lnum + 1, true)[1]
       local sort_diagnostics, is_next
       if search_forward then
         sort_diagnostics = function(a, b) return a.col < b.col end
-        is_next = function(diagnostic) return diagnostic.col > position[2] end
+        is_next = function(d) return math.min(d.col, line_length - 1) > position[2] end
       else
         sort_diagnostics = function(a, b) return a.col > b.col end
-        is_next = function(diagnostic) return diagnostic.col < position[2] end
+        is_next = function(d) return math.min(d.col, line_length - 1) < position[2] end
       end
       table.sort(line_diagnostics[lnum], sort_diagnostics)
       if i == 0 then
diff --git a/test/functional/lua/diagnostic_spec.lua b/test/functional/lua/diagnostic_spec.lua
index 00e261c01d..cdf9a8a834 100644
--- a/test/functional/lua/diagnostic_spec.lua
+++ b/test/functional/lua/diagnostic_spec.lua
@@ -686,6 +686,19 @@ describe('vim.diagnostic', function()
         return vim.diagnostic.get_prev_pos { namespace = diagnostic_ns }
       ]])
     end)
+
+    it('works with diagnostics past the end of the line #16349', function()
+      eq({4, 0}, exec_lua [[
+        vim.diagnostic.set(diagnostic_ns, diagnostic_bufnr, {
+          make_error('Diagnostic #1', 3, 9001, 3, 9001),
+          make_error('Diagnostic #2', 4, 0, 4, 0),
+        })
+        vim.api.nvim_win_set_buf(0, diagnostic_bufnr)
+        vim.api.nvim_win_set_cursor(0, {1, 1})
+        vim.diagnostic.goto_next { float = false }
+        return vim.diagnostic.get_next_pos { namespace = diagnostic_ns }
+      ]])
+    end)
   end)
 
   describe('get_prev_pos()', function()

From 5b0d8f85fdb705b07143fc4019189a9dcfe3c108 Mon Sep 17 00:00:00 2001
From: Gregory Anders <greg@gpanders.com>
Date: Fri, 19 Nov 2021 11:49:44 -0700
Subject: [PATCH 3/3] fix(diagnostic): clamp line numbers in setqflist and
 setloclist

---
 runtime/lua/vim/diagnostic.lua | 44 +++++++++++++++++-----------------
 1 file changed, 22 insertions(+), 22 deletions(-)

diff --git a/runtime/lua/vim/diagnostic.lua b/runtime/lua/vim/diagnostic.lua
index 2a26a36d10..ae0bbb7d05 100644
--- a/runtime/lua/vim/diagnostic.lua
+++ b/runtime/lua/vim/diagnostic.lua
@@ -374,28 +374,6 @@ local function clear_scheduled_display(namespace, bufnr)
   end
 end
 
----@private
-local function set_list(loclist, opts)
-  opts = opts or {}
-  local open = vim.F.if_nil(opts.open, true)
-  local title = opts.title or "Diagnostics"
-  local winnr = opts.winnr or 0
-  local bufnr
-  if loclist then
-    bufnr = vim.api.nvim_win_get_buf(winnr)
-  end
-  local diagnostics = M.get(bufnr, opts)
-  local items = M.toqflist(diagnostics)
-  if loclist then
-    vim.fn.setloclist(winnr, {}, ' ', { title = title, items = items })
-  else
-    vim.fn.setqflist({}, ' ', { title = title, items = items })
-  end
-  if open then
-    vim.api.nvim_command(loclist and "lopen" or "copen")
-  end
-end
-
 ---@private
 local function get_diagnostics(bufnr, opts, clamp)
   opts = opts or {}
@@ -449,6 +427,28 @@ local function get_diagnostics(bufnr, opts, clamp)
   return diagnostics
 end
 
+---@private
+local function set_list(loclist, opts)
+  opts = opts or {}
+  local open = vim.F.if_nil(opts.open, true)
+  local title = opts.title or "Diagnostics"
+  local winnr = opts.winnr or 0
+  local bufnr
+  if loclist then
+    bufnr = vim.api.nvim_win_get_buf(winnr)
+  end
+  local diagnostics = get_diagnostics(bufnr, opts, true)
+  local items = M.toqflist(diagnostics)
+  if loclist then
+    vim.fn.setloclist(winnr, {}, ' ', { title = title, items = items })
+  else
+    vim.fn.setqflist({}, ' ', { title = title, items = items })
+  end
+  if open then
+    vim.api.nvim_command(loclist and "lopen" or "copen")
+  end
+end
+
 ---@private
 local function next_diagnostic(position, search_forward, bufnr, opts, namespace)
   position[1] = position[1] - 1