3 files changed, 47 insertions, 20 deletions

diff --git a/src/main/java/cuchaz/enigma/Constants.java b/src/main/java/cuchaz/enigma/Constants.java
index b08438a3..afffb5d3 100644
--- a/src/main/java/cuchaz/enigma/Constants.java
+++ b/src/main/java/cuchaz/enigma/Constants.java
@@ -13,8 +13,8 @@ package cuchaz.enigma;
 
 public class Constants {
         public static final String NAME = "Enigma";
-        public static final String VERSION = "0.11.0 (Fabric Fork)";
-        public static final String URL = "http://www.cuchazinteractive.com/enigma";
+        public static final String VERSION = "@VERSION@/Fabric";
+        public static final String URL = "https://fabric.asie.pl";
         public static final int MiB = 1024 * 1024; // 1 mebibyte
         public static final int KiB = 1024; // 1 kebibyte
 }
diff --git a/src/main/java/cuchaz/enigma/gui/ClassSelector.java b/src/main/java/cuchaz/enigma/gui/ClassSelector.java
index 93c5d4ba..b032526f 100644
--- a/src/main/java/cuchaz/enigma/gui/ClassSelector.java
+++ b/src/main/java/cuchaz/enigma/gui/ClassSelector.java
@@ -141,7 +141,7 @@ public class ClassSelector extends JTree {
         }
 
         public void setClasses(Collection<ClassEntry> classEntries) {
-                String state = getExpansionState(this, 0);
+                List<StateEntry> state = getExpansionState(this);
                 if (classEntries == null) {
                         setModel(null);
                         return;
@@ -210,7 +210,7 @@ public class ClassSelector extends JTree {
                 // finally, update the tree control
                 setModel(new DefaultTreeModel(rootNodes));
 
-                restoreExpanstionState(this, 0, state);
+                restoreExpansionState(this, state);
         }
 
         public ClassEntry getSelectedClass() {
@@ -251,28 +251,48 @@ public class ClassSelector extends JTree {
                 return path1.equals(path2);
         }
 
-        public String getExpansionState(JTree tree, int row) {
-                TreePath rowPath = tree.getPathForRow(row);
-                StringBuilder buf = new StringBuilder();
+        public enum State {
+                EXPANDED,
+                SELECTED
+        }
+
+        public static class StateEntry {
+                public final State state;
+                public final TreePath path;
+
+                public StateEntry(State state, TreePath path) {
+                        this.state = state;
+                        this.path = path;
+                }
+        }
+
+        public List<StateEntry> getExpansionState(JTree tree) {
+                List<StateEntry> state = new ArrayList<>();
                 int rowCount = tree.getRowCount();
-                for (int i = row; i < rowCount; i++) {
+                for (int i = 0; i < rowCount; i++) {
                         TreePath path = tree.getPathForRow(i);
-                        if (i == row || isDescendant(path, rowPath)) {
-                                if (tree.isExpanded(path)) {
-                                        buf.append(",").append((i - row));
-                                }
-                        } else {
-                                break;
+                        if (tree.isPathSelected(path)) {
+                                state.add(new StateEntry(State.SELECTED, path));
+                        }
+                        if (tree.isExpanded(path)) {
+                                state.add(new StateEntry(State.EXPANDED, path));
                         }
                 }
-                return buf.toString();
+                return state;
         }
 
-        public void restoreExpanstionState(JTree tree, int row, String expansionState) {
-                StringTokenizer stok = new StringTokenizer(expansionState, ",");
-                while (stok.hasMoreTokens()) {
-                        int token = row + Integer.parseInt(stok.nextToken());
-                        tree.expandRow(token);
+        public void restoreExpansionState(JTree tree, List<StateEntry> expansionState) {
+                tree.clearSelection();
+
+                for (StateEntry entry : expansionState) {
+                        switch (entry.state) {
+                                case SELECTED:
+                                        tree.addSelectionPath(entry.path);
+                                        break;
+                                case EXPANDED:
+                                        tree.expandPath(entry.path);
+                                        break;
+                        }
                 }
         }
 
diff --git a/src/main/java/cuchaz/enigma/gui/Gui.java b/src/main/java/cuchaz/enigma/gui/Gui.java
index 1fc8bf34..71bf5dcc 100644
--- a/src/main/java/cuchaz/enigma/gui/Gui.java
+++ b/src/main/java/cuchaz/enigma/gui/Gui.java
@@ -792,11 +792,15 @@ public class Gui {
                         newEntry == null);
         }
 
+        // TODO: getExpansionState will *not* actually update itself based on name changes!
         public void moveClassTree(EntryReference<Entry, Entry> deobfReference, String newName, boolean isOldOb, boolean isNewOb) {
                 ClassEntry oldEntry = deobfReference.entry.getOwnerClassEntry();
                 ClassEntry newEntry = new ClassEntry(newName);
 
                 // Ob -> deob
+                List<ClassSelector.StateEntry> stateDeobf = this.deobfPanel.deobfClasses.getExpansionState(this.deobfPanel.deobfClasses);
+                List<ClassSelector.StateEntry> stateObf = this.obfPanel.obfClasses.getExpansionState(this.obfPanel.obfClasses);
+
                 if (isOldOb && !isNewOb) {
                         this.deobfPanel.deobfClasses.moveClassTree(oldEntry, newEntry, obfPanel.obfClasses);
                         ClassSelectorPackageNode packageNode = this.obfPanel.obfClasses.getPackageNode(oldEntry);
@@ -824,5 +828,8 @@ public class Gui {
                         this.deobfPanel.deobfClasses.removeNodeIfEmpty(this.deobfPanel.deobfClasses.getPackageNode(oldEntry));
                         this.deobfPanel.deobfClasses.reload();
                 }
+
+                this.deobfPanel.deobfClasses.restoreExpansionState(this.deobfPanel.deobfClasses, stateDeobf);
+                this.obfPanel.obfClasses.restoreExpansionState(this.obfPanel.obfClasses, stateObf);
         }
 }