2 files changed, 20 insertions, 0 deletions

diff --git a/src/main/java/cuchaz/enigma/gui/node/ClassSelectorClassNode.java b/src/main/java/cuchaz/enigma/gui/node/ClassSelectorClassNode.java
index e083572..1c9dad4 100644
--- a/src/main/java/cuchaz/enigma/gui/node/ClassSelectorClassNode.java
+++ b/src/main/java/cuchaz/enigma/gui/node/ClassSelectorClassNode.java
@@ -20,6 +20,7 @@ public class ClassSelectorClassNode extends DefaultMutableTreeNode {
 
     public ClassSelectorClassNode(ClassEntry classEntry) {
         this.classEntry = classEntry;
+        this.setUserObject(classEntry);
     }
 
     public ClassEntry getClassEntry() {
@@ -36,6 +37,18 @@ public class ClassSelectorClassNode extends DefaultMutableTreeNode {
         return other instanceof ClassSelectorClassNode && equals((ClassSelectorClassNode) other);
     }
 
+    @Override public void setUserObject(Object userObject)
+    {
+        String packageName = "";
+        if (classEntry.getPackageName() != null)
+            packageName = classEntry.getPackageName() + "/";
+        if (userObject instanceof String)
+            this.classEntry = new ClassEntry(packageName + userObject);
+        else if (userObject instanceof ClassEntry)
+            this.classEntry = (ClassEntry) userObject;
+        super.setUserObject(classEntry);
+    }
+
     public boolean equals(ClassSelectorClassNode other) {
         return this.classEntry.equals(other.classEntry);
     }
diff --git a/src/main/java/cuchaz/enigma/gui/node/ClassSelectorPackageNode.java b/src/main/java/cuchaz/enigma/gui/node/ClassSelectorPackageNode.java
index dfdc765..ee3009a 100644
--- a/src/main/java/cuchaz/enigma/gui/node/ClassSelectorPackageNode.java
+++ b/src/main/java/cuchaz/enigma/gui/node/ClassSelectorPackageNode.java
@@ -24,6 +24,13 @@ public class ClassSelectorPackageNode extends DefaultMutableTreeNode {
         return packageName;
     }
 
+    @Override public void setUserObject(Object userObject)
+    {
+        if (userObject instanceof String)
+            this.packageName = (String) userObject;
+        super.setUserObject(userObject);
+    }
+
     @Override
     public String toString() {
         return this.packageName;