From 42942957fe4e0af4f08c0bd2fd7d50ce325be8f8 Mon Sep 17 00:00:00 2001
From: Stephan Dilly <dilly.stephan@gmail.com>
Date: Fri, 12 Mar 2021 23:41:52 +0100
Subject: [PATCH] keep line selection after staging/unstagin/discard (closes
 #583)

---
 src/components/diff.rs | 30 +++++++++++++++++++++---------
 1 file changed, 21 insertions(+), 9 deletions(-)

diff --git a/src/components/diff.rs b/src/components/diff.rs
index 8cbfa59619..76069eb7eb 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 {
                 }
             };
 
-            let new_start = cmp::min(max, new_start);
+            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);
         }