fixed mapping conversion bug with class rename order

5 files changed, 103 insertions, 62 deletions

diff --git a/src/cuchaz/enigma/convert/ClassMatcher.java b/src/cuchaz/enigma/convert/ClassMatcher.java
index b551da25..0821bd3a 100644
--- a/src/cuchaz/enigma/convert/ClassMatcher.java
+++ b/src/cuchaz/enigma/convert/ClassMatcher.java
@@ -19,6 +19,8 @@ import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -89,9 +91,7 @@ public class ClassMatcher
 
                 // compute the matching
                 ClassMatching matching = computeMatching( sourceIndex, sourceLoader, destIndex, destLoader );
-                
-                // start the class conversion map with the unique and ambiguous matchings
-                Map<String,Map.Entry<ClassIdentity,List<ClassIdentity>>> conversionMap = matching.getConversionMap();
+                Map<String,Map.Entry<ClassIdentity,List<ClassIdentity>>> matchingIndex = matching.getIndex();
                 
                 // get all the obf class names used in the mappings
                 Set<String> usedClassNames = mappings.getAllObfClassNames();
@@ -101,9 +101,10 @@ public class ClassMatcher
                         allClassNames.add( classEntry.getName() );
                 }
                 usedClassNames.retainAll( allClassNames );
+                System.out.println( "Used " + usedClassNames.size() + " classes in the mappings" );
                 
                 // probabilistically match the non-uniquely-matched source classes
-                for( Map.Entry<ClassIdentity,List<ClassIdentity>> entry : conversionMap.values() )
+                for( Map.Entry<ClassIdentity,List<ClassIdentity>> entry : matchingIndex.values() )
                 {
                         ClassIdentity sourceClass = entry.getKey();
                         List<ClassIdentity> destClasses = entry.getValue();
@@ -148,20 +149,20 @@ public class ClassMatcher
                         }
                 }
                 
-                // use the matching to convert the mappings
-                BiMap<String,String> classConversion = HashBiMap.create();
-                Set<String> unmatchedSourceClasses = Sets.newHashSet();
+                // group the matching into unique and non-unique matches
+                BiMap<String,String> matchedClassNames = HashBiMap.create();
+                Set<String> unmatchedSourceClassNames = Sets.newHashSet();
                 for( String className : usedClassNames )
                 {
                         // is there a match for this class?
-                        Map.Entry<ClassIdentity,List<ClassIdentity>> entry = conversionMap.get( className );
+                        Map.Entry<ClassIdentity,List<ClassIdentity>> entry = matchingIndex.get( className );
                         ClassIdentity sourceClass = entry.getKey();
                         List<ClassIdentity> matches = entry.getValue();
                         
                         if( matches.size() == 1 )
                         {
                                 // unique match! We're good to go!
-                                classConversion.put(
+                                matchedClassNames.put(
                                         sourceClass.getClassEntry().getName(),
                                         matches.get( 0 ).getClassEntry().getName()
                                 );
@@ -172,35 +173,36 @@ public class ClassMatcher
                                 String fallbackMatch = fallbackMatching.get( className );
                                 if( fallbackMatch != null )
                                 {
-                                        classConversion.put(
+                                        matchedClassNames.put(
                                                 sourceClass.getClassEntry().getName(),
                                                 fallbackMatch
                                         );
                                 }
                                 else
                                 {
-                                        unmatchedSourceClasses.add( className );
+                                        unmatchedSourceClassNames.add( className );
                                 }
                         }
                 }
                 
-                // remove (and warn about) unmatched classes
-                if( !unmatchedSourceClasses.isEmpty() )
+                // report unmatched classes
+                if( !unmatchedSourceClassNames.isEmpty() )
                 {
-                        System.err.println( "WARNING: there were unmatched classes!" );
-                        for( String className : unmatchedSourceClasses )
+                        System.err.println( "ERROR: there were unmatched classes!" );
+                        for( String className : unmatchedSourceClassNames )
                         {
                                 System.err.println( "\t" + className );
-                                mappings.removeClassByObfName( className );
                         }
-                        System.err.println( "Mappings for these classes have been removed." );
+                        return;
                 }
                 
-                // show the class name changes
-                for( Map.Entry<String,String> entry : classConversion.entrySet() )
+                // get the class name changes from the matched class names
+                Map<String,String> classChanges = Maps.newHashMap();
+                for( Map.Entry<String,String> entry : matchedClassNames.entrySet() )
                 {
                         if( !entry.getKey().equals( entry.getValue() ) )
                         {
+                                classChanges.put( entry.getKey(), entry.getValue() );
                                 System.out.println( String.format( "Class change: %s -> %s", entry.getKey(), entry.getValue() ) );
                                 /* DEBUG
                                 System.out.println( String.format( "\n%s\n%s",
@@ -211,8 +213,45 @@ public class ClassMatcher
                         }
                 }
                 
-                // convert the mappings
-                mappings.renameObfClasses( classConversion );
+                // sort the changes so classes are renamed in the correct order
+                // ie. if we have the mappings a->b, b->c, we have to apply b->c before a->b
+                LinkedHashMap<String,String> orderedClassChanges = Maps.newLinkedHashMap();
+                int numChangesLeft = classChanges.size();
+                while( !classChanges.isEmpty() )
+                {
+                        Iterator<Map.Entry<String,String>> iter = classChanges.entrySet().iterator();
+                        while( iter.hasNext() )
+                        {
+                                Map.Entry<String,String> entry = iter.next();
+                                if( classChanges.get( entry.getValue() ) == null )
+                                {
+                                        orderedClassChanges.put( entry.getKey(), entry.getValue() );
+                                        iter.remove();
+                                }
+                        }
+                        
+                        // did we remove any changes?
+                        if( numChangesLeft - classChanges.size() > 0 )
+                        {
+                                // keep going
+                                numChangesLeft = classChanges.size();                           
+                        }
+                        else
+                        {
+                                // can't sort anymore. There must be a loop
+                                break;
+                        }
+                }
+                if( classChanges.size() > 0 )
+                {
+                        throw new Error( String.format( "Unable to sort %d/%d class changes!", classChanges.size(), matchedClassNames.size() ) );
+                }
+                
+                // convert the mappings in the correct class order
+                for( Map.Entry<String,String> entry : orderedClassChanges.entrySet() )
+                {
+                        mappings.renameObfClass( entry.getKey(), entry.getValue() );
+                }
                 
                 // check the method matches
                 System.out.println( "Checking methods..." );
@@ -250,7 +289,7 @@ public class ClassMatcher
                                         }
                                         
                                         System.err.println( "\tAvailable source methods:" );
-                                        c = sourceLoader.loadClass( classConversion.inverse().get( classMapping.getObfName() ) );
+                                        c = sourceLoader.loadClass( matchedClassNames.inverse().get( classMapping.getObfName() ) );
                                         for( CtBehavior behavior : c.getDeclaredBehaviors() )
                                         {
                                                 MethodEntry declaredMethodEntry = new MethodEntry(
diff --git a/src/cuchaz/enigma/convert/ClassMatching.java b/src/cuchaz/enigma/convert/ClassMatching.java
index 6ce8c887..5511902d 100644
--- a/src/cuchaz/enigma/convert/ClassMatching.java
+++ b/src/cuchaz/enigma/convert/ClassMatching.java
@@ -159,7 +159,7 @@ public class ClassMatching
                 return new ArrayList<ClassIdentity>( m_unmatchedDestClasses );
         }
         
-        public Map<String,Map.Entry<ClassIdentity,List<ClassIdentity>>> getConversionMap( )
+        public Map<String,Map.Entry<ClassIdentity,List<ClassIdentity>>> getIndex( )
         {
                 Map<String,Map.Entry<ClassIdentity,List<ClassIdentity>>> conversion = Maps.newHashMap();
                 for( Map.Entry<ClassIdentity,ClassIdentity> entry : getUniqueMatches().entrySet() )
diff --git a/src/cuchaz/enigma/mapping/ClassMapping.java b/src/cuchaz/enigma/mapping/ClassMapping.java
index 095cb385..1219e7ca 100644
--- a/src/cuchaz/enigma/mapping/ClassMapping.java
+++ b/src/cuchaz/enigma/mapping/ClassMapping.java
@@ -301,31 +301,37 @@ public class ClassMapping implements Serializable, Comparable<ClassMapping>
                 return m_obfName.compareTo( other.m_obfName );
         }
         
-        public void renameObfClasses( Map<String,String> nameMap )
+        public boolean renameObfClass( String oldObfClassName, String newObfClassName )
         {
-                // rename self
+                // rename inner classes
+                for( ClassMapping innerClassMapping : new ArrayList<ClassMapping>( m_innerClassesByObf.values() ) )
                 {
-                        String newName = nameMap.get( m_obfName );
-                        if( newName != null )
+                        if( innerClassMapping.renameObfClass( oldObfClassName, newObfClassName ) )
                         {
-                                m_obfName = newName;
+                                m_innerClassesByObf.remove( oldObfClassName );
+                                m_innerClassesByObf.put( newObfClassName, innerClassMapping );
+                                assert( m_innerClassesByObf.size() == m_innerClassesByDeobf.size() );
                         }
                 }
                 
-                // rename inner classes
-                for( ClassMapping classMapping : new ArrayList<ClassMapping>( m_innerClassesByObf.values() ) )
+                // rename method signatures
+                for( MethodMapping methodMapping : new ArrayList<MethodMapping>( m_methodsByObf.values() ) )
                 {
-                        m_innerClassesByObf.remove( classMapping.getObfName() );
-                        classMapping.renameObfClasses( nameMap );
-                        m_innerClassesByObf.put( classMapping.getObfName(), classMapping );
+                        String oldMethodKey = getMethodKey( methodMapping.getObfName(), methodMapping.getObfSignature() );
+                        if( methodMapping.renameObfClass( oldObfClassName, newObfClassName ) )
+                        {
+                                m_methodsByObf.remove( oldMethodKey );
+                                m_methodsByObf.put( getMethodKey( methodMapping.getObfName(), methodMapping.getObfSignature() ), methodMapping );
+                                assert( m_methodsByObf.size() == m_methodsByDeobf.size() );
+                        }
                 }
                 
-                // rename method signatures
-                for( MethodMapping methodMapping : new ArrayList<MethodMapping>( m_methodsByObf.values() ) )
+                if( m_obfName.equals( oldObfClassName ) )
                 {
-                        m_methodsByObf.remove( getMethodKey( methodMapping.getObfName(), methodMapping.getObfSignature() ) );
-                        methodMapping.renameObfClasses( nameMap );
-                        m_methodsByObf.put( getMethodKey( methodMapping.getObfName(), methodMapping.getObfSignature() ), methodMapping );
+                        // rename this class
+                        m_obfName = newObfClassName;
+                        return true;
                 }
+                return false;
         }
 }
diff --git a/src/cuchaz/enigma/mapping/Mappings.java b/src/cuchaz/enigma/mapping/Mappings.java
index 4b47d160..0b03abb0 100644
--- a/src/cuchaz/enigma/mapping/Mappings.java
+++ b/src/cuchaz/enigma/mapping/Mappings.java
@@ -15,6 +15,7 @@ import java.io.InputStream;
 import java.io.ObjectInputStream;
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Map;
 import java.util.Set;
 import java.util.zip.GZIPInputStream;
@@ -66,7 +67,7 @@ public class Mappings implements Serializable
                 }
         }
         
-        public Iterable<ClassMapping> classes( )
+        public Collection<ClassMapping> classes( )
         {
                 assert( m_classesByObf.size() == m_classesByDeobf.size() );
                 return m_classesByObf.values();
@@ -141,30 +142,19 @@ public class Mappings implements Serializable
                 return buf.toString();
         }
         
-        public void renameObfClasses( Map<String,String> nameMap )
+        public void renameObfClass( String oldObfName, String newObfName )
         {
-                for( ClassMapping classMapping : new ArrayList<ClassMapping>( m_classesByObf.values() ) )
+                for( ClassMapping classMapping : new ArrayList<ClassMapping>( classes() ) )
                 {
-                        String newName = nameMap.get( classMapping.getObfName() );
-                        if( newName != null )
+                        if( classMapping.renameObfClass( oldObfName, newObfName ) )
                         {
-                                m_classesByObf.remove( classMapping.getObfName() );
-                                classMapping.renameObfClasses( nameMap );
-                                m_classesByObf.put( classMapping.getObfName(), classMapping );
+                                m_classesByObf.remove( oldObfName );
+                                m_classesByObf.put( newObfName, classMapping );
+                                assert( m_classesByObf.size() == m_classesByDeobf.size() );
                         }
                 }
         }
-
-        public void removeClassByObfName( String obfName )
-        {
-                ClassMapping classMapping = m_classesByObf.get( obfName );
-                if( classMapping != null )
-                {
-                        m_classesByObf.remove( classMapping.getObfName() );
-                        m_classesByDeobf.remove( classMapping.getDeobfName() );
-                }
-        }
-
+        
         public Set<String> getAllObfClassNames( )
         {
                 final Set<String> classNames = Sets.newHashSet();
diff --git a/src/cuchaz/enigma/mapping/MethodMapping.java b/src/cuchaz/enigma/mapping/MethodMapping.java
index fe4e29b2..b0f7ba84 100644
--- a/src/cuchaz/enigma/mapping/MethodMapping.java
+++ b/src/cuchaz/enigma/mapping/MethodMapping.java
@@ -142,21 +142,27 @@ 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 )
+        public boolean renameObfClass( final String oldObfClassName, final String newObfClassName )
         {
                 // rename obf classes in the signature
-                m_obfSignature = SignatureUpdater.update( m_obfSignature, new ClassNameUpdater( )
+                String newSignature = SignatureUpdater.update( m_obfSignature, new ClassNameUpdater( )
                 {
                         @Override
                         public String update( String className )
                         {
-                                String newName = nameMap.get( className );
-                                if( newName != null )
+                                if( className.equals( oldObfClassName ) )
                                 {
-                                        return newName;
+                                        return newObfClassName;
                                 }
                                 return className;
                         }
                 } );
+                
+                if( newSignature != m_obfSignature )
+                {
+                        m_obfSignature = newSignature;
+                        return true;
+                }
+                return false;
         }
 }