finished class matching for now, need to work on class member matching

6 files changed, 319 insertions, 57 deletions

diff --git a/src/cuchaz/enigma/convert/ClassIdentity.java b/src/cuchaz/enigma/convert/ClassIdentity.java
index 0a3a4497..980f31f5 100644
--- a/src/cuchaz/enigma/convert/ClassIdentity.java
+++ b/src/cuchaz/enigma/convert/ClassIdentity.java
@@ -60,6 +60,7 @@ import cuchaz.enigma.mapping.SignatureUpdater.ClassNameUpdater;
 public class ClassIdentity
 {
         private ClassEntry m_classEntry;
+        private String m_rawName;
         private SidedClassNamer m_namer;
         private Set<String> m_fields;
         private Set<String> m_methods;
@@ -70,13 +71,18 @@ public class ClassIdentity
         private Set<String> m_implementations;
         private Set<String> m_references;
         
-        public ClassIdentity( CtClass c, SidedClassNamer namer, JarIndex index, boolean useReferences )
+        public ClassIdentity( CtClass c, SidedClassNamer namer, JarIndex index, boolean useReferences, boolean useRawNames )
         {
                 m_namer = namer;
                 
                 // stuff from the bytecode
                 
                 m_classEntry = new ClassEntry( Descriptor.toJvmName( c.getName() ) );
+                m_rawName = "";
+                if( useRawNames )
+                {
+                        m_rawName = m_classEntry.getName();
+                }
                 m_fields = Sets.newHashSet();
                 for( CtField field : c.getDeclaredFields() )
                 {
@@ -176,8 +182,16 @@ public class ClassIdentity
         {
                 StringBuilder buf = new StringBuilder();
                 buf.append( "class: " );
+                buf.append( m_classEntry.getName() );
+                buf.append( " " );
                 buf.append( hashCode() );
                 buf.append( "\n" );
+                if( m_rawName.length() > 0 )
+                {
+                        buf.append( "\traw name: " );
+                        buf.append( m_rawName );
+                        buf.append( "\n" );
+                }
                 for( String field : m_fields )
                 {
                         buf.append( "\tfield " );
@@ -425,7 +439,8 @@ public class ClassIdentity
         
         public boolean equals( ClassIdentity other )
         {
-                return m_fields.equals( other.m_fields )
+                return m_rawName.equals( other.m_rawName )
+                        && m_fields.equals( other.m_fields )
                         && m_methods.equals( other.m_methods )
                         && m_constructors.equals( other.m_constructors )
                         && m_staticInitializer.equals( other.m_staticInitializer )
@@ -439,6 +454,7 @@ public class ClassIdentity
         public int hashCode( )
         {
                 List<Object> objs = Lists.newArrayList();
+                objs.add( m_rawName );
                 objs.addAll( m_fields );
                 objs.addAll( m_methods );
                 objs.addAll( m_constructors );
@@ -449,4 +465,24 @@ public class ClassIdentity
                 objs.addAll( m_references );
                 return Util.combineHashesOrdered( objs );
         }
+        
+        public int getMatchScore( ClassIdentity other )
+        {
+                return getNumMatches( m_fields, other.m_fields )
+                        + getNumMatches( m_methods, other.m_methods )
+                        + getNumMatches( m_constructors, other.m_constructors );
+        }
+        
+        private int getNumMatches( Set<String> a, Set<String> b ) 
+        {
+                int numMatches = 0;
+                for( String val : a )
+                {
+                        if( b.contains( val ) )
+                        {
+                                numMatches++;
+                        }
+                }
+                return numMatches;
+        }
 }
diff --git a/src/cuchaz/enigma/convert/ClassMapper.java b/src/cuchaz/enigma/convert/ClassMapper.java
index fe48c505..fd6ab922 100644
--- a/src/cuchaz/enigma/convert/ClassMapper.java
+++ b/src/cuchaz/enigma/convert/ClassMapper.java
@@ -11,32 +11,130 @@
 package cuchaz.enigma.convert;
 
 import java.io.File;
+import java.io.FileReader;
+import java.io.FileWriter;
 import java.io.IOException;
 import java.util.Arrays;
+import java.util.Collections;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.TreeMap;
 import java.util.jar.JarFile;
 
 import javassist.CtClass;
+
+import com.beust.jcommander.internal.Sets;
+import com.google.common.collect.Maps;
+
+import cuchaz.enigma.Constants;
 import cuchaz.enigma.TranslatingTypeLoader;
 import cuchaz.enigma.analysis.JarIndex;
 import cuchaz.enigma.convert.ClassNamer.SidedClassNamer;
 import cuchaz.enigma.mapping.ClassEntry;
+import cuchaz.enigma.mapping.ClassMapping;
+import cuchaz.enigma.mapping.MappingParseException;
+import cuchaz.enigma.mapping.Mappings;
+import cuchaz.enigma.mapping.MappingsReader;
+import cuchaz.enigma.mapping.MappingsWriter;
 
 public class ClassMapper
 {
         public static void main( String[] args )
-        throws IOException
+        throws IOException, MappingParseException
         {
                 // TEMP
                 JarFile fromJar = new JarFile( new File( "input/1.8-pre1.jar" ) );
-                JarFile toJar = new JarFile( new File( "input/1.8-pre2.jar" ) );
+                JarFile toJar = new JarFile( new File( "input/1.8-pre3.jar" ) );
+                File inMappingsFile = new File( "../minecraft-mappings/1.8-pre.mappings" );
+                File outMappingsFile = new File( "../minecraft-mappings/1.8-pre3.mappings" );
                 
                 // compute the matching
                 ClassMatching matching = ClassMapper.computeMatching( fromJar, toJar );
                 
-                // TODO: use the matching to convert the mappings
+                // use the matching to convert the mappings
+                Mappings mappings = new MappingsReader().read( new FileReader( inMappingsFile ) );
+                Map<String,Map.Entry<ClassIdentity,List<ClassIdentity>>> conversionMap = matching.getConversionMap();
+                Map<String,String> finalConversion = Maps.newHashMap();
+                Set<String> unmatchedSourceClasses = Sets.newHashSet();
+                for( ClassMapping classMapping : mappings.classes() )
+                {
+                        // is there a match for this class?
+                        Map.Entry<ClassIdentity,List<ClassIdentity>> entry = conversionMap.get( classMapping.getObfName() );
+                        ClassIdentity sourceClass = entry.getKey();
+                        List<ClassIdentity> matches = entry.getValue();
+                        
+                        if( matches.isEmpty() )
+                        {
+                                // no match! =(
+                                System.out.println( "No exact match for source class " + classMapping.getObfName() );
+                                
+                                // find the closest classes
+                                TreeMap<Integer,ClassIdentity> scoredMatches = Maps.newTreeMap( Collections.reverseOrder() );
+                                for( ClassIdentity c : matching.getUnmatchedDestClasses() )
+                                {
+                                        scoredMatches.put( sourceClass.getMatchScore( c ), c );
+                                }
+                                Iterator<Map.Entry<Integer,ClassIdentity>> iter = scoredMatches.entrySet().iterator();
+                                for( int i=0; i<10 && iter.hasNext(); i++ )
+                                {
+                                        Map.Entry<Integer,ClassIdentity> score = iter.next();
+                                        System.out.println( String.format( "\tScore: %3d   %s", score.getKey(), score.getValue().getClassEntry().getName() ) );
+                                }
+                                
+                                // does the best match have a non-zero score and the same name?
+                                Map.Entry<Integer,ClassIdentity> bestMatch = scoredMatches.firstEntry();
+                                if( bestMatch.getKey() > 0 && bestMatch.getValue().getClassEntry().equals( sourceClass.getClassEntry() ) )
+                                {
+                                        // use it
+                                        System.out.println( "\tAutomatically choosing likely match: " + bestMatch.getValue().getClassEntry().getName() );
+                                        addFinalConversion( finalConversion, sourceClass, bestMatch.getValue() );
+                                }
+                                else
+                                {
+                                        unmatchedSourceClasses.add( classMapping.getObfName() );
+                                }
+                        }
+                        if( matches.size() == 1 )
+                        {
+                                // unique match! We're good to go!
+                                addFinalConversion( finalConversion, sourceClass, matches.get( 0 ) );
+                        }
+                        else if( matches.size() > 1 )
+                        {
+                                // too many matches! =(
+                                unmatchedSourceClasses.add( classMapping.getObfName() );
+                        }
+                }
+                
+                // remove (and warn about) unmatched classes
+                if( !unmatchedSourceClasses.isEmpty() )
+                {
+                        System.err.println( "WARNING: there were unmatched classes!" );
+                        for( String className : unmatchedSourceClasses )
+                        {
+                                System.err.println( "\t" + className );
+                                mappings.removeClassByObfName( className );
+                        }
+                        System.err.println( "Mappings for these classes have been removed." );
+                }
+                
+                // show the class name changes
+                for( Map.Entry<String,String> entry : finalConversion.entrySet() )
+                {
+                        if( !entry.getKey().equals( entry.getValue() ) )
+                        {
+                                System.out.println( String.format( "Class change: %s -> %s", entry.getKey(), entry.getValue() ) );
+                        }
+                }
+                
+                // do the final conversion
+                mappings.renameObfClasses( finalConversion );
+                FileWriter writer = new FileWriter( outMappingsFile );
+                new MappingsWriter().write( writer, mappings );
+                writer.close();
+                System.out.println( "Wrote converted mappings to:\n\t" + outMappingsFile.getAbsolutePath() );
         }
         
         public static ClassMatching computeMatching( JarFile sourceJar, JarFile destJar )
@@ -55,78 +153,81 @@ public class ClassMapper
                 TranslatingTypeLoader destLoader = new TranslatingTypeLoader( destJar, destIndex );
 
                 ClassMatching matching = null;
-                for( boolean useReferences : Arrays.asList( false, true ) )
+                for( boolean useRawNames : Arrays.asList( false, true ) )
                 {
-                        int numMatches = 0;
-                        do
+                        for( boolean useReferences : Arrays.asList( false, true ) )
                         {
-                                SidedClassNamer sourceNamer = null;
-                                SidedClassNamer destNamer = null;
-                                if( matching != null )
+                                int numMatches = 0;
+                                do
                                 {
-                                        // build a class namer
-                                        ClassNamer namer = new ClassNamer( matching.getUniqueMatches() );
-                                        sourceNamer = namer.getSourceNamer();
-                                        destNamer = namer.getDestNamer();
+                                        SidedClassNamer sourceNamer = null;
+                                        SidedClassNamer destNamer = null;
+                                        if( matching != null )
+                                        {
+                                                // build a class namer
+                                                ClassNamer namer = new ClassNamer( matching.getUniqueMatches() );
+                                                sourceNamer = namer.getSourceNamer();
+                                                destNamer = namer.getDestNamer();
+                                                
+                                                // note the number of matches
+                                                numMatches = matching.getUniqueMatches().size();
+                                        }
                                         
-                                        // note the number of matches
-                                        numMatches = matching.getUniqueMatches().size();
-                                }
-                                
-                                // get the entries left to match
-                                Set<ClassEntry> sourceClassEntries = sourceIndex.getObfClassEntries();
-                                Set<ClassEntry> destClassEntries = destIndex.getObfClassEntries();
-                                if( matching != null )
-                                {
-                                        sourceClassEntries.clear();
-                                        destClassEntries.clear();
-                                        for( Map.Entry<List<ClassIdentity>,List<ClassIdentity>> entry : matching.getAmbiguousMatches().entrySet() )
+                                        // get the entries left to match
+                                        Set<ClassEntry> sourceClassEntries = sourceIndex.getObfClassEntries();
+                                        Set<ClassEntry> destClassEntries = destIndex.getObfClassEntries();
+                                        if( matching != null )
                                         {
-                                                for( ClassIdentity c : entry.getKey() )
+                                                sourceClassEntries.clear();
+                                                destClassEntries.clear();
+                                                for( Map.Entry<List<ClassIdentity>,List<ClassIdentity>> entry : matching.getAmbiguousMatches().entrySet() )
+                                                {
+                                                        for( ClassIdentity c : entry.getKey() )
+                                                        {
+                                                                sourceClassEntries.add( c.getClassEntry() );
+                                                                matching.removeSource( c );
+                                                        }
+                                                        for( ClassIdentity c : entry.getValue() )
+                                                        {
+                                                                destClassEntries.add( c.getClassEntry() );
+                                                                matching.removeDest( c );
+                                                        }
+                                                }
+                                                for( ClassIdentity c : matching.getUnmatchedSourceClasses() )
                                                 {
                                                         sourceClassEntries.add( c.getClassEntry() );
                                                         matching.removeSource( c );
                                                 }
-                                                for( ClassIdentity c : entry.getValue() )
+                                                for( ClassIdentity c : matching.getUnmatchedDestClasses() )
                                                 {
                                                         destClassEntries.add( c.getClassEntry() );
                                                         matching.removeDest( c );
                                                 }
                                         }
-                                        for( ClassIdentity c : matching.getUnmatchedSourceClasses() )
+                                        else
                                         {
-                                                sourceClassEntries.add( c.getClassEntry() );
-                                                matching.removeSource( c );
+                                                matching = new ClassMatching();
                                         }
-                                        for( ClassIdentity c : matching.getUnmatchedDestClasses() )
+                                        
+                                        // compute a matching for the classes
+                                        for( ClassEntry classEntry : sourceClassEntries )
                                         {
-                                                destClassEntries.add( c.getClassEntry() );
-                                                matching.removeDest( c );
+                                                CtClass c = sourceLoader.loadClass( classEntry.getName() );
+                                                ClassIdentity sourceClass = new ClassIdentity( c, sourceNamer, sourceIndex, useReferences, useRawNames );
+                                                matching.addSource( sourceClass );
                                         }
+                                        for( ClassEntry classEntry : destClassEntries )
+                                        {
+                                                CtClass c = destLoader.loadClass( classEntry.getName() );
+                                                ClassIdentity destClass = new ClassIdentity( c, destNamer, destIndex, useReferences, useRawNames );
+                                                matching.matchDestClass( destClass );
+                                        }
+                                        
+                                        // TEMP
+                                        System.out.println( matching );
                                 }
-                                else
-                                {
-                                        matching = new ClassMatching();
-                                }
-                                
-                                // compute a matching for the classes
-                                for( ClassEntry classEntry : sourceClassEntries )
-                                {
-                                        CtClass c = sourceLoader.loadClass( classEntry.getName() );
-                                        ClassIdentity sourceClass = new ClassIdentity( c, sourceNamer, sourceIndex, useReferences );
-                                        matching.addSource( sourceClass );
-                                }
-                                for( ClassEntry classEntry : destClassEntries )
-                                {
-                                        CtClass c = destLoader.loadClass( classEntry.getName() );
-                                        ClassIdentity destClass = new ClassIdentity( c, destNamer, destIndex, useReferences );
-                                        matching.matchDestClass( destClass );
-                                }
-                                
-                                // TEMP
-                                System.out.println( matching );
+                                while( matching.getUniqueMatches().size() - numMatches > 0 );
                         }
-                        while( matching.getUniqueMatches().size() - numMatches > 0 );
                 }
                 
                 /* DEBUG: show some ambiguous matches
@@ -163,6 +264,24 @@ public class ClassMapper
                 return matching;
         }
         
+        private static void addFinalConversion( Map<String,String> finalConversion, ClassIdentity sourceClass, ClassIdentity destClass )
+        {
+                // flatten inner classes since these are all obf classes in the none package
+                String sourceClassName = sourceClass.getClassEntry().getName();
+                if( sourceClass.getClassEntry().isInnerClass() )
+                {
+                        sourceClassName = Constants.NonePackage + "/" + sourceClass.getClassEntry().getInnerClassName();
+                }
+                
+                String destClassName = destClass.getClassEntry().getName();
+                if( destClass.getClassEntry().isInnerClass() )
+                {
+                        destClassName = Constants.NonePackage + "/" + destClass.getClassEntry().getInnerClassName();
+                }
+                
+                finalConversion.put( sourceClassName, destClassName );
+        }
+        
         /* DEBUG
         private static String decompile( TranslatingTypeLoader loader, ClassEntry classEntry )
         {
diff --git a/src/cuchaz/enigma/convert/ClassMatching.java b/src/cuchaz/enigma/convert/ClassMatching.java
index fea84386..4e9fe398 100644
--- a/src/cuchaz/enigma/convert/ClassMatching.java
+++ b/src/cuchaz/enigma/convert/ClassMatching.java
@@ -10,12 +10,15 @@
  ******************************************************************************/
 package cuchaz.enigma.convert;
 
+import java.util.AbstractMap;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 
 import com.beust.jcommander.internal.Lists;
+import com.beust.jcommander.internal.Maps;
 import com.google.common.collect.ArrayListMultimap;
 import com.google.common.collect.BiMap;
 import com.google.common.collect.HashBiMap;
@@ -156,6 +159,36 @@ public class ClassMatching
                 return new ArrayList<ClassIdentity>( m_unmatchedDestClasses );
         }
         
+        public Map<String,Map.Entry<ClassIdentity,List<ClassIdentity>>> getConversionMap( )
+        {
+                Map<String,Map.Entry<ClassIdentity,List<ClassIdentity>>> conversion = Maps.newHashMap();
+                for( Map.Entry<ClassIdentity,ClassIdentity> entry : getUniqueMatches().entrySet() )
+                {
+                        conversion.put(
+                                entry.getKey().getClassEntry().getName(),
+                                new AbstractMap.SimpleEntry<ClassIdentity,List<ClassIdentity>>( entry.getKey(), Arrays.asList( entry.getValue() ) )
+                        );
+                }
+                for( Map.Entry<List<ClassIdentity>,List<ClassIdentity>> entry : getAmbiguousMatches().entrySet() )
+                {
+                        for( ClassIdentity sourceClass : entry.getKey() )
+                        {
+                                conversion.put(
+                                        sourceClass.getClassEntry().getName(),
+                                        new AbstractMap.SimpleEntry<ClassIdentity,List<ClassIdentity>>( sourceClass, entry.getValue() )
+                                );
+                        }
+                }
+                for( ClassIdentity sourceClass : getUnmatchedSourceClasses() )
+                {
+                        conversion.put(
+                                sourceClass.getClassEntry().getName(),
+                                new AbstractMap.SimpleEntry<ClassIdentity,List<ClassIdentity>>( sourceClass, new ArrayList<ClassIdentity>() )
+                        );
+                }
+                return conversion;
+        }
+        
         @Override
         public String toString( )
         {
diff --git a/src/cuchaz/enigma/mapping/ClassMapping.java b/src/cuchaz/enigma/mapping/ClassMapping.java
index 59365129..095cb385 100644
--- a/src/cuchaz/enigma/mapping/ClassMapping.java
+++ b/src/cuchaz/enigma/mapping/ClassMapping.java
@@ -11,6 +11,7 @@
 package cuchaz.enigma.mapping;
 
 import java.io.Serializable;
+import java.util.ArrayList;
 import java.util.Map;
 
 import com.google.common.collect.Maps;
@@ -299,4 +300,32 @@ public class ClassMapping implements Serializable, Comparable<ClassMapping>
         {
                 return m_obfName.compareTo( other.m_obfName );
         }
+        
+        public void renameObfClasses( Map<String,String> nameMap )
+        {
+                // rename self
+                {
+                        String newName = nameMap.get( m_obfName );
+                        if( newName != null )
+                        {
+                                m_obfName = newName;
+                        }
+                }
+                
+                // rename inner classes
+                for( ClassMapping classMapping : new ArrayList<ClassMapping>( m_innerClassesByObf.values() ) )
+                {
+                        m_innerClassesByObf.remove( classMapping.getObfName() );
+                        classMapping.renameObfClasses( nameMap );
+                        m_innerClassesByObf.put( classMapping.getObfName(), classMapping );
+                }
+                
+                // rename method signatures
+                for( MethodMapping methodMapping : new ArrayList<MethodMapping>( m_methodsByObf.values() ) )
+                {
+                        m_methodsByObf.remove( getMethodKey( methodMapping.getObfName(), methodMapping.getObfSignature() ) );
+                        methodMapping.renameObfClasses( nameMap );
+                        m_methodsByObf.put( getMethodKey( methodMapping.getObfName(), methodMapping.getObfSignature() ), methodMapping );
+                }
+        }
 }
diff --git a/src/cuchaz/enigma/mapping/Mappings.java b/src/cuchaz/enigma/mapping/Mappings.java
index f3b8fad1..70bea25f 100644
--- a/src/cuchaz/enigma/mapping/Mappings.java
+++ b/src/cuchaz/enigma/mapping/Mappings.java
@@ -14,6 +14,7 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.io.ObjectInputStream;
 import java.io.Serializable;
+import java.util.ArrayList;
 import java.util.Map;
 import java.util.zip.GZIPInputStream;
 
@@ -136,4 +137,28 @@ public class Mappings implements Serializable
                 }
                 return buf.toString();
         }
+        
+        public void renameObfClasses( Map<String,String> nameMap )
+        {
+                for( ClassMapping classMapping : new ArrayList<ClassMapping>( m_classesByObf.values() ) )
+                {
+                        String newName = nameMap.get( classMapping.getObfName() );
+                        if( newName != null )
+                        {
+                                m_classesByObf.remove( classMapping.getObfName() );
+                                classMapping.renameObfClasses( nameMap );
+                                m_classesByObf.put( classMapping.getObfName(), classMapping );
+                        }
+                }
+        }
+
+        public void removeClassByObfName( String obfName )
+        {
+                ClassMapping classMapping = m_classesByObf.get( obfName );
+                if( classMapping != null )
+                {
+                        m_classesByObf.remove( classMapping.getObfName() );
+                        m_classesByDeobf.remove( classMapping.getDeobfName() );
+                }
+        }
 }
diff --git a/src/cuchaz/enigma/mapping/MethodMapping.java b/src/cuchaz/enigma/mapping/MethodMapping.java
index 6e6bec46..fe4e29b2 100644
--- a/src/cuchaz/enigma/mapping/MethodMapping.java
+++ b/src/cuchaz/enigma/mapping/MethodMapping.java
@@ -14,6 +14,8 @@ import java.io.Serializable;
 import java.util.Map;
 import java.util.TreeMap;
 
+import cuchaz.enigma.mapping.SignatureUpdater.ClassNameUpdater;
+
 public class MethodMapping implements Serializable, Comparable<MethodMapping>
 {
         private static final long serialVersionUID = -4409570216084263978L;
@@ -139,4 +141,22 @@ public class MethodMapping implements Serializable, Comparable<MethodMapping>
         {
                 return ( m_obfName + m_obfSignature ).compareTo( ( other.m_obfName + other.m_obfSignature ) );
         }
+
+        public void renameObfClasses( final Map<String,String> nameMap )
+        {
+                // rename obf classes in the signature
+                m_obfSignature = SignatureUpdater.update( m_obfSignature, new ClassNameUpdater( )
+                {
+                        @Override
+                        public String update( String className )
+                        {
+                                String newName = nameMap.get( className );
+                                if( newName != null )
+                                {
+                                        return newName;
+                                }
+                                return className;
+                        }
+                } );
+        }
 }