fix: Stay at next hunk/file after staging

Hunk operations (stage/unstage/reset) stay at the same index position,
effectively moving to the next hunk. When a file is fully staged or
unstaged, jump to the next file of the same type rather than the first.

Author: MaxFangX

lua/vgit/features/screens/ProjectDiffScreen/init.lua

@@ -87,6 +87,38 @@ function ProjectDiffScreen:move_to(query_fn)
   return self.status_list_view:move_to(query_fn)
 end
 
+-- Find the next file of given entry_type after current filename
+function ProjectDiffScreen:find_next_file(filename, target_entry_type)
+  local next_filename = nil
+  local found_current = false
+  self.status_list_view:each_status(function(status, entry_type)
+    if entry_type == target_entry_type then
+      if found_current and not next_filename then
+        next_filename = status.filename
+      end
+      if status.filename == filename then
+        found_current = true
+      end
+    end
+  end)
+  return next_filename
+end
+
+-- Restore cursor to same hunk index after staging/unstaging/resetting a hunk
+function ProjectDiffScreen:restore_hunk_position(filename, entry_type, hunk_index)
+  local new_entry = self.model:get_entry()
+  if new_entry
+      and new_entry.type == entry_type
+      and new_entry.status.filename == filename then
+    local diff = self.model:get_diff()
+    if diff and diff.marks and #diff.marks > 0 then
+      local target = math.min(hunk_index, #diff.marks)
+      local hunk_alignment = project_diff_preview_setting:get('hunk_alignment')
+      self.diff_view:move_to_hunk(target, hunk_alignment)
+    end
+  end
+end
+
 function ProjectDiffScreen:stage_hunk()
   local entry = self.model:get_entry()
   if not entry then return end
@@ -97,6 +129,8 @@ function ProjectDiffScreen:stage_hunk()
   if not hunk then return end
 
   local filename = entry.status.filename
+  local next_file = self:find_next_file(filename, 'unstaged')
+
   local _, err = self.model:stage_hunk(filename, hunk)
   if err then
     console.debug.error(err)
@@ -112,35 +146,21 @@ function ProjectDiffScreen:stage_hunk()
     end)
 
     if has_unstaged then
-      -- Stay on the unstaged entry for this file
       self:move_to(function(status, entry_type)
         return status.filename == filename and entry_type == 'unstaged'
       end)
+    elseif next_file then
+      self:move_to(function(status, entry_type)
+        return status.filename == next_file and entry_type == 'unstaged'
+      end)
     else
-      -- File fully staged - jump to next unstaged file, else this file's staged
-      local found = self:move_to(function(_, entry_type)
-        return entry_type == 'unstaged'
+      self:move_to(function(status)
+        return status.filename == filename
       end)
-      if not found then
-        self:move_to(function(status)
-          return status.filename == filename
-        end)
-      end
     end
   end)
 
-  -- If still on the same file, stay at same hunk index position
-  local new_entry = self.model:get_entry()
-  if new_entry
-      and new_entry.type == 'unstaged'
-      and new_entry.status.filename == filename then
-    local diff = self.model:get_diff()
-    if diff and diff.marks and #diff.marks > 0 then
-      local target = math.min(hunk_index, #diff.marks)
-      local hunk_alignment = project_diff_preview_setting:get('hunk_alignment')
-      self.diff_view:move_to_hunk(target, hunk_alignment)
-    end
-  end
+  self:restore_hunk_position(filename, 'unstaged', hunk_index)
 end
 
 function ProjectDiffScreen:unstage_hunk()
@@ -153,6 +173,8 @@ function ProjectDiffScreen:unstage_hunk()
   if not hunk then return end
 
   local filename = entry.status.filename
+  local next_file = self:find_next_file(filename, 'staged')
+
   local _, err = self.model:unstage_hunk(filename, hunk)
   if err then
     console.debug.error(err)
@@ -168,35 +190,21 @@ function ProjectDiffScreen:unstage_hunk()
     end)
 
     if has_staged then
-      -- Stay on the staged entry for this file
       self:move_to(function(status, entry_type)
         return status.filename == filename and entry_type == 'staged'
       end)
+    elseif next_file then
+      self:move_to(function(status, entry_type)
+        return status.filename == next_file and entry_type == 'staged'
+      end)
     else
-      -- File fully unstaged - jump to next staged file, else this file's unstaged
-      local found = self:move_to(function(_, entry_type)
-        return entry_type == 'staged'
+      self:move_to(function(status)
+        return status.filename == filename
       end)
-      if not found then
-        self:move_to(function(status)
-          return status.filename == filename
-        end)
-      end
     end
   end)
 
-  -- If still on the same file, stay at same hunk index position
-  local new_entry = self.model:get_entry()
-  if new_entry
-      and new_entry.type == 'staged'
-      and new_entry.status.filename == filename then
-    local diff = self.model:get_diff()
-    if diff and diff.marks and #diff.marks > 0 then
-      local target = math.min(hunk_index, #diff.marks)
-      local hunk_alignment = project_diff_preview_setting:get('hunk_alignment')
-      self.diff_view:move_to_hunk(target, hunk_alignment)
-    end
-  end
+  self:restore_hunk_position(filename, 'staged', hunk_index)
 end
 
 function ProjectDiffScreen:reset_hunk()
@@ -209,6 +217,8 @@ function ProjectDiffScreen:reset_hunk()
   if not hunk then return end
 
   local filename = entry.status.filename
+  local next_file = self:find_next_file(filename, 'unstaged')
+
   loop.free_textlock()
   local decision = console.input('Are you sure you want to discard this hunk? (y/N) '):lower()
   if decision ~= 'yes' and decision ~= 'y' then return end
@@ -233,30 +243,18 @@ function ProjectDiffScreen:reset_hunk()
       self:move_to(function(status, entry_type)
         return status.filename == filename and entry_type == 'unstaged'
       end)
+    elseif next_file then
+      self:move_to(function(status, entry_type)
+        return status.filename == next_file and entry_type == 'unstaged'
+      end)
     else
-      local found = self:move_to(function(_, entry_type)
-        return entry_type == 'unstaged'
+      self:move_to(function(status)
+        return status.filename == filename
       end)
-      if not found then
-        self:move_to(function(status)
-          return status.filename == filename
-        end)
-      end
     end
   end)
 
-  -- If still on the same file, stay at same hunk index position
-  local new_entry = self.model:get_entry()
-  if new_entry
-      and new_entry.type == 'unstaged'
-      and new_entry.status.filename == filename then
-    local diff = self.model:get_diff()
-    if diff and diff.marks and #diff.marks > 0 then
-      local target = math.min(hunk_index, #diff.marks)
-      local hunk_alignment = project_diff_preview_setting:get('hunk_alignment')
-      self.diff_view:move_to_hunk(target, hunk_alignment)
-    end
-  end
+  self:restore_hunk_position(filename, 'unstaged', hunk_index)
 end
 
 function ProjectDiffScreen:stage_file()