From 90e3abed8d742d7ab8dbe27ffc2cbfa5375f05cb Mon Sep 17 00:00:00 2001
From: leo <longshuang@msn.cn>
Date: Tue, 6 Aug 2024 10:42:55 +0800
Subject: [PATCH] enhance: keep selection after unstaging (#320)

---
 src/ViewModels/WorkingCopy.cs | 39 +++++++++++++++++++++++++++++++----
 1 file changed, 35 insertions(+), 4 deletions(-)

diff --git a/src/ViewModels/WorkingCopy.cs b/src/ViewModels/WorkingCopy.cs
index 8f352fbe..08ee0b2a 100644
--- a/src/ViewModels/WorkingCopy.cs
+++ b/src/ViewModels/WorkingCopy.cs
@@ -218,7 +218,22 @@ namespace SourceGit.ViewModels
 
         public bool SetData(List<Models.Change> changes)
         {
+            if (!IsChanged(_cached, changes))
+            {
+                // Just force refresh selected changes.
+                Dispatcher.UIThread.Invoke(() =>
+                {
+                    if (_selectedUnstaged.Count > 0)
+                        SelectedUnstaged = new List<Models.Change>(_selectedUnstaged);
+                    else if (_selectedStaged.Count > 0)
+                        SelectedStaged = new List<Models.Change>(_selectedStaged);
+                });
+
+                return _cached.Find(x => x.IsConflit) != null;
+            }
+
             _cached = changes;
+            _count = _cached.Count;
 
             var unstaged = new List<Models.Change>();
             var staged = new List<Models.Change>();
@@ -227,12 +242,12 @@ namespace SourceGit.ViewModels
 
             var lastSelectedUnstaged = new HashSet<string>();
             var lastSelectedStaged = new HashSet<string>();
-            if (_selectedUnstaged != null)
+            if (_selectedUnstaged != null && _selectedUnstaged.Count > 0)
             {
                 foreach (var c in _selectedUnstaged)
                     lastSelectedUnstaged.Add(c.Path);
             }
-            else if (_selectedStaged != null)
+            else if (_selectedStaged != null && _selectedStaged.Count > 0)
             {
                 foreach (var c in _selectedStaged)
                     lastSelectedStaged.Add(c.Path);
@@ -258,8 +273,6 @@ namespace SourceGit.ViewModels
                     selectedStaged.Add(c);
             }
 
-            _count = changes.Count;
-
             Dispatcher.UIThread.Invoke(() =>
             {
                 _isLoadingData = true;
@@ -1273,6 +1286,24 @@ namespace SourceGit.ViewModels
             });
         }
 
+        private bool IsChanged(List<Models.Change> old, List<Models.Change> cur)
+        {
+            if (old.Count != cur.Count)
+                return true;
+
+            var oldSet = new HashSet<string>();
+            foreach (var c in old)
+                oldSet.Add($"{c.Path}\n{c.WorkTree}\n{c.Index}");
+
+            foreach (var c in cur)
+            {
+                if (!oldSet.Contains($"{c.Path}\n{c.WorkTree}\n{c.Index}"))
+                    return true;
+            }
+
+            return false;
+        }
+
         private Repository _repo = null;
         private bool _isLoadingData = false;
         private bool _isStaging = false;