From 07be7d7d0fa37a79b473ae1ec13e82cac89091fc Mon Sep 17 00:00:00 2001
From: Hannes Braun <hannes@hannesbraun.net>
Date: Fri, 6 Sep 2024 15:47:58 +0200
Subject: [PATCH] Automatically select next change after (un)stage

---
 src/ViewModels/WorkingCopy.cs | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/src/ViewModels/WorkingCopy.cs b/src/ViewModels/WorkingCopy.cs
index 92330eef..c116ca1c 100644
--- a/src/ViewModels/WorkingCopy.cs
+++ b/src/ViewModels/WorkingCopy.cs
@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using System.IO;
+using System.Linq;
 using System.Threading.Tasks;
 
 using Avalonia.Controls;
@@ -322,8 +323,9 @@ namespace SourceGit.ViewModels
 
         public void StageSelected()
         {
+            var nextSelection = CalculateNextSelection(_selectedUnstaged, Unstaged);
             StageChanges(_selectedUnstaged);
-            SelectedUnstaged = [];
+            SelectedUnstaged = nextSelection;
         }
 
         public void StageAll()
@@ -359,8 +361,9 @@ namespace SourceGit.ViewModels
 
         public void UnstageSelected()
         {
+            var nextSelection = CalculateNextSelection(_selectedStaged, Staged);
             UnstageChanges(_selectedStaged);
-            SelectedStaged = [];
+            SelectedStaged = nextSelection;
         }
 
         public void UnstageAll()
@@ -1357,6 +1360,19 @@ namespace SourceGit.ViewModels
             return false;
         }
 
+        private static List<Models.Change> CalculateNextSelection(List<Models.Change> selected, List<Models.Change> all)
+        {
+            var indices = selected.Select(x => all.IndexOf(x)).ToArray();
+            var min = indices.Min();
+            var next = min + 1;
+            while (indices.Contains(next))
+                next++;
+            if (next >= all.Count)
+                next = min - 1;
+            List<Models.Change> nextSelection = next >= 0 ? [all[next]] : [];
+            return nextSelection;
+        }
+
         private Repository _repo = null;
         private bool _isLoadingData = false;
         private bool _isStaging = false;