From b80df36cdff98323940591f21fb488618fe8a271 Mon Sep 17 00:00:00 2001
From: Stephan Dilly <dilly.stephan@gmail.com>
Date: Fri, 12 Jun 2020 16:39:06 +0200
Subject: [PATCH] fix selection disappearing in tree sometimes (closes #120)

---
 CHANGELOG.md                       |  1 +
 src/components/utils/statustree.rs | 41 ++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 37d1ae1d7d..94522b09a3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 - reset file inside folder failed when running `gitui` in a subfolder too ([#118](https://github.com/extrawurst/gitui/issues/118))
+- selection could disappear into collapsed folder ([#120](https://github.com/extrawurst/gitui/issues/120))
 
 ## [0.6.0] - 2020-06-09
 
diff --git a/src/components/utils/statustree.rs b/src/components/utils/statustree.rs
index c2462d5a6e..04caaef69d 100644
--- a/src/components/utils/statustree.rs
+++ b/src/components/utils/statustree.rs
@@ -58,9 +58,26 @@ impl StatusTree {
 
         self.update_visibility(None, 0, true);
 
+        //NOTE: now that visibility is set we can make sure selection is visible
+        if let Some(idx) = self.selection {
+            self.selection = Some(self.find_visible_idx(idx));
+        }
+
         Ok(())
     }
 
+    fn find_visible_idx(&self, mut idx: usize) -> usize {
+        while idx > 0 {
+            if self.is_visible_index(idx) {
+                break;
+            }
+
+            idx -= 1;
+        }
+
+        idx
+    }
+
     ///
     pub fn move_selection(&mut self, dir: MoveSelection) -> bool {
         if let Some(selection) = self.selection {
@@ -394,6 +411,30 @@ mod tests {
         assert_eq!(res.selection, Some(1));
     }
 
+    #[test]
+    fn test_keep_selected_index_if_not_collapsed() {
+        let mut res = StatusTree::default();
+        res.update(&string_vec_to_status(&["a/b", "c"])).unwrap();
+
+        res.collapse("a/b", 0);
+
+        res.selection = Some(2);
+
+        res.update(&string_vec_to_status(&["a/b"])).unwrap();
+        assert_eq!(
+            get_visibles(&res),
+            vec![
+                true,  //
+                false, //
+            ]
+        );
+        assert_eq!(
+            res.is_visible_index(res.selection.unwrap()),
+            true
+        );
+        assert_eq!(res.selection, Some(0));
+    }
+
     #[test]
     fn test_keep_collapsed_states() {
         let mut res = StatusTree::default();