fixed bugs with saving mappings

got argument renaming to work

8 files changed, 112 insertions, 77 deletions

diff --git a/src/cuchaz/enigma/TranslatingTypeLoader.java b/src/cuchaz/enigma/TranslatingTypeLoader.java
index cd36e8d..e57a09d 100644
--- a/src/cuchaz/enigma/TranslatingTypeLoader.java
+++ b/src/cuchaz/enigma/TranslatingTypeLoader.java
@@ -70,8 +70,8 @@ public class TranslatingTypeLoader implements ITypeLoader
                         try
                         {
                                 CtClass c = classPool.get( name );
-                                new ClassTranslator( m_deobfuscatingTranslator ).translate( c );
                                 new MethodParameterWriter( m_deobfuscatingTranslator ).writeMethodArguments( c );
+                                new ClassTranslator( m_deobfuscatingTranslator ).translate( c );
                                 buf = c.toBytecode();
                         }
                         catch( Exception ex )
diff --git a/src/cuchaz/enigma/analysis/SourcedAst.java b/src/cuchaz/enigma/analysis/SourcedAst.java
index 04c6f03..52a3453 100644
--- a/src/cuchaz/enigma/analysis/SourcedAst.java
+++ b/src/cuchaz/enigma/analysis/SourcedAst.java
@@ -13,7 +13,10 @@ package cuchaz.enigma.analysis;
 import java.io.IOException;
 import java.util.HashMap;
 
+import javassist.bytecode.Descriptor;
+
 import com.google.common.collect.Maps;
+import com.sun.source.tree.ClassTree;
 import com.sun.source.tree.CompilationUnitTree;
 import com.sun.source.tree.ImportTree;
 import com.sun.source.tree.Tree;
@@ -34,7 +37,7 @@ public class SourcedAst
                 m_positions = m_trees.getSourcePositions();
                 m_classNameIndex = Maps.newHashMap();
                 
-                // index all the class names
+                // index all the class names from package imports
                 for( ImportTree importTree : m_tree.getImports() )
                 {
                         // ignore static imports for now
@@ -44,9 +47,9 @@ public class SourcedAst
                         }
                         
                         // get the full and simple class names
-                        String fullName = importTree.getQualifiedIdentifier().toString();
+                        String fullName = Descriptor.toJvmName( importTree.getQualifiedIdentifier().toString() );
                         String simpleName = fullName;
-                        String[] parts = fullName.split( "\\." );
+                        String[] parts = fullName.split( "/" );
                         if( parts.length > 0 )
                         {
                                 simpleName = parts[parts.length - 1];
@@ -54,6 +57,18 @@ public class SourcedAst
                         
                         m_classNameIndex.put( simpleName, fullName );
                 }
+                
+                // index the self class using the package name
+                String packageName = Descriptor.toJvmName( m_tree.getPackageName().toString() );
+                for( Tree typeTree : m_tree.getTypeDecls() )
+                {
+                        if( typeTree instanceof ClassTree )
+                        {
+                                ClassTree classTree = (ClassTree)typeTree;
+                                String className = classTree.getSimpleName().toString();
+                                m_classNameIndex.put( className, packageName + "/" + className );
+                        }
+                }
         }
 
         public int getStart( Tree node )
diff --git a/src/cuchaz/enigma/bytecode/MethodParameterWriter.java b/src/cuchaz/enigma/bytecode/MethodParameterWriter.java
index 1e5d1f0..a8d3983 100644
--- a/src/cuchaz/enigma/bytecode/MethodParameterWriter.java
+++ b/src/cuchaz/enigma/bytecode/MethodParameterWriter.java
@@ -10,9 +10,15 @@
  ******************************************************************************/
 package cuchaz.enigma.bytecode;
 
+import java.util.ArrayList;
+import java.util.List;
+
 import javassist.CtBehavior;
 import javassist.CtClass;
-import javassist.bytecode.AttributeInfo;
+import javassist.bytecode.Descriptor;
+import cuchaz.enigma.mapping.ArgumentEntry;
+import cuchaz.enigma.mapping.ClassEntry;
+import cuchaz.enigma.mapping.MethodEntry;
 import cuchaz.enigma.mapping.Translator;
 
 public class MethodParameterWriter
@@ -27,9 +33,25 @@ public class MethodParameterWriter
         public void writeMethodArguments( CtClass c )
         {
                 // Procyon will read method arguments from the "MethodParameters" attribute, so write those
+                ClassEntry classEntry = new ClassEntry( Descriptor.toJvmName( c.getName() ) );
                 for( CtBehavior behavior : c.getDeclaredBehaviors() )
                 {
-                        AttributeInfo attribute = behavior.getMethodInfo().getAttribute( "MethodParameter" );
+                        int numParams = Descriptor.numOfParameters( behavior.getMethodInfo().getDescriptor() );
+                        if( numParams <= 0 )
+                        {
+                                continue;
+                        }
+                        
+                        // get the list of parameter names
+                        MethodEntry methodEntry = new MethodEntry( classEntry, behavior.getMethodInfo().getName(), behavior.getSignature() );
+                        List<String> names = new ArrayList<String>( numParams );
+                        for( int i=0; i<numParams; i++ )
+                        {
+                                names.add( m_translator.translate( new ArgumentEntry( methodEntry, i, "" ) ) );
+                        }
+                        
+                        // save the mappings to the class
+                        MethodParametersAttribute.updateClass( behavior.getMethodInfo(), names );
                 }
         }
 }
diff --git a/src/cuchaz/enigma/bytecode/MethodParametersAttribute.java b/src/cuchaz/enigma/bytecode/MethodParametersAttribute.java
index 0b29403..baf1ac1 100644
--- a/src/cuchaz/enigma/bytecode/MethodParametersAttribute.java
+++ b/src/cuchaz/enigma/bytecode/MethodParametersAttribute.java
@@ -13,24 +13,45 @@ package cuchaz.enigma.bytecode;
 import java.io.ByteArrayOutputStream;
 import java.io.DataOutputStream;
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.List;
 
 import javassist.bytecode.AttributeInfo;
 import javassist.bytecode.ConstPool;
+import javassist.bytecode.MethodInfo;
 
 public class MethodParametersAttribute extends AttributeInfo
 {
-        public MethodParametersAttribute( ConstPool pool, int attributeNameIndex, List<Integer> parameterNameIndices )
+        private MethodParametersAttribute( ConstPool pool, List<Integer> parameterNameIndices )
         {
-                super( pool, "MethodParameters", writeStruct( attributeNameIndex, parameterNameIndices ) );
+                super( pool, "MethodParameters", writeStruct( parameterNameIndices ) );
         }
         
-        private static byte[] writeStruct( int attributeNameIndex, List<Integer> parameterNameIndices )
+        public static void updateClass( MethodInfo info, List<String> names )
         {
-                // JVM Spec says the struct looks like this:
+                // add the names to the class const pool
+                ConstPool constPool = info.getConstPool();
+                List<Integer> parameterNameIndices = new ArrayList<Integer>();
+                for( String name : names )
+                {
+                        if( name != null )
+                        {
+                                parameterNameIndices.add( constPool.addUtf8Info( name ) );
+                        }
+                        else
+                        {
+                                parameterNameIndices.add( 0 );
+                        }
+                }
+                
+                // add the attribute to the method
+                info.addAttribute( new MethodParametersAttribute( constPool, parameterNameIndices ) );
+        }
+        
+        private static byte[] writeStruct( List<Integer> parameterNameIndices )
+        {
+                // JVM 8 Spec says the struct looks like this:
                 // http://cr.openjdk.java.net/~mr/se/8/java-se-8-fr-spec-01/java-se-8-jvms-fr-diffs.pdf
-                // uint16 name_index -> points to UTF8 entry in constant pool that says "MethodParameters"
-                // uint32 length -> length of this struct, minus 6 bytes (ie, length of num_params and parameter array)
                 // uint8 num_params
                 // for each param:
                 //    uint16 name_index -> points to UTF8 entry in constant pool, or 0 for no entry
@@ -43,20 +64,13 @@ public class MethodParametersAttribute extends AttributeInfo
                 // the writeShort(), writeByte() methods will read 16,8 low-order bits from the int argument
                 // as long as the int argument is in range of the unsigned short/byte type, it will be written as an unsigned short/byte
                 // if the int is out of range, the byte stream won't look the way we want and weird things will happen
+                final int SIZEOF_UINT8 = 1;
                 final int SIZEOF_UINT16 = 2;
                 final int MAX_UINT8 = ( 1 << 8 ) - 1;
                 final int MAX_UINT16 = ( 1 << 16 ) - 1;
-                final long MAX_UINT32 = ( 1 << 32 ) - 1;
                 
                 try
                 {
-                        assert( attributeNameIndex >= 0 && attributeNameIndex <= MAX_UINT16 );
-                        out.writeShort( attributeNameIndex );
-                        
-                        long length = SIZEOF_UINT16 + parameterNameIndices.size()*( SIZEOF_UINT16 + SIZEOF_UINT16 );
-                        assert( length >= 0 && length <= MAX_UINT32 );
-                        out.writeInt( (int)length );
-                        
                         assert( parameterNameIndices.size() >= 0 && parameterNameIndices.size() <= MAX_UINT8 );
                         out.writeByte( parameterNameIndices.size() );
                         
@@ -70,7 +84,9 @@ public class MethodParametersAttribute extends AttributeInfo
                         }
                         
                         out.close();
-                        return buf.toByteArray();
+                        byte[] data = buf.toByteArray();
+                        assert( data.length == SIZEOF_UINT8 + parameterNameIndices.size()*( SIZEOF_UINT16 + SIZEOF_UINT16 ) );
+                        return data;
                 }
                 catch( IOException ex )
                 {
diff --git a/src/cuchaz/enigma/gui/Gui.java b/src/cuchaz/enigma/gui/Gui.java
index d448dc2..3f46b6e 100644
--- a/src/cuchaz/enigma/gui/Gui.java
+++ b/src/cuchaz/enigma/gui/Gui.java
@@ -20,6 +20,7 @@ import java.awt.event.ActionEvent;
 import java.awt.event.ActionListener;
 import java.awt.event.MouseAdapter;
 import java.awt.event.MouseEvent;
+import java.io.File;
 import java.io.IOException;
 import java.util.Collections;
 import java.util.Comparator;
@@ -54,6 +55,7 @@ import cuchaz.enigma.Constants;
 import cuchaz.enigma.analysis.SourceIndex;
 import cuchaz.enigma.mapping.ArgumentEntry;
 import cuchaz.enigma.mapping.ClassEntry;
+import cuchaz.enigma.mapping.Entry;
 import cuchaz.enigma.mapping.EntryPair;
 import cuchaz.enigma.mapping.FieldEntry;
 import cuchaz.enigma.mapping.MethodEntry;
@@ -101,9 +103,9 @@ public class Gui
         private JMenuItem m_closeMappingsMenu;
         
         // state
-        private EntryPair m_selectedEntryPair;
+        private EntryPair<Entry> m_selectedEntryPair;
         private JFileChooser m_jarFileChooser;
-        private JFileChooser m_mappingFileChooser;
+        private JFileChooser m_mappingsFileChooser;
         
         public Gui( )
         {
@@ -111,7 +113,7 @@ public class Gui
                 
                 // init file choosers
                 m_jarFileChooser = new JFileChooser();
-                m_mappingFileChooser = new JFileChooser();
+                m_mappingsFileChooser = new JFileChooser();
                 
                 // init frame
                 m_frame = new JFrame( Constants.Name );
@@ -279,11 +281,11 @@ public class Gui
                                         @Override
                                         public void actionPerformed( ActionEvent event )
                                         {
-                                                if( m_mappingFileChooser.showOpenDialog( m_frame ) == JFileChooser.APPROVE_OPTION )
+                                                if( m_mappingsFileChooser.showOpenDialog( m_frame ) == JFileChooser.APPROVE_OPTION )
                                                 {
                                                         try
                                                         {
-                                                                m_controller.openMappings( m_mappingFileChooser.getSelectedFile() );
+                                                                m_controller.openMappings( m_mappingsFileChooser.getSelectedFile() );
                                                         }
                                                         catch( IOException ex )
                                                         {
@@ -304,7 +306,7 @@ public class Gui
                                         {
                                                 try
                                                 {
-                                                        m_controller.saveMappings( m_mappingFileChooser.getSelectedFile() );
+                                                        m_controller.saveMappings( m_mappingsFileChooser.getSelectedFile() );
                                                 }
                                                 catch( IOException ex )
                                                 {
@@ -322,11 +324,11 @@ public class Gui
                                         @Override
                                         public void actionPerformed( ActionEvent event )
                                         {
-                                                if( m_mappingFileChooser.showSaveDialog( m_frame ) == JFileChooser.APPROVE_OPTION )
+                                                if( m_mappingsFileChooser.showSaveDialog( m_frame ) == JFileChooser.APPROVE_OPTION )
                                                 {
                                                         try
                                                         {
-                                                                m_controller.saveMappings( m_mappingFileChooser.getSelectedFile() );
+                                                                m_controller.saveMappings( m_mappingsFileChooser.getSelectedFile() );
                                                                 m_saveMappingsMenu.setEnabled( true );
                                                         }
                                                         catch( IOException ex )
@@ -440,9 +442,10 @@ public class Gui
                 }
         }
         
-        public void setMappingsLoaded( boolean isLoaded )
+        public void setMappingsFile( File file )
         {
-                m_saveMappingsMenu.setEnabled( isLoaded );
+                m_mappingsFileChooser.setSelectedFile( file );
+                m_saveMappingsMenu.setEnabled( file != null );
         }
         
         public void setSource( String source )
@@ -493,7 +496,8 @@ public class Gui
                 redraw();
         }
         
-        private void showEntryPair( EntryPair pair )
+        @SuppressWarnings( "unchecked" )
+        private void showEntryPair( EntryPair<Entry> pair )
         {
                 if( pair == null )
                 {
@@ -514,19 +518,19 @@ public class Gui
                 m_actionPanel.add( dynamicPanel );
                 if( pair.deobf instanceof ClassEntry )
                 {
-                        showEntry( (ClassEntry)pair.deobf, dynamicPanel );
+                        showClassEntryPair( (EntryPair<? extends ClassEntry>)pair, dynamicPanel );
                 }
                 else if( pair.deobf instanceof FieldEntry )
                 {
-                        showEntry( (FieldEntry)pair.deobf, dynamicPanel );
+                        showFieldEntryPair( (EntryPair<? extends FieldEntry>)pair, dynamicPanel );
                 }
                 else if( pair.deobf instanceof MethodEntry )
                 {
-                        showEntry( (MethodEntry)pair.deobf, dynamicPanel );
+                        showMethodEntryPair( (EntryPair<? extends MethodEntry>)pair, dynamicPanel );
                 }
                 else if( pair.deobf instanceof ArgumentEntry )
                 {
-                        showEntry( (ArgumentEntry)pair.deobf, dynamicPanel );
+                        showArgumentEntryPair( (EntryPair<? extends ArgumentEntry>)pair, dynamicPanel );
                 }
                 else
                 {
@@ -536,30 +540,30 @@ public class Gui
                 redraw();
         }
         
-        private void showEntry( ClassEntry entry, JPanel panel )
+        private void showClassEntryPair( EntryPair<? extends ClassEntry> pair, JPanel panel )
         {
                 m_typeLabel.setText( "Class: " );
         }
         
-        private void showEntry( FieldEntry entry, JPanel panel )
+        private void showFieldEntryPair( EntryPair<? extends FieldEntry> pair, JPanel panel )
         {
                 m_typeLabel.setText( "Field: " );
-                addNameValue( panel, "Class", entry.getClassEntry().getName() );
+                addNameValue( panel, "Class", pair.obf.getClassEntry().getName() + " <-> " + pair.deobf.getClassEntry().getName() );
         }
         
-        private void showEntry( MethodEntry entry, JPanel panel )
+        private void showMethodEntryPair( EntryPair<? extends MethodEntry> pair, JPanel panel )
         {
                 m_typeLabel.setText( "Method: " );
-                addNameValue( panel, "Class", entry.getClassEntry().getName() );
-                addNameValue( panel, "Signature", entry.getSignature() );
+                addNameValue( panel, "Class", pair.obf.getClassEntry().getName() + " <-> " + pair.deobf.getClassEntry().getName() );
+                addNameValue( panel, "Signature", pair.obf.getSignature() + " <-> " + pair.deobf.getSignature() );
         }
         
-        private void showEntry( ArgumentEntry entry, JPanel panel )
+        private void showArgumentEntryPair( EntryPair<? extends ArgumentEntry> pair, JPanel panel )
         {
                 m_typeLabel.setText( "Argument: " );
-                addNameValue( panel, "Class", entry.getMethodEntry().getClassEntry().getName() );
-                addNameValue( panel, "Method", entry.getMethodEntry().getName() );
-                addNameValue( panel, "Index", Integer.toString( entry.getIndex() ) );
+                addNameValue( panel, "Class", pair.obf.getClassEntry().getName() + " <-> " + pair.deobf.getClassEntry().getName() );
+                addNameValue( panel, "Method", pair.obf.getMethodEntry().getName() + " <-> " + pair.deobf.getMethodEntry().getName() );
+                addNameValue( panel, "Index", Integer.toString( pair.obf.getIndex() ) );
         }
         
         private void addNameValue( JPanel container, String name, String value )
diff --git a/src/cuchaz/enigma/gui/GuiController.java b/src/cuchaz/enigma/gui/GuiController.java
index ce1c31b..6704ef8 100644
--- a/src/cuchaz/enigma/gui/GuiController.java
+++ b/src/cuchaz/enigma/gui/GuiController.java
@@ -62,7 +62,7 @@ public class GuiController
                 FileReader in = new FileReader( file );
                 m_deobfuscator.setMappings( new MappingsReader().read( in ) );
                 in.close();
-                m_gui.setMappingsLoaded( true );
+                m_gui.setMappingsFile( file );
                 refreshClasses();
                 refreshOpenFiles();
         }
@@ -78,7 +78,7 @@ public class GuiController
         public void closeMappings( )
         {
                 m_deobfuscator.setMappings( null );
-                m_gui.setMappingsLoaded( false );
+                m_gui.setMappingsFile( null );
                 refreshOpenFiles();
         }
         
@@ -88,7 +88,7 @@ public class GuiController
                 deobfuscate( m_currentFile );
         }
         
-        public EntryPair getEntryPair( int pos )
+        public EntryPair<Entry> getEntryPair( int pos )
         {
                 if( m_index == null )
                 {
@@ -100,7 +100,7 @@ public class GuiController
                 {
                         return null;
                 }
-                return new EntryPair( m_deobfuscator.obfuscate( deobfEntry ), deobfEntry );
+                return new EntryPair<Entry>( m_deobfuscator.obfuscate( deobfEntry ), deobfEntry );
         }
         
         public void rename( Entry obfsEntry, String newName )
diff --git a/src/cuchaz/enigma/mapping/EntryPair.java b/src/cuchaz/enigma/mapping/EntryPair.java
index e40e999..e3325b3 100644
--- a/src/cuchaz/enigma/mapping/EntryPair.java
+++ b/src/cuchaz/enigma/mapping/EntryPair.java
@@ -10,37 +10,15 @@
  ******************************************************************************/
 package cuchaz.enigma.mapping;
 
-import cuchaz.enigma.Util;
 
-public class EntryPair
+public class EntryPair<T extends Entry>
 {
-        public Entry obf;
-        public Entry deobf;
+        public T obf;
+        public T deobf;
         
-        public EntryPair( Entry obf, Entry deobf )
+        public EntryPair( T obf, T deobf )
         {
                 this.obf = obf;
                 this.deobf = deobf;
         }
-        
-        @Override
-        public int hashCode( )
-        {
-                return Util.combineHashesOrdered( obf, deobf );
-        }
-        
-        @Override
-        public boolean equals( Object other )
-        {
-                if( other instanceof EntryPair )
-                {
-                        return equals( (EntryPair)other );
-                }
-                return false;
-        }
-        
-        public boolean equals( EntryPair other )
-        {
-                return obf.equals( other.obf ) && deobf.equals( other.deobf );
-        }
 }
diff --git a/src/cuchaz/enigma/mapping/MethodMapping.java b/src/cuchaz/enigma/mapping/MethodMapping.java
index f2bc54d..1cdc38a 100644
--- a/src/cuchaz/enigma/mapping/MethodMapping.java
+++ b/src/cuchaz/enigma/mapping/MethodMapping.java
@@ -127,7 +127,7 @@ public class MethodMapping implements Serializable
                 {
                         buf.append( "\t\t" );
                         buf.append( argumentMapping.getIndex() );
-                        buf.append( " <-> " );
+                        buf.append( " -> " );
                         buf.append( argumentMapping.getName() );
                         buf.append( "\n" );
                 }