made obfuscated/deobfuscated class selector a bit easier to use

6 files changed, 282 insertions, 113 deletions

diff --git a/src/cuchaz/enigma/Deobfuscator.java b/src/cuchaz/enigma/Deobfuscator.java
index 9a78f38..49aa1ff 100644
--- a/src/cuchaz/enigma/Deobfuscator.java
+++ b/src/cuchaz/enigma/Deobfuscator.java
@@ -151,12 +151,12 @@ public class Deobfuscator
                 return m_mappings.getTranslator( m_jarIndex.getAncestries(), direction );
         }
         
-        public void getSeparatedClasses( List<String> obfClasses, List<String> deobfClasses )
+        public void getSeparatedClasses( List<ClassEntry> obfClasses, List<ClassEntry> deobfClasses )
         {
                 for( ClassEntry obfClassEntry : m_jarIndex.getObfClassEntries() )
                 {
                         // skip inner classes
-                        if( m_jarIndex.getOuterClass( obfClassEntry.getName() ) != null )
+                        if( obfClassEntry.isInnerClass() )
                         {
                                 continue;
                         }
@@ -166,17 +166,17 @@ public class Deobfuscator
                         if( !deobfClassEntry.equals( obfClassEntry ) )
                         {
                                 // if the class has a mapping, clearly it's deobfuscated
-                                deobfClasses.add( deobfClassEntry.getName() );
+                                deobfClasses.add( deobfClassEntry );
                         }
                         else if( !obfClassEntry.getPackageName().equals( Constants.NonePackage ) )
                         {
                                 // also call it deobufscated if it's not in the none package
-                                deobfClasses.add( obfClassEntry.getName() );
+                                deobfClasses.add( obfClassEntry );
                         }
                         else
                         {
                                 // otherwise, assume it's still obfuscated
-                                obfClasses.add( obfClassEntry.getName() );
+                                obfClasses.add( obfClassEntry );
                         }
                 }
         }
diff --git a/src/cuchaz/enigma/gui/ClassSelector.java b/src/cuchaz/enigma/gui/ClassSelector.java
new file mode 100644
index 0000000..338ad80
--- /dev/null
+++ b/src/cuchaz/enigma/gui/ClassSelector.java
@@ -0,0 +1,181 @@
+package cuchaz.enigma.gui;
+
+import java.awt.event.MouseAdapter;
+import java.awt.event.MouseEvent;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.List;
+import java.util.Map;
+
+import javax.swing.JTree;
+import javax.swing.tree.DefaultMutableTreeNode;
+import javax.swing.tree.DefaultTreeModel;
+import javax.swing.tree.TreePath;
+
+import com.beust.jcommander.internal.Lists;
+import com.beust.jcommander.internal.Maps;
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+
+import cuchaz.enigma.mapping.ClassEntry;
+
+public class ClassSelector extends JTree
+{
+        private static final long serialVersionUID = -7632046902384775977L;
+        
+        public interface ClassSelectionListener
+        {
+                void onSelectClass( ClassEntry classEntry );
+        }
+        
+        public static Comparator<ClassEntry> ObfuscatedClassEntryComparator;
+        public static Comparator<ClassEntry> DeobfuscatedClassEntryComparator;
+        
+        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();
+                                }
+                                return a.getName().compareTo( b.getName() );
+                        }
+                };
+                
+                DeobfuscatedClassEntryComparator = new Comparator<ClassEntry>( )
+                {
+                        @Override
+                        public int compare( ClassEntry a, ClassEntry b )
+                        {
+                                return a.getName().compareTo( b.getName() );
+                        }
+                };
+        }
+        
+        private ClassSelectionListener m_listener;
+        private Comparator<ClassEntry> m_comparator;
+        
+        public ClassSelector( Comparator<ClassEntry> comparator )
+        {
+                m_comparator = comparator;
+                
+                // configure the tree control
+                setRootVisible( false );
+                setShowsRootHandles( false );
+                setModel( null );
+                
+                // hook events
+                addMouseListener( new MouseAdapter()
+                {
+                        @Override
+                        public void mouseClicked( MouseEvent event )
+                        {
+                                if( m_listener != null && event.getClickCount() == 2 )
+                                {
+                                        // get the selected node
+                                        TreePath path = getSelectionPath();
+                                        if( path != null && path.getLastPathComponent() instanceof ClassSelectorClassNode )
+                                        {
+                                                ClassSelectorClassNode node = (ClassSelectorClassNode)path.getLastPathComponent();
+                                                m_listener.onSelectClass( node.getClassEntry() );
+                                        }
+                                }
+                        }
+                } );
+                
+                // init defaults
+                m_listener = null;
+        }
+        
+        public void setListener( ClassSelectionListener val )
+        {
+                m_listener = val;
+        }
+        
+        public void setClasses( Collection<ClassEntry> classEntries )
+        {
+                if( classEntries == null )
+                {
+                        setModel( null );
+                        return;
+                }
+                
+                // build the package names
+                Map<String,ClassSelectorPackageNode> packages = Maps.newHashMap();
+                for( ClassEntry classEntry : classEntries )
+                {
+                        packages.put( classEntry.getPackageName(), null );
+                }
+                
+                // sort the packages
+                List<String> sortedPackageNames = Lists.newArrayList( packages.keySet() );
+                Collections.sort( sortedPackageNames, new Comparator<String>( )
+                {
+                        @Override
+                        public int compare( String a, String b )
+                        {
+                                // I can never keep this rule straight when writing these damn things...
+                                // a < b => -1, a == b => 0, a > b => +1
+                                
+                                String[] aparts = a.split( "/" );
+                                String[] bparts = b.split( "/" );
+                                for( int i=0; true; i++ )
+                                {
+                                        if( i >= aparts.length )
+                                        {
+                                                return -1;
+                                        }
+                                        else if( i >= bparts.length )
+                                        {
+                                                return 1;
+                                        }
+                                        
+                                        int result = aparts[i].compareTo( bparts[i] );
+                                        if( result != 0 )
+                                        {
+                                                return result;
+                                        }
+                                }
+                        }
+                } );
+                
+                // create the root node and the package nodes
+                DefaultMutableTreeNode root = new DefaultMutableTreeNode();
+                for( String packageName : sortedPackageNames )
+                {
+                        ClassSelectorPackageNode node = new ClassSelectorPackageNode( packageName );
+                        packages.put( packageName, node );
+                        root.add( node );
+                }
+                
+                // put the classes into packages
+                Multimap<String,ClassEntry> packagedClassEntries = ArrayListMultimap.create();
+                for( ClassEntry classEntry : classEntries )
+                {
+                        packagedClassEntries.put( classEntry.getPackageName(), classEntry );
+                }
+                
+                // build the class nodes
+                for( String packageName : packagedClassEntries.keySet() )
+                {
+                        // sort the class entries
+                        List<ClassEntry> classEntriesInPackage = Lists.newArrayList( packagedClassEntries.get( packageName ) );
+                        Collections.sort( classEntriesInPackage, m_comparator );
+                        
+                        // create the nodes in order
+                        for( ClassEntry classEntry : classEntriesInPackage )
+                        {
+                                ClassSelectorPackageNode node = packages.get( packageName );
+                                node.add( new ClassSelectorClassNode( classEntry ) );
+                        }
+                }
+                
+                // finally, update the tree control
+                setModel( new DefaultTreeModel( root ) );
+        }
+}
diff --git a/src/cuchaz/enigma/gui/ClassSelectorClassNode.java b/src/cuchaz/enigma/gui/ClassSelectorClassNode.java
new file mode 100644
index 0000000..cffa795
--- /dev/null
+++ b/src/cuchaz/enigma/gui/ClassSelectorClassNode.java
@@ -0,0 +1,38 @@
+/*******************************************************************************
+ * Copyright (c) 2014 Jeff Martin.
+ * All rights reserved. This program and the accompanying materials
+ * are made available under the terms of the GNU Public License v3.0
+ * which accompanies this distribution, and is available at
+ * http://www.gnu.org/licenses/gpl.html
+ * 
+ * Contributors:
+ *     Jeff Martin - initial API and implementation
+ ******************************************************************************/
+package cuchaz.enigma.gui;
+
+import javax.swing.tree.DefaultMutableTreeNode;
+
+import cuchaz.enigma.mapping.ClassEntry;
+
+public class ClassSelectorClassNode extends DefaultMutableTreeNode
+{
+        private static final long serialVersionUID = -8956754339813257380L;
+        
+        private ClassEntry m_classEntry;
+        
+        public ClassSelectorClassNode( ClassEntry classEntry )
+        {
+                m_classEntry = classEntry;
+        }
+        
+        public ClassEntry getClassEntry( )
+        {
+                return m_classEntry;
+        }
+        
+        @Override
+        public String toString( )
+        {
+                return m_classEntry.getSimpleName();
+        }
+}
diff --git a/src/cuchaz/enigma/gui/ClassSelectorPackageNode.java b/src/cuchaz/enigma/gui/ClassSelectorPackageNode.java
new file mode 100644
index 0000000..ad88fb4
--- /dev/null
+++ b/src/cuchaz/enigma/gui/ClassSelectorPackageNode.java
@@ -0,0 +1,36 @@
+/*******************************************************************************
+ * Copyright (c) 2014 Jeff Martin.
+ * All rights reserved. This program and the accompanying materials
+ * are made available under the terms of the GNU Public License v3.0
+ * which accompanies this distribution, and is available at
+ * http://www.gnu.org/licenses/gpl.html
+ * 
+ * Contributors:
+ *     Jeff Martin - initial API and implementation
+ ******************************************************************************/
+package cuchaz.enigma.gui;
+
+import javax.swing.tree.DefaultMutableTreeNode;
+
+public class ClassSelectorPackageNode extends DefaultMutableTreeNode
+{
+        private static final long serialVersionUID = -3730868701219548043L;
+        
+        private String m_packageName;
+        
+        public ClassSelectorPackageNode( String packageName )
+        {
+                m_packageName = packageName;
+        }
+        
+        public String getPackageName( )
+        {
+                return m_packageName;
+        }
+        
+        @Override
+        public String toString( )
+        {
+                return m_packageName;
+        }
+}
diff --git a/src/cuchaz/enigma/gui/Gui.java b/src/cuchaz/enigma/gui/Gui.java
index 46395ac..8ae16f4 100644
--- a/src/cuchaz/enigma/gui/Gui.java
+++ b/src/cuchaz/enigma/gui/Gui.java
@@ -31,7 +31,6 @@ import java.io.IOException;
 import java.lang.Thread.UncaughtExceptionHandler;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.Comparator;
 import java.util.List;
 import java.util.Vector;
 
@@ -79,6 +78,7 @@ import cuchaz.enigma.analysis.MethodImplementationsTreeNode;
 import cuchaz.enigma.analysis.MethodInheritanceTreeNode;
 import cuchaz.enigma.analysis.ReferenceTreeNode;
 import cuchaz.enigma.analysis.Token;
+import cuchaz.enigma.gui.ClassSelector.ClassSelectionListener;
 import cuchaz.enigma.mapping.ArgumentEntry;
 import cuchaz.enigma.mapping.ClassEntry;
 import cuchaz.enigma.mapping.ConstructorEntry;
@@ -90,61 +90,12 @@ import cuchaz.enigma.mapping.MethodEntry;
 
 public class Gui
 {
-        private static Comparator<String> m_obfClassSorter;
-        private static Comparator<String> m_deobfClassSorter;
-        
-        static
-        {
-                m_obfClassSorter = new Comparator<String>( )
-                {
-                        @Override
-                        public int compare( String a, String b )
-                        {
-                                if( a.length() != b.length() )
-                                {
-                                        return a.length() - b.length();
-                                }
-                                return a.compareTo( b );
-                        }
-                };
-                
-                m_deobfClassSorter = new Comparator<String>( )
-                {
-                        @Override
-                        public int compare( String a, String b )
-                        {
-                                // I can never keep this rule straight when writing these damn things...
-                                // a < b => -1, a == b => 0, a > b => +1
-                                
-                                String[] aparts = a.split( "\\." );
-                                String[] bparts = b.split( "\\." );
-                                for( int i=0; true; i++ )
-                                {
-                                        if( i >= aparts.length )
-                                        {
-                                                return -1;
-                                        }
-                                        else if( i >= bparts.length )
-                                        {
-                                                return 1;
-                                        }
-                                        
-                                        int result = aparts[i].compareTo( bparts[i] );
-                                        if( result != 0 )
-                                        {
-                                                return result;
-                                        }
-                                }
-                        }
-                };
-        }
-        
         private GuiController m_controller;
         
         // controls
         private JFrame m_frame;
-        private JList<String> m_obfClasses;
-        private JList<String> m_deobfClasses;
+        private ClassSelector m_obfClasses;
+        private ClassSelector m_deobfClasses;
         private JEditorPane m_editor;
         private JPanel m_classesPanel;
         private JSplitPane m_splitClasses;
@@ -205,27 +156,17 @@ public class Gui
                 m_exportFileChooser.setFileSelectionMode( JFileChooser.DIRECTORIES_ONLY );
                 
                 // init obfuscated classes list
-                m_obfClasses = new JList<String>();
-                m_obfClasses.setSelectionMode( ListSelectionModel.SINGLE_SELECTION );
-                m_obfClasses.setLayoutOrientation( JList.VERTICAL );
-                m_obfClasses.setCellRenderer( new ClassListCellRenderer() );
-                m_obfClasses.addMouseListener( new MouseAdapter()
+                m_obfClasses = new ClassSelector( ClassSelector.ObfuscatedClassEntryComparator );
+                m_obfClasses.setListener( new ClassSelectionListener( )
                 {
                         @Override
-                        public void mouseClicked( MouseEvent event )
+                        public void onSelectClass( ClassEntry classEntry )
                         {
-                                if( event.getClickCount() == 2 )
+                                if( m_reference != null )
                                 {
-                                        String selected = m_obfClasses.getSelectedValue();
-                                        if( selected != null )
-                                        {
-                                                if( m_reference != null )
-                                                {
-                                                        m_controller.savePreviousReference( m_reference );
-                                                }
-                                                m_controller.openDeclaration( new ClassEntry( selected ) );
-                                        }
+                                        m_controller.savePreviousReference( m_reference );
                                 }
+                                m_controller.openDeclaration( classEntry );
                         }
                 } );
                 JScrollPane obfScroller = new JScrollPane( m_obfClasses );
@@ -235,27 +176,17 @@ public class Gui
                 obfPanel.add( obfScroller, BorderLayout.CENTER );
                 
                 // init deobfuscated classes list
-                m_deobfClasses = new JList<String>();
-                m_deobfClasses.setSelectionMode( ListSelectionModel.SINGLE_SELECTION );
-                m_deobfClasses.setLayoutOrientation( JList.VERTICAL );
-                m_deobfClasses.setCellRenderer( new ClassListCellRenderer() );
-                m_deobfClasses.addMouseListener( new MouseAdapter()
+                m_deobfClasses = new ClassSelector( ClassSelector.DeobfuscatedClassEntryComparator );
+                m_deobfClasses.setListener( new ClassSelectionListener( )
                 {
                         @Override
-                        public void mouseClicked( MouseEvent event )
+                        public void onSelectClass( ClassEntry classEntry )
                         {
-                                if( event.getClickCount() == 2 )
+                                if( m_reference != null )
                                 {
-                                        String selected = m_deobfClasses.getSelectedValue();
-                                        if( selected != null )
-                                        {
-                                                if( m_reference != null )
-                                                {
-                                                        m_controller.savePreviousReference( m_reference );
-                                                }
-                                                m_controller.openDeclaration( new ClassEntry( selected ) );
-                                        }
+                                        m_controller.savePreviousReference( m_reference );
                                 }
+                                m_controller.openDeclaration( classEntry );
                         }
                 } );
                 JScrollPane deobfScroller = new JScrollPane( m_deobfClasses );
@@ -266,6 +197,7 @@ public class Gui
                 
                 // set up classes panel (don't add the splitter yet)
                 m_splitClasses = new JSplitPane( JSplitPane.VERTICAL_SPLIT, true, obfPanel, deobfPanel );
+                m_splitClasses.setResizeWeight( 0.3 );
                 m_classesPanel = new JPanel();
                 m_classesPanel.setLayout( new BorderLayout() );
                 m_classesPanel.setPreferredSize( new Dimension( 250, 0 ) );
@@ -847,32 +779,14 @@ public class Gui
                 redraw();
         }
         
-        public void setObfClasses( List<String> obfClasses )
+        public void setObfClasses( Collection<ClassEntry> obfClasses )
         {
-                if( obfClasses != null )
-                {
-                        Vector<String> sortedClasses = new Vector<String>( obfClasses );
-                        Collections.sort( sortedClasses, m_obfClassSorter );
-                        m_obfClasses.setListData( sortedClasses );
-                }
-                else
-                {
-                        m_obfClasses.setListData( new Vector<String>() );
-                }
+                m_obfClasses.setClasses( obfClasses );
         }
         
-        public void setDeobfClasses( List<String> deobfClasses )
+        public void setDeobfClasses( Collection<ClassEntry> deobfClasses )
         {
-                if( deobfClasses != null )
-                {
-                        Vector<String> sortedClasses = new Vector<String>( deobfClasses );
-                        Collections.sort( sortedClasses, m_deobfClassSorter );
-                        m_deobfClasses.setListData( sortedClasses );
-                }
-                else
-                {
-                        m_deobfClasses.setListData( new Vector<String>() );
-                }
+                m_deobfClasses.setClasses( deobfClasses );
         }
         
         public void setMappingsFile( File file )
diff --git a/src/cuchaz/enigma/gui/GuiController.java b/src/cuchaz/enigma/gui/GuiController.java
index 2879483..c0fb2e4 100644
--- a/src/cuchaz/enigma/gui/GuiController.java
+++ b/src/cuchaz/enigma/gui/GuiController.java
@@ -315,8 +315,8 @@ public class GuiController
         
         private void refreshClasses( )
         {
-                List<String> obfClasses = Lists.newArrayList();
-                List<String> deobfClasses = Lists.newArrayList();
+                List<ClassEntry> obfClasses = Lists.newArrayList();
+                List<ClassEntry> deobfClasses = Lists.newArrayList();
                 m_deobfuscator.getSeparatedClasses( obfClasses, deobfClasses );
                 m_gui.setObfClasses( obfClasses );
                 m_gui.setDeobfClasses( deobfClasses );