finished most of the matching gui

1 files changed, 15 insertions, 4 deletions

diff --git a/src/cuchaz/enigma/gui/ClassSelector.java b/src/cuchaz/enigma/gui/ClassSelector.java
index 654bfbe..e5f550b 100644
--- a/src/cuchaz/enigma/gui/ClassSelector.java
+++ b/src/cuchaz/enigma/gui/ClassSelector.java
@@ -41,21 +41,32 @@ public class ClassSelector extends JTree {
         public static Comparator<ClassEntry> ObfuscatedClassEntryComparator;
         public static Comparator<ClassEntry> DeobfuscatedClassEntryComparator;
         
+        private static String getClassEntryDisplayName(ClassEntry entry) {
+                if (entry instanceof DecoratedClassEntry) {
+                        return ((DecoratedClassEntry)entry).getDecoration() + entry.getName();
+                }
+                return entry.getName();
+        }
+        
         static {
                 ObfuscatedClassEntryComparator = new Comparator<ClassEntry>() {
                         @Override
                         public int compare(ClassEntry a, ClassEntry b) {
-                                if (a.getName().length() != b.getName().length()) {
-                                        return a.getName().length() - b.getName().length();
+                                String aname = getClassEntryDisplayName(a);
+                                String bname = getClassEntryDisplayName(b);
+                                if (aname.length() != bname.length()) {
+                                        return aname.length() - bname.length();
                                 }
-                                return a.getName().compareTo(b.getName());
+                                return aname.compareTo(bname);
                         }
                 };
                 
                 DeobfuscatedClassEntryComparator = new Comparator<ClassEntry>() {
                         @Override
                         public int compare(ClassEntry a, ClassEntry b) {
-                                return a.getName().compareTo(b.getName());
+                                String aname = getClassEntryDisplayName(a);
+                                String bname = getClassEntryDisplayName(b);
+                                return aname.compareTo(bname);
                         }
                 };
         }