3 files changed, 152 insertions, 18 deletions

diff --git a/src/main/java/cuchaz/enigma/gui/ClassSelector.java b/src/main/java/cuchaz/enigma/gui/ClassSelector.java
index ed0f6c8d..98880cd6 100644
--- a/src/main/java/cuchaz/enigma/gui/ClassSelector.java
+++ b/src/main/java/cuchaz/enigma/gui/ClassSelector.java
@@ -22,7 +22,8 @@ import cuchaz.enigma.throwables.IllegalNameException;
 import javax.swing.*;
 import javax.swing.event.CellEditorListener;
 import javax.swing.event.ChangeEvent;
-import javax.swing.event.TreeSelectionEvent;
+import javax.swing.event.TreeModelEvent;
+import javax.swing.event.TreeModelListener;
 import javax.swing.tree.*;
 import java.awt.event.MouseAdapter;
 import java.awt.event.MouseEvent;
@@ -294,7 +295,7 @@ public class ClassSelector extends JTree {
         }
     }
 
-    public Iterable<ClassSelectorPackageNode> packageNodes() {
+    public List<ClassSelectorPackageNode> packageNodes() {
         List<ClassSelectorPackageNode> nodes = Lists.newArrayList();
         DefaultMutableTreeNode root = (DefaultMutableTreeNode)getModel().getRoot();
         Enumeration<?> children = root.children();
@@ -305,7 +306,7 @@ public class ClassSelector extends JTree {
         return nodes;
     }
 
-    public Iterable<ClassSelectorClassNode> classNodes(ClassSelectorPackageNode packageNode) {
+    public List<ClassSelectorClassNode> classNodes(ClassSelectorPackageNode packageNode) {
         List<ClassSelectorClassNode> nodes = Lists.newArrayList();
         Enumeration<?> children = packageNode.children();
         while (children.hasMoreElements()) {
@@ -351,6 +352,15 @@ public class ClassSelector extends JTree {
         return null;
     }
 
+    public ClassSelectorPackageNode getPackageNode(ClassSelector selector, ClassEntry entry)
+    {
+        ClassSelectorPackageNode packageNode = getPackageNode(entry);
+
+        if (selector != null && packageNode == null && selector.getPackageNode(entry) != null)
+                return selector.getPackageNode(entry);
+        return packageNode;
+    }
+
     public ClassEntry getNextClass(ClassEntry entry) {
         boolean foundIt = false;
         for (ClassSelectorPackageNode packageNode : packageNodes()) {
@@ -388,4 +398,82 @@ public class ClassSelector extends JTree {
             }
         }
     }
+
+    public void removeNode(ClassSelectorPackageNode packageNode, ClassEntry entry)
+    {
+        DefaultTreeModel model = (DefaultTreeModel) getModel();
+
+        if (packageNode == null)
+            return;
+
+        for (int i = 0; i < packageNode.getChildCount(); i++)
+        {
+            DefaultMutableTreeNode childNode = (DefaultMutableTreeNode) packageNode.getChildAt(i);
+            if (childNode.getUserObject() instanceof ClassEntry && childNode.getUserObject().equals(entry))
+            {
+                model.removeNodeFromParent(childNode);
+                break;
+            }
+        }
+    }
+
+    public void removeNodeIfEmpty(ClassSelectorPackageNode packageNode)
+    {
+        if (packageNode != null && packageNode.getChildCount() == 0)
+            ((DefaultTreeModel) getModel()).removeNodeFromParent(packageNode);
+    }
+
+    public void moveClassTree(ClassEntry oldClassEntry, ClassEntry newClassEntry, ClassSelector otherSelector)
+    {
+        if (otherSelector == null)
+            removeNode(getPackageNode(oldClassEntry), oldClassEntry);
+        insertNode(getOrCreate(newClassEntry), newClassEntry);
+    }
+
+    public ClassSelectorPackageNode getOrCreate(ClassEntry entry)
+    {
+        DefaultTreeModel model = (DefaultTreeModel) getModel();
+        ClassSelectorPackageNode newPackageNode = getPackageNode(entry);
+        if (newPackageNode == null)
+        {
+            newPackageNode = new ClassSelectorPackageNode(entry.getPackageName());
+            model.insertNodeInto(newPackageNode, (MutableTreeNode) model.getRoot(), getPlacementIndex(newPackageNode));
+        }
+        return newPackageNode;
+    }
+
+    public void insertNode(ClassSelectorPackageNode packageNode, ClassEntry entry)
+    {
+        DefaultTreeModel model = (DefaultTreeModel) getModel();
+        ClassSelectorClassNode classNode = new ClassSelectorClassNode(entry);
+        model.insertNodeInto(classNode, packageNode, getPlacementIndex(packageNode, classNode));
+    }
+
+    private int getPlacementIndex(ClassSelectorPackageNode newPackageNode, ClassSelectorClassNode classNode)
+    {
+        List<ClassSelectorClassNode> classNodes = classNodes(newPackageNode);
+        classNodes.add(classNode);
+        Collections.sort(classNodes, (a, b) -> a.toString().compareTo(b.toString()));
+        for (int i = 0; i < classNodes.size(); i++)
+            if (classNodes.get(i) == classNode)
+                return i;
+
+        return 0;
+    }
+
+    private int getPlacementIndex(ClassSelectorPackageNode newPackageNode)
+    {
+        List<ClassSelectorPackageNode> packageNodes = packageNodes();
+        if (!packageNodes.contains(newPackageNode))
+        {
+            packageNodes.add(newPackageNode);
+            Collections.sort(packageNodes, (a, b) -> a.toString().compareTo(b.toString()));
+        }
+
+        for (int i = 0; i < packageNodes.size(); i++)
+            if (packageNodes.get(i) == newPackageNode)
+                return i;
+
+        return 0;
+    }
 }
diff --git a/src/main/java/cuchaz/enigma/gui/Gui.java b/src/main/java/cuchaz/enigma/gui/Gui.java
index bc49a3ff..9fd824b8 100644
--- a/src/main/java/cuchaz/enigma/gui/Gui.java
+++ b/src/main/java/cuchaz/enigma/gui/Gui.java
@@ -24,6 +24,7 @@ import cuchaz.enigma.gui.highlight.DeobfuscatedHighlightPainter;
 import cuchaz.enigma.gui.highlight.ObfuscatedHighlightPainter;
 import cuchaz.enigma.gui.highlight.OtherHighlightPainter;
 import cuchaz.enigma.gui.highlight.SelectionHighlightPainter;
+import cuchaz.enigma.gui.node.ClassSelectorPackageNode;
 import cuchaz.enigma.gui.panels.PanelDeobf;
 import cuchaz.enigma.gui.panels.PanelEditor;
 import cuchaz.enigma.gui.panels.PanelIdentifier;
@@ -49,7 +50,6 @@ import java.util.Collections;
 import java.util.List;
 import java.util.Vector;
 import java.util.function.Function;
-import java.util.function.Supplier;
 
 public class Gui {
 
@@ -762,11 +762,53 @@ public class Gui {
                 DefaultMutableTreeNode childNode = (DefaultMutableTreeNode) node.getChildAt(i);
                 ClassEntry prevDataChild = (ClassEntry) childNode.getUserObject();
                 ClassEntry dataChild = new ClassEntry(data + "/" + prevDataChild.getSimpleName());
-                this.controller.rename(new EntryReference<>(prevDataChild, prevDataChild.getName()), dataChild.getName());
+                this.controller.rename(new EntryReference<>(prevDataChild, prevDataChild.getName()), dataChild.getName(), false);
             }
         }
         // class rename
         else if (data instanceof ClassEntry)
-            this.controller.rename(new EntryReference<>((ClassEntry) prevData, ((ClassEntry) prevData).getName()), ((ClassEntry) data).getName());
+            this.controller.rename(new EntryReference<>((ClassEntry) prevData, ((ClassEntry) prevData).getName()), ((ClassEntry) data).getName(), false);
+    }
+
+    public void moveClassTree(EntryReference<Entry, Entry> deobfReference, String newName)
+    {
+        ClassEntry oldEntry = deobfReference.entry.getClassEntry();
+        ClassEntry newEntry = new ClassEntry(newName);
+        moveClassTree(deobfReference, newName, oldEntry.getPackageName().equals(Constants.NONE_PACKAGE),
+                newEntry.getClassEntry().getPackageName().equals(Constants.NONE_PACKAGE));
+    }
+
+    public void moveClassTree(EntryReference<Entry, Entry> deobfReference, String newName, boolean isOldOb, boolean isNewOb)
+    {
+        ClassEntry oldEntry = deobfReference.entry.getClassEntry();
+        ClassEntry newEntry = new ClassEntry(newName);
+
+        // Ob -> deob
+        if (isOldOb && !isNewOb)
+        {
+            this.deobfPanel.deobfClasses.moveClassTree(oldEntry, newEntry, obfPanel.obfClasses);
+            ClassSelectorPackageNode packageNode = this.obfPanel.obfClasses.getPackageNode(oldEntry);
+            this.obfPanel.obfClasses.removeNode(packageNode, oldEntry);
+            this.obfPanel.obfClasses.removeNodeIfEmpty(packageNode);
+        }
+        // Deob -> ob
+        else if (isNewOb && !isOldOb)
+        {
+            this.obfPanel.obfClasses.moveClassTree(oldEntry, newEntry, deobfPanel.deobfClasses);
+            ClassSelectorPackageNode packageNode = this.deobfPanel.deobfClasses.getPackageNode(oldEntry);
+            this.deobfPanel.deobfClasses.removeNode(packageNode, oldEntry);
+            this.deobfPanel.deobfClasses.removeNodeIfEmpty(packageNode);
+        }
+        // Local move
+        else if (isOldOb)
+        {
+            this.obfPanel.obfClasses.moveClassTree(oldEntry, newEntry, null);
+            this.obfPanel.obfClasses.removeNodeIfEmpty(this.obfPanel.obfClasses.getPackageNode(oldEntry));
+        }
+        else
+        {
+            this.deobfPanel.deobfClasses.moveClassTree(oldEntry, newEntry, null);
+            this.deobfPanel.deobfClasses.removeNodeIfEmpty(this.deobfPanel.deobfClasses.getPackageNode(oldEntry));
+        }
     }
 }
diff --git a/src/main/java/cuchaz/enigma/gui/GuiController.java b/src/main/java/cuchaz/enigma/gui/GuiController.java
index cc645d8a..e466f637 100644
--- a/src/main/java/cuchaz/enigma/gui/GuiController.java
+++ b/src/main/java/cuchaz/enigma/gui/GuiController.java
@@ -12,24 +12,21 @@ package cuchaz.enigma.gui;
 
 import com.google.common.collect.Lists;
 import com.google.common.collect.Queues;
-
 import com.strobel.decompiler.languages.java.ast.CompilationUnit;
+import cuchaz.enigma.Deobfuscator;
+import cuchaz.enigma.analysis.*;
+import cuchaz.enigma.gui.dialog.ProgressDialog;
+import cuchaz.enigma.mapping.*;
+import cuchaz.enigma.throwables.MappingParseException;
+import cuchaz.enigma.utils.ReadableToken;
 
 import java.io.File;
-import java.io.FileWriter;
 import java.io.IOException;
 import java.util.Collection;
 import java.util.Deque;
 import java.util.List;
 import java.util.jar.JarFile;
 
-import cuchaz.enigma.Deobfuscator;
-import cuchaz.enigma.analysis.*;
-import cuchaz.enigma.gui.dialog.ProgressDialog;
-import cuchaz.enigma.mapping.*;
-import cuchaz.enigma.throwables.MappingParseException;
-import cuchaz.enigma.utils.ReadableToken;
-
 public class GuiController {
 
     private Deobfuscator deobfuscator;
@@ -192,18 +189,26 @@ public class GuiController {
     }
 
     public void rename(EntryReference<Entry, Entry> deobfReference, String newName) {
+        rename(deobfReference, newName, true);
+    }
+
+    public void rename(EntryReference<Entry, Entry> deobfReference, String newName, boolean refreshClassTree)
+    {
         EntryReference<Entry, Entry> obfReference = this.deobfuscator.obfuscateReference(deobfReference);
         this.deobfuscator.rename(obfReference.getNameableEntry(), newName);
         this.isDirty = true;
-        refreshClasses();
+
+        if (refreshClassTree)
+            this.gui.moveClassTree(deobfReference, newName);
         refreshCurrentClass(obfReference);
+
     }
 
     public void removeMapping(EntryReference<Entry, Entry> deobfReference) {
         EntryReference<Entry, Entry> obfReference = this.deobfuscator.obfuscateReference(deobfReference);
         this.deobfuscator.removeMapping(obfReference.getNameableEntry());
         this.isDirty = true;
-        refreshClasses();
+        this.gui.moveClassTree(deobfReference, obfReference.entry.getName(), false, true);
         refreshCurrentClass(obfReference);
     }
 
@@ -211,7 +216,6 @@ public class GuiController {
         EntryReference<Entry, Entry> obfReference = this.deobfuscator.obfuscateReference(deobfReference);
         this.deobfuscator.markAsDeobfuscated(obfReference.getNameableEntry());
         this.isDirty = true;
-        refreshClasses();
         refreshCurrentClass(obfReference);
     }