From 0381ce2fa0331756372ef0d94bc1bb6938bef72b Mon Sep 17 00:00:00 2001
From: Jamie Mansfield <dev@jamierocks.uk>
Date: Sun, 11 Sep 2016 12:57:49 +0100
Subject: [PATCH] Prevent the tree list collapsing every rename

---
 build.gradle                                  |  2 +-
 .../lapis/nocturne/gui/MainController.java    | 31 +++++++++++++------
 2 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/build.gradle b/build.gradle
index 5bdefa5..ca827d2 100644
--- a/build.gradle
+++ b/build.gradle
@@ -16,7 +16,7 @@ defaultTasks 'clean', 'licenseFormat', 'build'
 // Project information
 allprojects {
     group = 'blue.lapis.nocturne'
-    version = '1.0.3'
+    version = '1.0.4'
 
     sourceCompatibility = '1.8'
     targetCompatibility = '1.8'
diff --git a/src/main/java/blue/lapis/nocturne/gui/MainController.java b/src/main/java/blue/lapis/nocturne/gui/MainController.java
index 8b4be0a..fd64b60 100644
--- a/src/main/java/blue/lapis/nocturne/gui/MainController.java
+++ b/src/main/java/blue/lapis/nocturne/gui/MainController.java
@@ -41,13 +41,11 @@
 import blue.lapis.nocturne.util.Constants;
 import blue.lapis.nocturne.util.helper.PropertiesHelper;
 import blue.lapis.nocturne.util.helper.SceneHelper;
-
 import javafx.event.ActionEvent;
 import javafx.fxml.FXMLLoader;
 import javafx.fxml.Initializable;
 import javafx.scene.Node;
 import javafx.scene.control.Alert;
-import javafx.scene.control.Dialog;
 import javafx.scene.control.MenuItem;
 import javafx.scene.control.RadioMenuItem;
 import javafx.scene.control.TabPane;
@@ -67,7 +65,6 @@
 import java.util.Optional;
 import java.util.ResourceBundle;
 import java.util.function.BiConsumer;
-import java.util.stream.Collectors;
 
 /**
  * The main JavaFX controller.
@@ -266,7 +263,7 @@ public void onLanguageSelect(ActionEvent actionEvent) throws IOException {
 
     public void updateObfuscatedClassListView() {
         if (Main.getLoadedJar() != null) {
-            TreeItem<String> root = generateTreeItem(Main.getLoadedJar().getObfuscatedHierarchy());
+            TreeItem<String> root = generateTreeItem(Main.getLoadedJar().getObfuscatedHierarchy(), obfTree.getRoot());
             root.setExpanded(true);
             obfTree.setRoot(root);
         } else {
@@ -277,7 +274,7 @@ public void updateObfuscatedClassListView() {
 
     public void updateDeobfuscatedClassListView() {
         if (Main.getLoadedJar() != null) {
-            TreeItem<String> root = generateTreeItem(Main.getLoadedJar().getDeobfuscatedHierarchy());
+            TreeItem<String> root = generateTreeItem(Main.getLoadedJar().getDeobfuscatedHierarchy(), deobfTree.getRoot());
             root.setExpanded(true);
             deobfTree.setRoot(root);
         } else {
@@ -285,7 +282,7 @@ public void updateDeobfuscatedClassListView() {
         }
     }
 
-    public TreeItem<String> generateTreeItem(HierarchyElement element) {
+    public TreeItem<String> generateTreeItem(HierarchyElement element, TreeItem<String> oldTreeItem) {
         TreeItem<String> treeItem;
         if (element instanceof HierarchyNode) {
             HierarchyNode node = (HierarchyNode) element;
@@ -297,11 +294,27 @@ public TreeItem<String> generateTreeItem(HierarchyElement element) {
         } else {
             treeItem = new TreeItem<>("(root)");
         }
+        if (oldTreeItem != null) {
+            treeItem.setExpanded(oldTreeItem.isExpanded());
+        }
         if (element instanceof Hierarchy
                 || (element instanceof HierarchyNode && !((HierarchyNode) element).isTerminal())) {
-            treeItem.getChildren().addAll(
-                    element.getChildren().stream().map(this::generateTreeItem).collect(Collectors.toList())
-            );
+            for (HierarchyNode node : element.getChildren()) {
+                if (oldTreeItem != null) {
+                    boolean added = false;
+                    for (TreeItem<String> child : oldTreeItem.getChildren()) {
+                        if (node.getDisplayName().equalsIgnoreCase(child.getValue())) {
+                            treeItem.getChildren().add(this.generateTreeItem(node, child));
+                            added = true;
+                        }
+                    }
+                    if (!added) {
+                        treeItem.getChildren().add(this.generateTreeItem(node, null));
+                    }
+                } else {
+                    treeItem.getChildren().add(this.generateTreeItem(node, null));
+                }
+            }
         }
         treeItem.getChildren().setAll(treeItem.getChildren().sorted((t1, t2) -> {
             boolean c1 = t1.getChildren().size() > 0;