got a decent class matcher working

9 files changed, 748 insertions, 145 deletions

diff --git a/src/cuchaz/enigma/TranslatingTypeLoader.java b/src/cuchaz/enigma/TranslatingTypeLoader.java
index 1e0e95a..e70093e 100644
--- a/src/cuchaz/enigma/TranslatingTypeLoader.java
+++ b/src/cuchaz/enigma/TranslatingTypeLoader.java
@@ -47,6 +47,11 @@ public class TranslatingTypeLoader implements ITypeLoader
         private Map<String,byte[]> m_cache;
         private ClasspathTypeLoader m_defaultTypeLoader;
         
+        public TranslatingTypeLoader( JarFile jar, JarIndex jarIndex )
+        {
+                this( jar, jarIndex, new Translator(), new Translator() );
+        }
+        
         public TranslatingTypeLoader( JarFile jar, JarIndex jarIndex, Translator obfuscatingTranslator, Translator deobfuscatingTranslator )
         {
                 m_jar = jar;
@@ -90,6 +95,28 @@ public class TranslatingTypeLoader implements ITypeLoader
                 return true;
         }
         
+        public CtClass loadClass( String deobfClassName )
+        {
+                byte[] data = loadType( deobfClassName );
+                if( data == null )
+                {
+                        return null;
+                }
+                
+                // return a javassist handle for the class
+                String javaClassFileName = Descriptor.toJavaName( deobfClassName );
+                ClassPool classPool = new ClassPool();
+                classPool.insertClassPath( new ByteArrayClassPath( javaClassFileName, data ) );
+                try
+                {
+                        return classPool.get( javaClassFileName );
+                }
+                catch( NotFoundException ex )
+                {
+                        throw new Error( ex );
+                }
+        }
+        
         private byte[] loadType( String deobfClassName )
         {
                 // what class file should we actually load?
diff --git a/src/cuchaz/enigma/analysis/JarClassIterator.java b/src/cuchaz/enigma/analysis/JarClassIterator.java
index 6c9f124..10ae805 100644
--- a/src/cuchaz/enigma/analysis/JarClassIterator.java
+++ b/src/cuchaz/enigma/analysis/JarClassIterator.java
@@ -67,33 +67,7 @@ public class JarClassIterator implements Iterator<CtClass>
                 JarEntry entry = m_iter.next();
                 try
                 {
-                        // read the class into a buffer
-                        ByteArrayOutputStream bos = new ByteArrayOutputStream();
-                        byte[] buf = new byte[Constants.KiB];
-                        int totalNumBytesRead = 0;
-                        InputStream in = m_jar.getInputStream( entry );
-                        while( in.available() > 0 )
-                        {
-                                int numBytesRead = in.read( buf );
-                                if( numBytesRead < 0 )
-                                {
-                                        break;
-                                }
-                                bos.write( buf, 0, numBytesRead );
-                                
-                                // sanity checking
-                                totalNumBytesRead += numBytesRead;
-                                if( totalNumBytesRead > Constants.MiB )
-                                {
-                                        throw new Error( "Class file " + entry.getName() + " larger than 1 MiB! Something is wrong!" );
-                                }
-                        }
-                        
-                        // get a javassist handle for the class
-                        String className = Descriptor.toJavaName( getClassEntry( entry ).getName() );
-                        ClassPool classPool = new ClassPool();
-                        classPool.insertClassPath( new ByteArrayClassPath( className, bos.toByteArray() ) );
-                        return classPool.get( className );
+                        return getClass( m_jar, entry );
                 }
                 catch( IOException | NotFoundException ex )
                 {
@@ -136,6 +110,50 @@ public class JarClassIterator implements Iterator<CtClass>
                 };
         }
         
+        public static CtClass getClass( JarFile jar, ClassEntry classEntry )
+        {
+                try
+                {
+                        return getClass( jar, new JarEntry( classEntry.getName() + ".class" ) );
+                }
+                catch( IOException | NotFoundException ex )
+                {
+                        throw new Error( "Unable to load class: " + classEntry.getName() );
+                }
+        }
+        
+        private static CtClass getClass( JarFile jar, JarEntry entry )
+        throws IOException, NotFoundException
+        {
+                // read the class into a buffer
+                ByteArrayOutputStream bos = new ByteArrayOutputStream();
+                byte[] buf = new byte[Constants.KiB];
+                int totalNumBytesRead = 0;
+                InputStream in = jar.getInputStream( entry );
+                while( in.available() > 0 )
+                {
+                        int numBytesRead = in.read( buf );
+                        if( numBytesRead < 0 )
+                        {
+                                break;
+                        }
+                        bos.write( buf, 0, numBytesRead );
+                        
+                        // sanity checking
+                        totalNumBytesRead += numBytesRead;
+                        if( totalNumBytesRead > Constants.MiB )
+                        {
+                                throw new Error( "Class file " + entry.getName() + " larger than 1 MiB! Something is wrong!" );
+                        }
+                }
+                
+                // get a javassist handle for the class
+                String className = Descriptor.toJavaName( getClassEntry( entry ).getName() );
+                ClassPool classPool = new ClassPool();
+                classPool.insertClassPath( new ByteArrayClassPath( className, bos.toByteArray() ) );
+                return classPool.get( className );
+        }
+        
         private static ClassEntry getClassEntry( JarEntry entry )
         {
                 return new ClassEntry( entry.getName().substring( 0, entry.getName().length() - ".class".length() ) );
diff --git a/src/cuchaz/enigma/analysis/JarIndex.java b/src/cuchaz/enigma/analysis/JarIndex.java
index deacf16..b479b69 100644
--- a/src/cuchaz/enigma/analysis/JarIndex.java
+++ b/src/cuchaz/enigma/analysis/JarIndex.java
@@ -95,7 +95,7 @@ public class JarIndex
                 // step 2: index method/field access
                 for( CtClass c : JarClassIterator.classes( jar ) )
                 {
-                        fixClass( c );
+                        ClassRenamer.moveAllClassesOutOfDefaultPackage( c, Constants.NonePackage );
                         ClassEntry classEntry = new ClassEntry( Descriptor.toJvmName( c.getName() ) );
                         for( CtField field : c.getDeclaredFields() )
                         {
@@ -112,7 +112,7 @@ public class JarIndex
                 // step 3: index the types, methods
                 for( CtClass c : JarClassIterator.classes( jar ) )
                 {
-                        fixClass( c );
+                        ClassRenamer.moveAllClassesOutOfDefaultPackage( c, Constants.NonePackage );
                         String className = Descriptor.toJvmName( c.getName() );
                         m_ancestries.addSuperclass( className, Descriptor.toJvmName( c.getClassFile().getSuperclass() ) );
                         for( String interfaceName : c.getClassFile().getInterfaces() )
@@ -128,8 +128,7 @@ public class JarIndex
                 // step 4: index inner classes and anonymous classes
                 for( CtClass c : JarClassIterator.classes( jar ) )
                 {
-                        fixClass( c );
-                        
+                        ClassRenamer.moveAllClassesOutOfDefaultPackage( c, Constants.NonePackage );
                         String outerClassName = findOuterClass( c );
                         if( outerClassName != null )
                         {
@@ -164,17 +163,6 @@ public class JarIndex
                 renameMethods( m_bridgeMethods );
         }
 
-        private void fixClass( CtClass c )
-        {
-                ClassEntry classEntry = new ClassEntry( Descriptor.toJvmName( c.getName() ) );
-                if( classEntry.isInDefaultPackage() )
-                {
-                        // move class out of default package
-                        classEntry = new ClassEntry( Constants.NonePackage + "/" + classEntry.getName() );
-                        ClassRenamer.moveAllClassesOutOfDefaultPackage( c, Constants.NonePackage );
-                }
-        }
-
         private void indexBehavior( CtBehavior behavior )
         {
                 // get the method entry
@@ -729,6 +717,22 @@ public class JarIndex
         
         private void renameClasses( Map<String,String> renames )
         {
+                // rename class entries
+                Set<ClassEntry> obfClassEntries = Sets.newHashSet();
+                for( ClassEntry classEntry : m_obfClassEntries )
+                {
+                        if( renames.containsKey( classEntry.getName() ) )
+                        {
+                                obfClassEntries.add( new ClassEntry( renames.get( classEntry.getName() ) ) );
+                        }
+                        else
+                        {
+                                obfClassEntries.add( classEntry );
+                        }
+                }
+                m_obfClassEntries = obfClassEntries;
+                
+                // rename others
                 m_ancestries.renameClasses( renames );
                 renameClassesInMultimap( renames, m_methodImplementations );
                 renameClassesInMultimap( renames, m_behaviorReferences );
diff --git a/src/cuchaz/enigma/bytecode/ClassRenamer.java b/src/cuchaz/enigma/bytecode/ClassRenamer.java
index f3a8c0e..efe22a1 100644
--- a/src/cuchaz/enigma/bytecode/ClassRenamer.java
+++ b/src/cuchaz/enigma/bytecode/ClassRenamer.java
@@ -14,6 +14,7 @@ import java.util.Map;
 import java.util.Set;
 
 import javassist.ClassMap;
+import javassist.CtBehavior;
 import javassist.CtClass;
 import javassist.bytecode.ConstPool;
 import javassist.bytecode.Descriptor;
@@ -23,6 +24,8 @@ import com.beust.jcommander.internal.Sets;
 import com.google.common.collect.Maps;
 
 import cuchaz.enigma.mapping.ClassEntry;
+import cuchaz.enigma.mapping.SignatureUpdater;
+import cuchaz.enigma.mapping.SignatureUpdater.ClassNameUpdater;
 
 public class ClassRenamer
 {
@@ -115,5 +118,37 @@ public class ClassRenamer
                         }
                 }
                 ClassRenamer.renameClasses( c, map );
+                
+                // TEMP
+                for( ClassEntry classEntry : ClassRenamer.getAllClassEntries( c ) )
+                {
+                        if( classEntry.isInDefaultPackage() )
+                        {
+                                throw new Error( "!!! " + classEntry );
+                        }
+                }
+                
+                // TEMP
+                for( CtBehavior behavior : c.getDeclaredBehaviors() )
+                {
+                        if( behavior.getSignature() == null )
+                        {
+                                continue;
+                        }
+                        
+                        SignatureUpdater.update( behavior.getSignature(), new ClassNameUpdater( )
+                        {
+                                @Override
+                                public String update( String className )
+                                {
+                                        ClassEntry classEntry = new ClassEntry( className );
+                                        if( classEntry.isInDefaultPackage() )
+                                        {
+                                                throw new Error( "!!! " + className );
+                                        }
+                                        return className;
+                                }
+                        } );
+                }
         }
 }
diff --git a/src/cuchaz/enigma/convert/ClassIdentity.java b/src/cuchaz/enigma/convert/ClassIdentity.java
index aecf7fc..0a3a449 100644
--- a/src/cuchaz/enigma/convert/ClassIdentity.java
+++ b/src/cuchaz/enigma/convert/ClassIdentity.java
@@ -13,10 +13,12 @@ package cuchaz.enigma.convert;
 import java.io.UnsupportedEncodingException;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
+import java.util.Enumeration;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
+import javassist.CannotCompileException;
 import javassist.CtBehavior;
 import javassist.CtClass;
 import javassist.CtConstructor;
@@ -27,34 +29,58 @@ import javassist.bytecode.CodeIterator;
 import javassist.bytecode.ConstPool;
 import javassist.bytecode.Descriptor;
 import javassist.bytecode.Opcode;
+import javassist.expr.ConstructorCall;
+import javassist.expr.ExprEditor;
+import javassist.expr.FieldAccess;
+import javassist.expr.MethodCall;
+import javassist.expr.NewExpr;
 
 import com.beust.jcommander.internal.Maps;
 import com.beust.jcommander.internal.Sets;
 import com.google.common.collect.Lists;
 
+import cuchaz.enigma.Constants;
 import cuchaz.enigma.Util;
+import cuchaz.enigma.analysis.ClassImplementationsTreeNode;
+import cuchaz.enigma.analysis.EntryReference;
+import cuchaz.enigma.analysis.JarIndex;
 import cuchaz.enigma.bytecode.ConstPoolEditor;
 import cuchaz.enigma.bytecode.InfoType;
 import cuchaz.enigma.bytecode.accessors.ConstInfoAccessor;
+import cuchaz.enigma.convert.ClassNamer.SidedClassNamer;
+import cuchaz.enigma.mapping.BehaviorEntry;
 import cuchaz.enigma.mapping.ClassEntry;
+import cuchaz.enigma.mapping.ConstructorEntry;
+import cuchaz.enigma.mapping.Entry;
+import cuchaz.enigma.mapping.FieldEntry;
+import cuchaz.enigma.mapping.MethodEntry;
 import cuchaz.enigma.mapping.SignatureUpdater;
 import cuchaz.enigma.mapping.SignatureUpdater.ClassNameUpdater;
 
 public class ClassIdentity
 {
         private ClassEntry m_classEntry;
+        private SidedClassNamer m_namer;
         private Set<String> m_fields;
         private Set<String> m_methods;
         private Set<String> m_constructors;
         private String m_staticInitializer;
+        private String m_extends;
+        private Set<String> m_implements;
+        private Set<String> m_implementations;
+        private Set<String> m_references;
         
-        public ClassIdentity( CtClass c )
+        public ClassIdentity( CtClass c, SidedClassNamer namer, JarIndex index, boolean useReferences )
         {
+                m_namer = namer;
+                
+                // stuff from the bytecode
+                
                 m_classEntry = new ClassEntry( Descriptor.toJvmName( c.getName() ) );
                 m_fields = Sets.newHashSet();
                 for( CtField field : c.getDeclaredFields() )
                 {
-                        m_fields.add( scrubSignature( scrubSignature( field.getSignature() ) ) );
+                        m_fields.add( scrubSignature( field.getSignature() ) );
                 }
                 m_methods = Sets.newHashSet();
                 for( CtMethod method : c.getDeclaredMethods() )
@@ -71,6 +97,73 @@ public class ClassIdentity
                 {
                         m_staticInitializer = getBehaviorSignature( c.getClassInitializer() );
                 }
+                m_extends = "";
+                if( c.getClassFile().getSuperclass() != null )
+                {
+                        m_extends = scrubClassName( c.getClassFile().getSuperclass() );
+                }
+                m_implements = Sets.newHashSet();
+                for( String interfaceName : c.getClassFile().getInterfaces() )
+                {
+                        m_implements.add( scrubClassName( interfaceName ) );
+                }
+                
+                // stuff from the jar index
+                
+                m_implementations = Sets.newHashSet();
+                @SuppressWarnings( "unchecked" )
+                Enumeration<ClassImplementationsTreeNode> implementations = index.getClassImplementations( null, m_classEntry ).children();
+                while( implementations.hasMoreElements() )
+                {
+                        ClassImplementationsTreeNode node = implementations.nextElement();
+                        m_implementations.add( scrubClassName( node.getClassEntry().getName() ) );
+                }
+                
+                m_references = Sets.newHashSet();
+                if( useReferences )
+                {
+                        for( CtField field : c.getDeclaredFields() )
+                        {
+                                FieldEntry fieldEntry = new FieldEntry( m_classEntry, field.getName() );
+                                for( EntryReference<FieldEntry,BehaviorEntry> reference : index.getFieldReferences( fieldEntry ) )
+                                {
+                                        addReference( reference );
+                                }
+                        }
+                        for( CtMethod method : c.getDeclaredMethods() )
+                        {
+                                MethodEntry methodEntry = new MethodEntry( m_classEntry, method.getName(), method.getSignature() );
+                                for( EntryReference<BehaviorEntry,BehaviorEntry> reference : index.getBehaviorReferences( methodEntry ) )
+                                {
+                                        addReference( reference );
+                                }
+                        }
+                        for( CtConstructor constructor : c.getDeclaredConstructors() )
+                        {
+                                ConstructorEntry constructorEntry = new ConstructorEntry( m_classEntry, constructor.getSignature() );
+                                for( EntryReference<BehaviorEntry,BehaviorEntry> reference : index.getBehaviorReferences( constructorEntry ) )
+                                {
+                                        addReference( reference );
+                                }
+                        }
+                }
+        }
+
+        private void addReference( EntryReference<? extends Entry,BehaviorEntry> reference )
+        {
+                if( reference.context.getSignature() != null )
+                {
+                        m_references.add( String.format( "%s_%s",
+                                scrubClassName( reference.context.getClassName() ),
+                                scrubSignature( reference.context.getSignature() )
+                        ) );
+                }
+                else
+                {
+                        m_references.add( String.format( "%s_<clinit>",
+                                scrubClassName( reference.context.getClassName() )
+                        ) );
+                }
         }
 
         public ClassEntry getClassEntry( )
@@ -109,9 +202,38 @@ public class ClassIdentity
                         buf.append( m_staticInitializer );
                         buf.append( "\n" );
                 }
+                if( m_extends.length() > 0 )
+                {
+                        buf.append( "\textends " );
+                        buf.append( m_extends );
+                        buf.append( "\n" );
+                }
+                for( String interfaceName : m_implements )
+                {
+                        buf.append( "\timplements " );
+                        buf.append( interfaceName );
+                        buf.append( "\n" );
+                }
+                for( String implementation : m_implementations )
+                {
+                        buf.append( "\timplemented by " );
+                        buf.append( implementation );
+                        buf.append( "\n" );
+                }
+                for( String reference : m_references )
+                {
+                        buf.append( "\treference " );
+                        buf.append( reference );
+                        buf.append( "\n" );
+                }
                 return buf.toString();
         }
         
+        private String scrubClassName( String className )
+        {
+                return scrubSignature( "L" + Descriptor.toJvmName( className ) + ";" );
+        }
+        
         private String scrubSignature( String signature )
         {
                 return SignatureUpdater.update( signature, new ClassNameUpdater( )
@@ -121,12 +243,30 @@ public class ClassIdentity
                         @Override
                         public String update( String className )
                         {
-                                // does the class have a package?
-                                if( className.indexOf( '/' ) >= 0 )
+                                // classes not in the none package can be passed through
+                                ClassEntry classEntry = new ClassEntry( className );
+                                if( !classEntry.getPackageName().equals( Constants.NonePackage ) )
                                 {
                                         return className;
                                 }
                                 
+                                // is this class ourself?
+                                if( className.equals( m_classEntry.getName() ) )
+                                {
+                                        return "CSelf";
+                                }
+                                
+                                // try the namer
+                                if( m_namer != null )
+                                {
+                                        String newName = m_namer.getName( className );
+                                        if( newName != null )
+                                        {
+                                                return newName;
+                                        }
+                                }
+                                
+                                // otherwise, use local naming
                                 if( !m_classNames.containsKey( className ) )
                                 {
                                         m_classNames.put( className, getNewClassName() );
@@ -141,6 +281,11 @@ public class ClassIdentity
                 } );
         }
         
+        private boolean isClassMatchedUniquely( String className )
+        {
+                return m_namer != null && m_namer.getName( Descriptor.toJvmName( className ) ) != null;
+        }
+        
         private String getBehaviorSignature( CtBehavior behavior )
         {
                 try
@@ -153,8 +298,7 @@ public class ClassIdentity
                         
                         // compute the hash from the opcodes
                         ConstPool constants = behavior.getMethodInfo().getConstPool();
-                        ConstPoolEditor editor = new ConstPoolEditor( constants );
-                        MessageDigest digest = MessageDigest.getInstance( "MD5" );
+                        final MessageDigest digest = MessageDigest.getInstance( "MD5" );
                         CodeIterator iter = behavior.getMethodInfo().getCodeAttribute().iterator();
                         while( iter.hasNext() )
                         {
@@ -164,46 +308,91 @@ public class ClassIdentity
                                 int opcode = iter.byteAt( pos );
                                 digest.update( (byte)opcode );
                                 
-                                // is there a constant value here?
-                                int constIndex = -1;
                                 switch( opcode )
                                 {
                                         case Opcode.LDC:
-                                                constIndex = iter.byteAt( pos + 1 );
+                                        {
+                                                int constIndex = iter.byteAt( pos + 1 );
+                                                updateHashWithConstant( digest, constants, constIndex );
+                                        }
                                         break;
                                         
                                         case Opcode.LDC_W:
-                                                constIndex = ( iter.byteAt( pos + 1 ) << 8 ) | iter.byteAt( pos + 2 );
-                                        break;
-                                        
                                         case Opcode.LDC2_W:
-                                                constIndex = ( iter.byteAt( pos + 1 ) << 8 ) | iter.byteAt( pos + 2 );
+                                        {
+                                                int constIndex = ( iter.byteAt( pos + 1 ) << 8 ) | iter.byteAt( pos + 2 );
+                                                updateHashWithConstant( digest, constants, constIndex );
+                                        }
                                         break;
                                 }
+                        }
+                        
+                        // update hash with method and field accesses
+                        behavior.instrument( new ExprEditor( )
+                        {
+                                @Override
+                                public void edit( MethodCall call )
+                                {
+                                        updateHashWithString( digest, scrubClassName( call.getClassName() ) );
+                                        updateHashWithString( digest, scrubSignature( call.getSignature() ) );
+                                        if( isClassMatchedUniquely( call.getClassName() ) )
+                                        {
+                                                updateHashWithString( digest, call.getMethodName() );
+                                        }
+                                }
                                 
-                                if( constIndex >= 0 )
+                                @Override
+                                public void edit( FieldAccess access )
                                 {
-                                        // update the hash with the constant value
-                                        ConstInfoAccessor item = editor.getItem( constIndex );
-                                        if( item.getType() == InfoType.StringInfo )
+                                        updateHashWithString( digest, scrubClassName( access.getClassName() ) );
+                                        updateHashWithString( digest, scrubSignature( access.getSignature() ) );
+                                        if( isClassMatchedUniquely( access.getClassName() ) )
                                         {
-                                                String val = constants.getStringInfo( constIndex );
-                                                try
-                                                {
-                                                        digest.update( val.getBytes( "UTF8" ) );
-                                                }
-                                                catch( UnsupportedEncodingException ex )
-                                                {
-                                                        throw new Error( ex );
-                                                }
+                                                updateHashWithString( digest, access.getFieldName() );
                                         }
                                 }
-                        }
+                                
+                                @Override
+                                public void edit( ConstructorCall call )
+                                {
+                                        updateHashWithString( digest, scrubClassName( call.getClassName() ) );
+                                        updateHashWithString( digest, scrubSignature( call.getSignature() ) );
+                                }
+                                
+                                @Override
+                                public void edit( NewExpr expr )
+                                {
+                                        updateHashWithString( digest, scrubClassName( expr.getClassName() ) );
+                                }
+                        } );
                         
                         // convert the hash to a hex string
                         return toHex( digest.digest() );
                 }
-                catch( BadBytecode | NoSuchAlgorithmException ex )
+                catch( BadBytecode | NoSuchAlgorithmException | CannotCompileException ex )
+                {
+                        throw new Error( ex );
+                }
+        }
+        
+        private void updateHashWithConstant( MessageDigest digest, ConstPool constants, int index )
+        {
+                ConstPoolEditor editor = new ConstPoolEditor( constants );
+                ConstInfoAccessor item = editor.getItem( index );
+                if( item.getType() == InfoType.StringInfo )
+                {
+                        updateHashWithString( digest, constants.getStringInfo( index ) );
+                }
+                // TODO: other constants
+        }
+        
+        private void updateHashWithString( MessageDigest digest, String val )
+        {
+                try
+                {
+                        digest.update( val.getBytes( "UTF8" ) );
+                }
+                catch( UnsupportedEncodingException ex )
                 {
                         throw new Error( ex );
                 }
@@ -239,7 +428,11 @@ public class ClassIdentity
                 return m_fields.equals( other.m_fields )
                         && m_methods.equals( other.m_methods )
                         && m_constructors.equals( other.m_constructors )
-                        && m_staticInitializer.equals( other.m_staticInitializer );
+                        && m_staticInitializer.equals( other.m_staticInitializer )
+                        && m_extends.equals( other.m_extends )
+                        && m_implements.equals( other.m_implements )
+                        && m_implementations.equals( other.m_implementations )
+                        && m_references.equals( other.m_references );
         }
         
         @Override
@@ -250,6 +443,10 @@ public class ClassIdentity
                 objs.addAll( m_methods );
                 objs.addAll( m_constructors );
                 objs.add( m_staticInitializer );
+                objs.add( m_extends );
+                objs.addAll( m_implements );
+                objs.addAll( m_implementations );
+                objs.addAll( m_references );
                 return Util.combineHashesOrdered( objs );
         }
 }
diff --git a/src/cuchaz/enigma/convert/ClassMapper.java b/src/cuchaz/enigma/convert/ClassMapper.java
index 5a16a1c..fe48c50 100644
--- a/src/cuchaz/enigma/convert/ClassMapper.java
+++ b/src/cuchaz/enigma/convert/ClassMapper.java
@@ -12,31 +12,20 @@ package cuchaz.enigma.convert;
 
 import java.io.File;
 import java.io.IOException;
-import java.util.Collection;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.jar.JarFile;
 
 import javassist.CtClass;
-
-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.analysis.JarClassIterator;
+import cuchaz.enigma.TranslatingTypeLoader;
+import cuchaz.enigma.analysis.JarIndex;
+import cuchaz.enigma.convert.ClassNamer.SidedClassNamer;
 import cuchaz.enigma.mapping.ClassEntry;
 
 public class ClassMapper
 {
-        private int m_numSourceClasses;
-        private int m_numDestClasses;
-        private Multimap<ClassIdentity,ClassIdentity> m_sourceClasses;
-        private Multimap<ClassIdentity,ClassIdentity> m_destClasses;
-        private List<ClassIdentity> m_unmatchedSourceClasses;
-        private List<ClassIdentity> m_unmatchedDestClasses;
-        private Map<ClassEntry,ClassIdentity> m_sourceKeyIndex;
-        
         public static void main( String[] args )
         throws IOException
         {
@@ -44,81 +33,146 @@ public class ClassMapper
                 JarFile fromJar = new JarFile( new File( "input/1.8-pre1.jar" ) );
                 JarFile toJar = new JarFile( new File( "input/1.8-pre2.jar" ) );
                 
-                ClassMapper mapper = new ClassMapper( fromJar, toJar );
-                System.out.println( String.format( "Mapped %d/%d source classes (%d unmatched) to %d/%d dest classes (%d unmatched)",
-                        mapper.m_sourceClasses.size(), mapper.m_numSourceClasses, mapper.m_unmatchedSourceClasses.size(),
-                        mapper.m_destClasses.size(), mapper.m_numDestClasses, mapper.m_unmatchedDestClasses.size()
-                ) );
+                // compute the matching
+                ClassMatching matching = ClassMapper.computeMatching( fromJar, toJar );
+                
+                // TODO: use the matching to convert the mappings
         }
         
-        public ClassMapper( JarFile sourceJar, JarFile destJar )
+        public static ClassMatching computeMatching( JarFile sourceJar, JarFile destJar )
         {
-                m_numSourceClasses = JarClassIterator.getClassEntries( sourceJar ).size();
-                m_numDestClasses = JarClassIterator.getClassEntries( destJar ).size();
+                // index jars
+                System.out.println( "Indexing source jar..." );
+                JarIndex sourceIndex = new JarIndex();
+                sourceIndex.indexJar( sourceJar );
+                System.out.println( "Indexing dest jar..." );
+                JarIndex destIndex = new JarIndex();
+                destIndex.indexJar( destJar );
                 
-                // compute identities for the source classes
-                m_sourceClasses = ArrayListMultimap.create();
-                m_sourceKeyIndex = Maps.newHashMap();
-                for( CtClass c : JarClassIterator.classes( sourceJar ) )
-                {
-                        ClassIdentity sourceClass = new ClassIdentity( c );
-                        m_sourceClasses.put( sourceClass, sourceClass );
-                        m_sourceKeyIndex.put( sourceClass.getClassEntry(), sourceClass );
-                }
+                System.out.println( "Computing matching..." );
                 
-                // match the dest classes to the source classes
-                m_destClasses = ArrayListMultimap.create();
-                m_unmatchedDestClasses = Lists.newArrayList();
-                for( CtClass c : JarClassIterator.classes( destJar ) )
+                TranslatingTypeLoader sourceLoader = new TranslatingTypeLoader( sourceJar, sourceIndex );
+                TranslatingTypeLoader destLoader = new TranslatingTypeLoader( destJar, destIndex );
+
+                ClassMatching matching = null;
+                for( boolean useReferences : Arrays.asList( false, true ) )
                 {
-                        ClassIdentity destClass = new ClassIdentity( c );
-                        Collection<ClassIdentity> matchedSourceClasses = m_sourceClasses.get( destClass );
-                        if( matchedSourceClasses.isEmpty() )
-                        {
-                                // unmatched dest class
-                                m_unmatchedDestClasses.add( destClass );
-                        }
-                        else
+                        int numMatches = 0;
+                        do
                         {
-                                ClassIdentity sourceClass = matchedSourceClasses.iterator().next();
-                                m_destClasses.put( sourceClass, destClass );
+                                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();
+                                }
+                                
+                                // 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() )
+                                        {
+                                                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 : matching.getUnmatchedDestClasses() )
+                                        {
+                                                destClassEntries.add( c.getClassEntry() );
+                                                matching.removeDest( c );
+                                        }
+                                }
+                                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 );
                 }
-
-                // get unmatched source classes
-                m_unmatchedSourceClasses = Lists.newArrayList();
-                for( ClassIdentity sourceClass : m_sourceClasses.keySet() )
+                
+                /* DEBUG: show some ambiguous matches
+                List<Map.Entry<List<ClassIdentity>,List<ClassIdentity>>> ambiguousMatches = new ArrayList<Map.Entry<List<ClassIdentity>,List<ClassIdentity>>>( matching.getAmbiguousMatches().entrySet() );
+                Collections.sort( ambiguousMatches, new Comparator<Map.Entry<List<ClassIdentity>,List<ClassIdentity>>>( )
                 {
-                        Collection<ClassIdentity> matchedSourceClasses = m_sourceClasses.get( sourceClass );
-                        Collection<ClassIdentity> matchedDestClasses = m_destClasses.get( sourceClass );
-                        if( matchedDestClasses.isEmpty() )
+                        @Override
+                        public int compare( Map.Entry<List<ClassIdentity>,List<ClassIdentity>> a, Map.Entry<List<ClassIdentity>,List<ClassIdentity>> b )
                         {
-                                m_unmatchedSourceClasses.add( sourceClass );
+                                String aName = a.getKey().get( 0 ).getClassEntry().getName();
+                                String bName = b.getKey().get( 0 ).getClassEntry().getName();
+                                return aName.compareTo( bName );
                         }
-                        else if( matchedDestClasses.size() > 1 )
+                } );
+                for( Map.Entry<List<ClassIdentity>,List<ClassIdentity>> entry : ambiguousMatches )
+                {
+                        for( ClassIdentity c : entry.getKey() )
                         {
-                                // warn about identity collisions
-                                System.err.println( String.format( "WARNING: identity collision:\n\tSource: %s\n\t  Dest: %s",
-                                        getClassEntries( matchedSourceClasses ),
-                                        getClassEntries( matchedDestClasses )
-                                ) );
+                                System.out.print( c.getClassEntry().getName() + " " );
                         }
+                        System.out.println();
+                }
+                Map.Entry<List<ClassIdentity>,List<ClassIdentity>> entry = ambiguousMatches.get( 7 );
+                for( ClassIdentity c : entry.getKey() )
+                {
+                        System.out.println( c );
                 }
+                for( ClassIdentity c : entry.getKey() )
+                {
+                        System.out.println( decompile( sourceLoader, c.getClassEntry() ) );
+                }
+                */
+                
+                return matching;
         }
         
-        public Map.Entry<Collection<ClassEntry>,Collection<ClassEntry>> getMapping( ClassEntry sourceEntry )
+        /* DEBUG
+        private static String decompile( TranslatingTypeLoader loader, ClassEntry classEntry )
         {
-                // TODO
-                return null;
-        }
-
-        private Collection<ClassEntry> getClassEntries( Collection<ClassIdentity> classes )
-        {
-                List<ClassEntry> entries = Lists.newArrayList();
-                for( ClassIdentity c : classes )
-                {
-                        entries.add( c.getClassEntry() );
-                }
-                return entries;
+                PlainTextOutput output = new PlainTextOutput();
+                DecompilerSettings settings = DecompilerSettings.javaDefaults();
+                settings.setForceExplicitImports( true );
+                settings.setShowSyntheticMembers( true );
+                settings.setTypeLoader( loader );
+                Decompiler.decompile( classEntry.getName(), output, settings );
+                return output.toString();
         }
+        */
 }
diff --git a/src/cuchaz/enigma/convert/ClassMatching.java b/src/cuchaz/enigma/convert/ClassMatching.java
new file mode 100644
index 0000000..fea8438
--- /dev/null
+++ b/src/cuchaz/enigma/convert/ClassMatching.java
@@ -0,0 +1,184 @@
+/*******************************************************************************
+ * 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.convert;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+
+import com.beust.jcommander.internal.Lists;
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.BiMap;
+import com.google.common.collect.HashBiMap;
+import com.google.common.collect.Multimap;
+
+public class ClassMatching
+{
+        private Multimap<ClassIdentity,ClassIdentity> m_sourceClasses;
+        private Multimap<ClassIdentity,ClassIdentity> m_matchedDestClasses;
+        private List<ClassIdentity> m_unmatchedDestClasses;
+        
+        public ClassMatching( )
+        {
+                m_sourceClasses = ArrayListMultimap.create();
+                m_matchedDestClasses = ArrayListMultimap.create();
+                m_unmatchedDestClasses = Lists.newArrayList();
+        }
+
+        public void addSource( ClassIdentity c )
+        {
+                m_sourceClasses.put( c, c );
+        }
+        
+        public void matchDestClass( ClassIdentity destClass )
+        {
+                Collection<ClassIdentity> matchedSourceClasses = m_sourceClasses.get( destClass );
+                if( matchedSourceClasses.isEmpty() )
+                {
+                        // no match
+                        m_unmatchedDestClasses.add( destClass );
+                }
+                else
+                {
+                        // found a match
+                        m_matchedDestClasses.put( destClass, destClass );
+                        
+                        // DEBUG
+                        ClassIdentity sourceClass = matchedSourceClasses.iterator().next();
+                        assert( sourceClass.hashCode() == destClass.hashCode() );
+                        assert( sourceClass.equals( destClass ) );
+                }
+        }
+        
+        public void removeSource( ClassIdentity sourceClass )
+        {
+                m_sourceClasses.remove( sourceClass, sourceClass );
+        }
+        
+        public void removeDest( ClassIdentity destClass )
+        {
+                m_matchedDestClasses.remove( destClass, destClass );
+                m_unmatchedDestClasses.remove( destClass );
+        }
+        
+        public List<ClassIdentity> getSourceClasses( )
+        {
+                return new ArrayList<ClassIdentity>( m_sourceClasses.values() );
+        }
+        
+        public List<ClassIdentity> getDestClasses( )
+        {
+                List<ClassIdentity> classes = Lists.newArrayList();
+                classes.addAll( m_matchedDestClasses.values() );
+                classes.addAll( m_unmatchedDestClasses );
+                return classes;
+        }
+        
+        public BiMap<ClassIdentity,ClassIdentity> getUniqueMatches( )
+        {
+                BiMap<ClassIdentity,ClassIdentity> uniqueMatches = HashBiMap.create();
+                for( ClassIdentity sourceClass : m_sourceClasses.keySet() )
+                {
+                        Collection<ClassIdentity> matchedSourceClasses = m_sourceClasses.get( sourceClass );
+                        Collection<ClassIdentity> matchedDestClasses = m_matchedDestClasses.get( sourceClass );
+                        if( matchedSourceClasses.size() == 1 && matchedDestClasses.size() == 1 )
+                        {
+                                ClassIdentity matchedSourceClass = matchedSourceClasses.iterator().next();
+                                ClassIdentity matchedDestClass = matchedSourceClasses.iterator().next();
+                                uniqueMatches.put( matchedSourceClass, matchedDestClass );
+                        }
+                }
+                return uniqueMatches;
+        }
+        
+        public BiMap<List<ClassIdentity>,List<ClassIdentity>> getAmbiguousMatches( )
+        {
+                BiMap<List<ClassIdentity>,List<ClassIdentity>> ambiguousMatches = HashBiMap.create();
+                for( ClassIdentity sourceClass : m_sourceClasses.keySet() )
+                {
+                        Collection<ClassIdentity> matchedSourceClasses = m_sourceClasses.get( sourceClass );
+                        Collection<ClassIdentity> matchedDestClasses = m_matchedDestClasses.get( sourceClass );
+                        if( matchedSourceClasses.size() > 1 && matchedDestClasses.size() > 1 )
+                        {
+                                ambiguousMatches.put(
+                                        new ArrayList<ClassIdentity>( matchedSourceClasses ),
+                                        new ArrayList<ClassIdentity>( matchedDestClasses )
+                                );
+                        }
+                }
+                return ambiguousMatches;
+        }
+        
+        public int getNumAmbiguousSourceMatches( )
+        {
+                int num = 0;
+                for( Map.Entry<List<ClassIdentity>,List<ClassIdentity>> entry : getAmbiguousMatches().entrySet() )
+                {
+                        num += entry.getKey().size();
+                }
+                return num;
+        }
+        
+        public int getNumAmbiguousDestMatches( )
+        {
+                int num = 0;
+                for( Map.Entry<List<ClassIdentity>,List<ClassIdentity>> entry : getAmbiguousMatches().entrySet() )
+                {
+                        num += entry.getValue().size();
+                }
+                return num;
+        }
+        
+        public List<ClassIdentity> getUnmatchedSourceClasses( )
+        {
+                List<ClassIdentity> classes = Lists.newArrayList();
+                for( ClassIdentity sourceClass : getSourceClasses() )
+                {
+                        if( m_matchedDestClasses.get( sourceClass ).isEmpty() )
+                        {
+                                classes.add( sourceClass );
+                        }
+                }
+                return classes;
+        }
+        
+        public List<ClassIdentity> getUnmatchedDestClasses( )
+        {
+                return new ArrayList<ClassIdentity>( m_unmatchedDestClasses );
+        }
+        
+        @Override
+        public String toString( )
+        {
+                StringBuilder buf = new StringBuilder();
+                
+                buf.append( "Source classes: " );
+                buf.append( getSourceClasses().size() );
+                buf.append( "\n\tUnique: " );
+                buf.append( getUniqueMatches().size() );
+                buf.append( "\n\tAmbiguous: " );
+                buf.append( getNumAmbiguousSourceMatches() );
+                buf.append( "\n\tUnmatched: " );
+                buf.append( getUnmatchedSourceClasses().size() );
+                
+                buf.append( "\nDest classes: " );
+                buf.append( getDestClasses().size() );
+                buf.append( "\n\tUnique: " );
+                buf.append( getUniqueMatches().size() );
+                buf.append( "\n\tAmbiguous: " );
+                buf.append( getNumAmbiguousDestMatches() );
+                buf.append( "\n\tUnmatched: " );
+                buf.append( getUnmatchedDestClasses().size() );
+                
+                return buf.toString();
+        }
+}
diff --git a/src/cuchaz/enigma/convert/ClassNamer.java b/src/cuchaz/enigma/convert/ClassNamer.java
new file mode 100644
index 0000000..1cd9665
--- /dev/null
+++ b/src/cuchaz/enigma/convert/ClassNamer.java
@@ -0,0 +1,75 @@
+/*******************************************************************************
+ * 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.convert;
+
+import java.util.Map;
+
+import com.beust.jcommander.internal.Maps;
+import com.google.common.collect.BiMap;
+
+public class ClassNamer
+{
+        public interface SidedClassNamer
+        {
+                String getName( String name );
+        }
+        
+        private Map<String,String> m_sourceNames;
+        private Map<String,String> m_destNames;
+        
+        public ClassNamer( BiMap<ClassIdentity,ClassIdentity> mappings )
+        {
+                // convert the identity mappings to name maps
+                m_sourceNames = Maps.newHashMap();
+                m_destNames = Maps.newHashMap();
+                int i = 0;
+                for( Map.Entry<ClassIdentity,ClassIdentity> entry : mappings.entrySet() )
+                {
+                        String name = String.format( "M%04d", i++ );
+                        m_sourceNames.put( entry.getKey().getClassEntry().getName(), name );
+                        m_destNames.put( entry.getValue().getClassEntry().getName(), name );
+                }
+        }
+        
+        public String getSourceName( String name )
+        {
+                return m_sourceNames.get( name );
+        }
+        
+        public String getDestName( String name )
+        {
+                return m_destNames.get( name );
+        }
+        
+        public SidedClassNamer getSourceNamer( )
+        {
+                return new SidedClassNamer( )
+                {
+                        @Override
+                        public String getName( String name )
+                        {
+                                return getSourceName( name );
+                        }
+                };
+        }
+        
+        public SidedClassNamer getDestNamer( )
+        {
+                return new SidedClassNamer( )
+                {
+                        @Override
+                        public String getName( String name )
+                        {
+                                return getDestName( name );
+                        }
+                };
+        }
+}
diff --git a/src/cuchaz/enigma/mapping/Translator.java b/src/cuchaz/enigma/mapping/Translator.java
index a671c27..23bf095 100644
--- a/src/cuchaz/enigma/mapping/Translator.java
+++ b/src/cuchaz/enigma/mapping/Translator.java
@@ -14,6 +14,8 @@ import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 
+import com.beust.jcommander.internal.Maps;
+
 import cuchaz.enigma.analysis.Ancestries;
 import cuchaz.enigma.mapping.SignatureUpdater.ClassNameUpdater;
 
@@ -23,6 +25,13 @@ public class Translator
         public Map<String,ClassMapping> m_classes;
         private Ancestries m_ancestries;
         
+        public Translator( )
+        {
+                m_direction = null;
+                m_classes = Maps.newHashMap();
+                m_ancestries = new Ancestries();
+        }
+        
         protected Translator( TranslationDirection direction, Map<String,ClassMapping> classes, Ancestries ancestries )
         {
                 m_direction = direction;