Keep current file selected after staging

lua/neogit/integrations/diffview.lua

@@ -5,10 +5,10 @@ local RevType = require("diffview.vcs.rev").RevType
 local CDiffView = require("diffview.api.views.diff.diff_view").CDiffView
 local dv_lib = require("diffview.lib")
 local dv_utils = require("diffview.utils")
+local a = require("plenary.async")
 
 local Watcher = require("neogit.watcher")
 local git = require("neogit.lib.git")
-local a = require("plenary.async")
 
 local function get_local_diff_view(section_name, item_name, opts)
   local left = Rev(RevType.STAGE)
@@ -21,20 +21,27 @@ local function get_local_diff_view(section_name, item_name, opts)
   local function update_files()
     local files = {}
 
+    git.repo:dispatch_refresh {
+      source = "diffview_update",
+      callback = function() end,
+    }
+
+    local repo_state = git.repo.state
+
     local sections = {
       conflicting = {
         items = vim.tbl_filter(function(item)
           return item.mode and item.mode:sub(2, 2) == "U"
-        end, git.repo.state.untracked.items),
+        end, repo_state.untracked and repo_state.untracked.items or {}),
       },
-      working = git.repo.state.unstaged,
-      staged = git.repo.state.staged,
+      working = repo_state.unstaged or { items = {} },
+      staged = repo_state.staged or { items = {} },
     }
 
     for kind, section in pairs(sections) do
       files[kind] = {}
 
-      for idx, item in ipairs(section.items) do
+      for _, item in ipairs(section.items or {}) do
         local file = {
           path = item.name,
           status = item.mode and item.mode:sub(1, 1),
@@ -44,12 +51,11 @@ local function get_local_diff_view(section_name, item_name, opts)
           } or nil,
           left_null = vim.tbl_contains({ "A", "?" }, item.mode),
           right_null = false,
-          selected = (item_name and item.name == item_name) or (not item_name and idx == 1),
         }
 
         -- restrict diff to only a particular section
         if opts.only then
-          if (item_name and file.selected) or (not item_name and section_name == kind) then
+          if (item_name and item.name == item_name) or (not item_name and section_name == kind) then
             table.insert(files[kind], file)
           end
         else
@@ -61,8 +67,37 @@ local function get_local_diff_view(section_name, item_name, opts)
     return files
   end
 
+  -- Get the initial file list
   local files = update_files()
 
+  -- Apply the "selected" logic only to the initial list
+  local file_was_selected = false
+  if item_name then
+    for _, section in pairs(files) do
+      for _, file in ipairs(section) do
+        if file.path == item_name then
+          file.selected = true
+          file_was_selected = true
+          break
+        end
+      end
+      if file_was_selected then
+        break
+      end
+    end
+  end
+
+  -- If no specific file was requested or found, select the first file as a fallback.
+  if not file_was_selected then
+    for _, section_key in ipairs({ "conflicting", "working", "staged" }) do
+      if files[section_key] and #files[section_key] > 0 then
+        files[section_key][1].selected = true
+        break
+      end
+    end
+  end
+
+  -- Pass the initial list (with `selected` flags) and the update callback (without them)
   local view = CDiffView {
     git_root = git.repo.worktree_root,
     left = left,
@@ -86,9 +121,7 @@ local function get_local_diff_view(section_name, item_name, opts)
 
   view:on_files_staged(a.void(function(_)
     Watcher.instance():dispatch_refresh()
-    view:update_files()
   end))
-
   dv_lib.add_view(view)
 
   return view