From 7937b80483123ff1ff37e0d65b71726b4d4726f8 Mon Sep 17 00:00:00 2001
From: Stephan Dilly <dilly.stephan@gmail.com>
Date: Fri, 12 Mar 2021 23:57:31 +0100
Subject: [PATCH] Fix 583 fix diff line selection (#585)

Preserve line selection after staging/unstaging/discard (closes #583)
---
 CHANGELOG.md           |  3 +++
 src/components/diff.rs | 30 +++++++++++++++++++++---------
 2 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 55dbb9ae..db894939 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,6 +17,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - support for pushing tags ([#568](https://github.com/extrawurst/gitui/issues/568))
 - visualize *conflicted* files differently ([#576](https://github.com/extrawurst/gitui/issues/576))
 
+### Fixed
+- keep diff line selection after staging/unstaging/discarding ([#583](https://github.com/extrawurst/gitui/issues/583))
+
 ## [0.12.0] - 2020-03-03
 
 **pull support (ff-merge or conflict-free merge-commit)**
diff --git a/src/components/diff.rs b/src/components/diff.rs
index 8cbfa596..76069eb7 100644
--- a/src/components/diff.rs
+++ b/src/components/diff.rs
@@ -169,20 +169,27 @@ impl DiffComponent {
         let hash = hash(&diff);
 
         if self.current.hash != hash {
+            let reset_selection = self.current.path != path;
+
             self.current = Current {
                 path,
                 is_stage,
                 hash,
             };
 
-            self.selected_hunk = Self::find_selected_hunk(
-                &diff,
-                self.selection.get_start(),
-            );
-
             self.diff = Some(diff);
-            self.scroll_top.set(0);
-            self.selection = Selection::Single(0);
+
+            if reset_selection {
+                self.scroll_top.set(0);
+                self.selection = Selection::Single(0);
+                self.update_selection(0);
+            } else {
+                let old_selection = match self.selection {
+                    Selection::Single(line) => line,
+                    Selection::Multiple(start, _) => start,
+                };
+                self.update_selection(old_selection);
+            }
         }
 
         Ok(())
@@ -215,10 +222,15 @@ impl DiffComponent {
                 }
             };
 
+            self.update_selection(new_start);
+        }
+    }
+
+    fn update_selection(&mut self, new_start: usize) {
+        if let Some(diff) = &self.diff {
+            let max = diff.lines.saturating_sub(1) as usize;
             let new_start = cmp::min(max, new_start);
-
             self.selection = Selection::Single(new_start);
-
             self.selected_hunk =
                 Self::find_selected_hunk(diff, new_start);
         }