Assign diagnostic version per node to reduce overhead

Signed-off-by: iusmac <[email protected]>

Author: iusmac

lua/nvim-tree/actions/moves/item.lua

@@ -3,6 +3,7 @@ local view = require "nvim-tree.view"
 local core = require "nvim-tree.core"
 local lib = require "nvim-tree.lib"
 local explorer_node = require "nvim-tree.explorer.node"
+local diagnostics = require "nvim-tree.diagnostics"
 
 local M = {}
 
@@ -33,7 +34,8 @@ function M.fn(opts)
         local git_status = explorer_node.get_git_status(node)
         valid = git_status ~= nil and (not opts.skip_gitignored or git_status[1] ~= "!!")
       elseif opts.what == "diag" then
-        valid = node.diag_status ~= nil
+        local diag_status = diagnostics.get_diag_status(node)
+        valid = diag_status ~= nil and diag_status.value ~= nil
       elseif opts.what == "opened" then
         valid = vim.fn.bufloaded(node.absolute_path) ~= 0
       end

lua/nvim-tree/diagnostics.lua

@@ -11,10 +11,18 @@ local severity_levels = {
   Hint = 4,
 }
 
+---@class DiagStatus
+---@field value integer|nil
+---@field cache_version integer
+
 --- A dictionary tree containing buffer-severity mappings.
 ---@type table
 local buffer_severity_dict = {}
 
+--- The cache version number of the buffer-severity mappings.
+---@type integer
+local BUFFER_SEVERITY_VERSION = 0
+
 ---@param path string
 ---@return string
 local function uniformize_path(path)
@@ -80,6 +88,31 @@ local function is_using_coc()
   return vim.g.coc_service_initialized == 1
 end
 
+---@param node Node
+---@return DiagStatus
+local function from_cache(node)
+  local nodepath = uniformize_path(node.absolute_path)
+  local max_severity = nil
+  if not node.nodes then
+    -- direct cache hit for files
+    max_severity = buffer_severity_dict[nodepath]
+  else
+    -- dirs should be searched in the list of cached buffer names by prefix
+    for bufname, severity in pairs(buffer_severity_dict) do
+      local node_contains_buf = vim.startswith(bufname, nodepath .. "/")
+      if node_contains_buf then
+        if severity == M.severity.max then
+          max_severity = severity
+          break
+        else
+          max_severity = math.min(max_severity or severity, severity)
+        end
+      end
+    end
+  end
+  return { value = max_severity, cache_version = BUFFER_SEVERITY_VERSION }
+end
+
 function M.update()
   if not M.enable then
     return
@@ -91,6 +124,7 @@ function M.update()
     else
       buffer_severity_dict = from_nvim_lsp()
     end
+    BUFFER_SEVERITY_VERSION = BUFFER_SEVERITY_VERSION + 1
     log.node("diagnostics", buffer_severity_dict, "update")
     log.profile_end(profile)
     if view.is_buf_valid(view.get_bufnr()) then
@@ -100,33 +134,34 @@ function M.update()
 end
 
 ---@param node Node
-function M.update_node_severity_level(node)
+---@return DiagStatus|nil
+function M.get_diag_status(node)
   if not M.enable then
-    return
+    return nil
   end
 
-  local is_folder = node.nodes ~= nil
-  local nodepath = uniformize_path(node.absolute_path)
+  -- dir but we shouldn't show on dirs at all
+  if node.nodes ~= nil and not M.show_on_dirs then
+    return nil
+  end
 
-  if is_folder then
-    local max_severity = nil
-    if M.show_on_dirs and (not node.open or M.show_on_open_dirs) then
-      for bufname, severity in pairs(buffer_severity_dict) do
-        local node_contains_buf = vim.startswith(bufname, nodepath .. "/")
-        if node_contains_buf then
-          if severity == M.severity.max then
-            max_severity = severity
-            break
-          else
-            max_severity = math.min(max_severity or severity, severity)
-          end
-        end
-      end
-    end
-    node.diag_status = max_severity
-  else
-    node.diag_status = buffer_severity_dict[nodepath]
+  -- here, we do a lazy update of the diagnostic status carried by the node.
+  -- This is by design, as diagnostics and nodes live in completely separate
+  -- worlds, and this module is the link between the two
+  if not node.diag_status or node.diag_status.cache_version < BUFFER_SEVERITY_VERSION then
+    node.diag_status = from_cache(node)
+  end
+
+  -- file
+  if not node.nodes then
+    return node.diag_status
+  end
+
+  -- dir is closed or we should show on open_dirs
+  if not node.open or M.show_on_open_dirs then
+    return node.diag_status
   end
+  return nil
 end
 
 function M.setup(opts)

lua/nvim-tree/node.lua

@@ -17,7 +17,7 @@
 ---@field parent DirNode
 ---@field type string
 ---@field watcher function|nil
----@field diag_status integer|nil
+---@field diag_status DiagStatus|nil
 
 ---@class DirNode: BaseNode
 ---@field has_children boolean

lua/nvim-tree/renderer/builder.lua

@@ -413,8 +413,6 @@ end
 function Builder:_build_line(node, idx, num_children, unloaded_bufnr)
   local copy_paste = require "nvim-tree.actions.fs.copy-paste"
 
-  require("nvim-tree.diagnostics").update_node_severity_level(node)
-
   -- various components
   local indent_markers = pad.get_indent_markers(self.depth, idx, num_children, node, self.markers)
   local arrows = pad.get_arrows(node)

lua/nvim-tree/renderer/components/diagnostics.lua

@@ -1,4 +1,5 @@
 local HL_POSITION = require("nvim-tree.enum").HL_POSITION
+local diagnostics = require "nvim-tree.diagnostics"
 
 local M = {
   HS_FILE = {},
@@ -17,10 +18,11 @@ function M.get_highlight(node)
   end
 
   local group
+  local diag_status = diagnostics.get_diag_status(node)
   if node.nodes then
-    group = M.HS_FOLDER[node.diag_status]
+    group = M.HS_FOLDER[diag_status and diag_status.value]
   else
-    group = M.HS_FILE[node.diag_status]
+    group = M.HS_FILE[diag_status and diag_status.value]
   end
 
   if group then
@@ -35,7 +37,8 @@ end
 ---@return HighlightedString|nil modified icon
 function M.get_icon(node)
   if node and M.config.diagnostics.enable and M.config.renderer.icons.show.diagnostics then
-    return M.ICON[node.diag_status]
+    local diag_status = diagnostics.get_diag_status(node)
+    return M.ICON[diag_status and diag_status.value]
   end
 end