4 files changed, 117 insertions, 25 deletions

diff --git a/src/cuchaz/enigma/Deobfuscator.java b/src/cuchaz/enigma/Deobfuscator.java
index 9c845325..03c35113 100644
--- a/src/cuchaz/enigma/Deobfuscator.java
+++ b/src/cuchaz/enigma/Deobfuscator.java
@@ -22,6 +22,7 @@ import java.util.jar.JarFile;
 
 import javassist.bytecode.Descriptor;
 
+import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
 import com.strobel.assembler.metadata.MetadataSystem;
@@ -44,6 +45,7 @@ import cuchaz.enigma.mapping.ClassMapping;
 import cuchaz.enigma.mapping.ConstructorEntry;
 import cuchaz.enigma.mapping.Entry;
 import cuchaz.enigma.mapping.FieldEntry;
+import cuchaz.enigma.mapping.FieldMapping;
 import cuchaz.enigma.mapping.Mappings;
 import cuchaz.enigma.mapping.MappingsRenamer;
 import cuchaz.enigma.mapping.MethodEntry;
@@ -113,41 +115,107 @@ public class Deobfuscator
                         val = new Mappings();
                 }
                 
-                // make sure all the mappings match the classes in the jar
+                // look for any classes that got moved to inner classes
+                Map<String,String> renames = Maps.newHashMap();
                 for( ClassMapping classMapping : val.classes() )
                 {
-                        ClassEntry classEntry = new ClassEntry( classMapping.getObfName() );
-                        if( !m_jarIndex.getObfClassEntries().contains( classEntry ) )
+                        String outerClassName = m_jarIndex.getOuterClass( classMapping.getObfName() );
+                        if( outerClassName != null )
                         {
-                                throw new Error( "Class " + classEntry + " not found in Jar!" );
+                                // build the composite class name
+                                String newName = outerClassName + "$" + new ClassEntry( classMapping.getObfName() ).getSimpleName();
+                                
+                                // add a rename
+                                renames.put( classMapping.getObfName(), newName );
+                                
+                                System.out.println( String.format( "Converted class mapping %s to %s", classMapping.getObfName(), newName ) );
                         }
-                        
-                        // and method implementations
-                        for( MethodMapping methodMapping : classMapping.methods() )
+                }
+                for( Map.Entry<String,String> entry : renames.entrySet() )
+                {
+                        val.renameObfClass( entry.getKey(), entry.getValue() );
+                }
+                
+                // drop mappings that don't match the jar
+                List<ClassEntry> unknownClasses = Lists.newArrayList();
+                for( ClassMapping classMapping : val.classes() )
+                {
+                        checkClassMapping( unknownClasses, classMapping );
+                }
+                if( !unknownClasses.isEmpty() )
+                {
+                        throw new Error( "Unable to find classes in jar: " + unknownClasses );
+                }
+                
+                m_mappings = val;
+                m_renamer = new MappingsRenamer( m_jarIndex, m_mappings );
+                m_translatorCache.clear();
+        }
+        
+        private void checkClassMapping( List<ClassEntry> unknownClasses, ClassMapping classMapping )
+        {
+                // check the class
+                ClassEntry classEntry = new ClassEntry( classMapping.getObfName() );
+                String outerClassName = m_jarIndex.getOuterClass( classMapping.getObfName() );
+                if( outerClassName != null )
+                {
+                        classEntry = new ClassEntry( outerClassName + "$" + classEntry.getSimpleName() );
+                }
+                if( !m_jarIndex.getObfClassEntries().contains( classEntry ) )
+                {
+                        unknownClasses.add( classEntry );
+                }
+                
+                // check the fields
+                for( FieldMapping fieldMapping : Lists.newArrayList( classMapping.fields() ) )
+                {
+                        FieldEntry fieldEntry = new FieldEntry( classEntry, fieldMapping.getObfName() );
+                        if( m_jarIndex.getAccess( fieldEntry ) == null )
+                        {
+                                System.err.println( "WARNING: unable to find field " + fieldEntry + ". dropping mapping." );
+                                classMapping.removeFieldMapping( fieldMapping );
+                        }
+                }
+                
+                // check methods
+                for( MethodMapping methodMapping : Lists.newArrayList( classMapping.methods() ) )
+                {
+                        if( methodMapping.getObfName().equals( "<clinit>" ) )
+                        {
+                                // skip static initializers
+                                continue;
+                        }
+                        else if( methodMapping.getObfName().equals( "<init>" ) )
                         {
-                                if( methodMapping.getObfName().startsWith( "<" ) )
+                                ConstructorEntry constructorEntry = new ConstructorEntry( classEntry, methodMapping.getObfSignature() );
+                                if( m_jarIndex.getAccess( constructorEntry ) == null )
                                 {
-                                        // skip constructors and static initializers
-                                        continue;
+                                        System.err.println( "WARNING: unable to find constructor " + constructorEntry + ". dropping mapping." );
+                                        classMapping.removeMethodMapping( methodMapping );
                                 }
-                                
+                        }
+                        else
+                        {
                                 MethodEntry methodEntry = new MethodEntry(
                                         classEntry,
                                         methodMapping.getObfName(),
                                         methodMapping.getObfSignature()
                                 );
-                                if( !m_jarIndex.isMethodImplemented( methodEntry ) )
+                                if( m_jarIndex.getAccess( methodEntry ) == null )
                                 {
-                                        throw new Error( "Method " + methodEntry + " not found in Jar!" );
+                                        System.err.println( "WARNING: unable to find method " + methodEntry + ". dropping mapping." );
+                                        classMapping.removeMethodMapping( methodMapping );
                                 }
                         }
                 }
                 
-                m_mappings = val;
-                m_renamer = new MappingsRenamer( m_jarIndex, m_mappings );
-                m_translatorCache.clear();
+                // check inner classes
+                for( ClassMapping innerClassMapping : classMapping.innerClasses() )
+                {
+                        checkClassMapping( unknownClasses, innerClassMapping );
+                }
         }
-        
+
         public Translator getTranslator( TranslationDirection direction )
         {
                 Translator translator = m_translatorCache.get( direction );
diff --git a/src/cuchaz/enigma/analysis/EntryRenamer.java b/src/cuchaz/enigma/analysis/EntryRenamer.java
index e9483caa..b82b2547 100644
--- a/src/cuchaz/enigma/analysis/EntryRenamer.java
+++ b/src/cuchaz/enigma/analysis/EntryRenamer.java
@@ -189,10 +189,6 @@ public class EntryRenamer
                         reference.context = renameClassesInThing( renames, reference.context );
                         return thing;
                 }
-                else
-                {
-                        throw new Error( "Not an entry: " + thing );
-                }
                 
                 return thing;
         }
diff --git a/src/cuchaz/enigma/analysis/JarIndex.java b/src/cuchaz/enigma/analysis/JarIndex.java
index 8ebce35e..b51428a2 100644
--- a/src/cuchaz/enigma/analysis/JarIndex.java
+++ b/src/cuchaz/enigma/analysis/JarIndex.java
@@ -95,7 +95,7 @@ public class JarIndex
                         m_obfClassEntries.add( classEntry );
                 }
                 
-                // step 2: index method/field access
+                // step 2: index field/method/constructor access
                 for( CtClass c : JarClassIterator.classes( jar ) )
                 {
                         ClassRenamer.moveAllClassesOutOfDefaultPackage( c, Constants.NonePackage );
@@ -105,10 +105,15 @@ public class JarIndex
                                 FieldEntry fieldEntry = new FieldEntry( classEntry, field.getName() );
                                 m_access.put( fieldEntry, Access.get( field ) );
                         }
-                        for( CtBehavior behavior : c.getDeclaredBehaviors() )
+                        for( CtMethod method : c.getDeclaredMethods() )
+                        {
+                                MethodEntry methodEntry = new MethodEntry( classEntry, method.getName(), method.getSignature() );
+                                m_access.put( methodEntry, Access.get( method ) );
+                        }
+                        for( CtConstructor constructor : c.getDeclaredConstructors() )
                         {
-                                MethodEntry methodEntry = new MethodEntry( classEntry, behavior.getName(), behavior.getSignature() );
-                                m_access.put( methodEntry, Access.get( behavior ) );
+                                ConstructorEntry constructorEntry = new ConstructorEntry( classEntry, constructor.getSignature() );
+                                m_access.put( constructorEntry, Access.get( constructor ) );
                         }
                 }
                 
@@ -190,6 +195,7 @@ public class JarIndex
                         EntryRenamer.renameClassesInMultimap( renames, m_behaviorReferences );
                         EntryRenamer.renameClassesInMultimap( renames, m_fieldReferences );
                         EntryRenamer.renameClassesInMap( renames, m_bridgeMethods );
+                        EntryRenamer.renameClassesInMap( renames, m_access );
                 }
                 
                 // step 6: update other indices with bridge method info
diff --git a/src/cuchaz/enigma/mapping/ClassMapping.java b/src/cuchaz/enigma/mapping/ClassMapping.java
index 200d9ca2..88006cff 100644
--- a/src/cuchaz/enigma/mapping/ClassMapping.java
+++ b/src/cuchaz/enigma/mapping/ClassMapping.java
@@ -155,6 +155,17 @@ public class ClassMapping implements Serializable, Comparable<ClassMapping>
                 assert( deobfWasAdded );
                 assert( m_fieldsByObf.size() == m_fieldsByDeobf.size() );
         }
+        
+        public void removeFieldMapping( FieldMapping fieldMapping )
+        {
+                boolean obfWasRemoved = m_fieldsByObf.remove( fieldMapping.getObfName() ) != null;
+                assert( obfWasRemoved );
+                if( fieldMapping.getDeobfName() != null )
+                {
+                        boolean deobfWasRemoved = m_fieldsByDeobf.remove( fieldMapping.getDeobfName() ) != null;
+                        assert( deobfWasRemoved );
+                }
+        }
 
         public String getObfFieldName( String deobfName )
         {
@@ -225,6 +236,17 @@ public class ClassMapping implements Serializable, Comparable<ClassMapping>
                 assert( m_methodsByObf.size() >= m_methodsByDeobf.size() );
         }
         
+        public void removeMethodMapping( MethodMapping methodMapping )
+        {
+                boolean obfWasRemoved = m_methodsByObf.remove( getMethodKey( methodMapping.getObfName(), methodMapping.getObfSignature() ) ) != null;
+                assert( obfWasRemoved );
+                if( methodMapping.getDeobfName() != null )
+                {
+                        boolean deobfWasRemoved = m_methodsByDeobf.remove( getMethodKey( methodMapping.getDeobfName(), methodMapping.getObfSignature() ) ) != null;
+                        assert( deobfWasRemoved );
+                }
+        }
+        
         public MethodMapping getMethodByObf( String obfName, String signature )
         {
                 return m_methodsByObf.get( getMethodKey( obfName, signature ) );