fix(diagnostic): deepcopy diagnostics before clamping line numbers

gpanders · gpanders · commit 2abc799ffd8b · 2021-11-19T11:37:45.000-07:00
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.
diff --git 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
@@ -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
@@ -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)
 
