1 files changed, 63 insertions, 177 deletions

diff --git a/src/cuchaz/enigma/analysis/JarIndex.java b/src/cuchaz/enigma/analysis/JarIndex.java
index e7c92be..a8ac001 100644
--- a/src/cuchaz/enigma/analysis/JarIndex.java
+++ b/src/cuchaz/enigma/analysis/JarIndex.java
@@ -11,9 +11,8 @@
 package cuchaz.enigma.analysis;
 
 import java.lang.reflect.Modifier;
-import java.util.AbstractMap;
 import java.util.Collection;
-import java.util.Iterator;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -43,7 +42,6 @@ import com.google.common.collect.Sets;
 
 import cuchaz.enigma.Constants;
 import cuchaz.enigma.bytecode.ClassRenamer;
-import cuchaz.enigma.mapping.ArgumentEntry;
 import cuchaz.enigma.mapping.BehaviorEntry;
 import cuchaz.enigma.mapping.ClassEntry;
 import cuchaz.enigma.mapping.ConstructorEntry;
@@ -55,7 +53,8 @@ import cuchaz.enigma.mapping.Translator;
 public class JarIndex
 {
         private Set<ClassEntry> m_obfClassEntries;
-        private Ancestries m_ancestries;
+        private TranslationIndex m_translationIndex;
+        private Multimap<String,String> m_interfaces;
         private Map<Entry,Access> m_access;
         private Multimap<String,MethodEntry> m_methodImplementations;
         private Multimap<BehaviorEntry,EntryReference<BehaviorEntry,BehaviorEntry>> m_behaviorReferences;
@@ -68,7 +67,8 @@ public class JarIndex
         public JarIndex( )
         {
                 m_obfClassEntries = Sets.newHashSet();
-                m_ancestries = new Ancestries();
+                m_translationIndex = new TranslationIndex();
+                m_interfaces = HashMultimap.create();
                 m_access = Maps.newHashMap();
                 m_methodImplementations = HashMultimap.create();
                 m_behaviorReferences = HashMultimap.create();
@@ -109,15 +109,25 @@ public class JarIndex
                         }
                 }
                 
-                // step 3: index the types, methods
+                // step 3: index extends, implements, fields, and methods
                 for( CtClass c : JarClassIterator.classes( jar ) )
                 {
                         ClassRenamer.moveAllClassesOutOfDefaultPackage( c, Constants.NonePackage );
                         String className = Descriptor.toJvmName( c.getName() );
-                        m_ancestries.addSuperclass( className, Descriptor.toJvmName( c.getClassFile().getSuperclass() ) );
+                        m_translationIndex.addSuperclass( className, Descriptor.toJvmName( c.getClassFile().getSuperclass() ) );
                         for( String interfaceName : c.getClassFile().getInterfaces() )
                         {
-                                m_ancestries.addInterface( className, Descriptor.toJvmName( interfaceName ) );
+                                className = Descriptor.toJvmName( className );
+                                interfaceName = Descriptor.toJvmName( interfaceName );
+                                if( className.equals( interfaceName ) )
+                                {
+                                        throw new IllegalArgumentException( "Class cannot be its own interface! " + className );
+                                }
+                                m_interfaces.put( className, interfaceName );
+                        }
+                        for( CtField field : c.getDeclaredFields() )
+                        {
+                                indexField( field );
                         }
                         for( CtBehavior behavior : c.getDeclaredBehaviors() )
                         {
@@ -159,11 +169,24 @@ public class JarIndex
                         {
                                 renames.put( entry.getKey(), entry.getValue() + "$" + new ClassEntry( entry.getKey() ).getSimpleName() );
                         }
-                        renameClasses( renames );
+                        EntryRenamer.renameClassesInSet( renames, m_obfClassEntries );
+                        m_translationIndex.renameClasses( renames );
+                        EntryRenamer.renameClassesInMultimap( renames, m_interfaces );
+                        EntryRenamer.renameClassesInMultimap( renames, m_methodImplementations );
+                        EntryRenamer.renameClassesInMultimap( renames, m_behaviorReferences );
+                        EntryRenamer.renameClassesInMultimap( renames, m_fieldReferences );
                 }
                 
-                // step 5: update other indices with bridge method info
-                renameMethods( m_bridgeMethods );
+                // step 6: update other indices with bridge method info
+                EntryRenamer.renameMethodsInMultimap( m_bridgeMethods, m_methodImplementations );
+                EntryRenamer.renameMethodsInMultimap( m_bridgeMethods, m_behaviorReferences );
+                EntryRenamer.renameMethodsInMultimap( m_bridgeMethods, m_fieldReferences );
+        }
+        
+        private void indexField( CtField field )
+        {
+                String className = Descriptor.toJvmName( field.getDeclaringClass().getName() );
+                m_translationIndex.addField( className, field.getName() );
         }
 
         private void indexBehavior( CtBehavior behavior )
@@ -528,9 +551,9 @@ public class JarIndex
                 return m_obfClassEntries;
         }
         
-        public Ancestries getAncestries( )
+        public TranslationIndex getTranslationIndex( )
         {
-                return m_ancestries;
+                return m_translationIndex;
         }
         
         public Access getAccess( Entry entry )
@@ -553,11 +576,11 @@ public class JarIndex
                 // get the root node
                 List<String> ancestry = Lists.newArrayList();
                 ancestry.add( obfClassEntry.getName() );
-                ancestry.addAll( m_ancestries.getAncestry( obfClassEntry.getName() ) );
+                ancestry.addAll( m_translationIndex.getAncestry( obfClassEntry.getName() ) );
                 ClassInheritanceTreeNode rootNode = new ClassInheritanceTreeNode( deobfuscatingTranslator, ancestry.get( ancestry.size() - 1 ) );
                 
                 // expand all children recursively
-                rootNode.load( m_ancestries, true );
+                rootNode.load( m_translationIndex, true );
                 
                 return rootNode;
         }
@@ -565,7 +588,7 @@ public class JarIndex
         public ClassImplementationsTreeNode getClassImplementations( Translator deobfuscatingTranslator, ClassEntry obfClassEntry )
         {
                 ClassImplementationsTreeNode node = new ClassImplementationsTreeNode( deobfuscatingTranslator, obfClassEntry );
-                node.load( m_ancestries );
+                node.load( this );
                 return node;
         }
         
@@ -573,7 +596,7 @@ public class JarIndex
         {
                 // travel to the ancestor implementation
                 String baseImplementationClassName = obfMethodEntry.getClassName();
-                for( String ancestorClassName : m_ancestries.getAncestry( obfMethodEntry.getClassName() ) )
+                for( String ancestorClassName : m_translationIndex.getAncestry( obfMethodEntry.getClassName() ) )
                 {
                         MethodEntry ancestorMethodEntry = new MethodEntry(
                                 new ClassEntry( ancestorClassName ),
@@ -609,7 +632,7 @@ public class JarIndex
                 MethodEntry interfaceMethodEntry;
                 
                 // is this method on an interface?
-                if( m_ancestries.isInterface( obfMethodEntry.getClassName() ) )
+                if( isInterface( obfMethodEntry.getClassName() ) )
                 {
                         interfaceMethodEntry = obfMethodEntry;
                 }
@@ -617,7 +640,7 @@ public class JarIndex
                 {
                         // get the interface class
                         List<MethodEntry> methodInterfaces = Lists.newArrayList();
-                        for( String interfaceName : m_ancestries.getInterfaces( obfMethodEntry.getClassName() ) )
+                        for( String interfaceName : getInterfaces( obfMethodEntry.getClassName() ) )
                         {
                                 // is this method defined in this interface?
                                 MethodEntry methodInterface = new MethodEntry(
@@ -717,181 +740,44 @@ public class JarIndex
                 return m_anonymousClasses.contains( obfInnerClassName );
         }
         
-        public MethodEntry getBridgeMethod( MethodEntry methodEntry )
-        {
-                return m_bridgeMethods.get( methodEntry );
-        }
-        
-        private void renameClasses( Map<String,String> renames )
+        public Set<String> getInterfaces( String className )
         {
-                // rename class entries
-                Set<ClassEntry> obfClassEntries = Sets.newHashSet();
-                for( ClassEntry classEntry : m_obfClassEntries )
+                Set<String> interfaceNames = new HashSet<String>();
+                interfaceNames.addAll( m_interfaces.get( className ) );
+                for( String ancestor : m_translationIndex.getAncestry( className ) )
                 {
-                        if( renames.containsKey( classEntry.getName() ) )
-                        {
-                                obfClassEntries.add( new ClassEntry( renames.get( classEntry.getName() ) ) );
-                        }
-                        else
-                        {
-                                obfClassEntries.add( classEntry );
-                        }
+                        interfaceNames.addAll( m_interfaces.get( ancestor ) );
                 }
-                m_obfClassEntries = obfClassEntries;
-                
-                // rename others
-                m_ancestries.renameClasses( renames );
-                renameClassesInMultimap( renames, m_methodImplementations );
-                renameClassesInMultimap( renames, m_behaviorReferences );
-                renameClassesInMultimap( renames, m_fieldReferences );
+                return interfaceNames;
         }
         
-        private void renameMethods( Map<MethodEntry,MethodEntry> renames )
+        public Set<String> getImplementingClasses( String targetInterfaceName )
         {
-                renameMethodsInMultimap( renames, m_methodImplementations );
-                renameMethodsInMultimap( renames, m_behaviorReferences );
-                renameMethodsInMultimap( renames, m_fieldReferences );
-        }
-        
-        private <Key,Val> void renameClassesInMultimap( Map<String,String> renames, Multimap<Key,Val> map )
-        {
-                // for each key/value pair...
-                Set<Map.Entry<Key,Val>> entriesToAdd = Sets.newHashSet();
-                for( Map.Entry<Key,Val> entry : map.entries() )
+                // linear search is fast enough for now
+                Set<String> classNames = Sets.newHashSet();
+                for( Map.Entry<String,String> entry : m_interfaces.entries() )
                 {
-                        entriesToAdd.add( new AbstractMap.SimpleEntry<Key,Val>(
-                                renameClassesInThing( renames, entry.getKey() ),
-                                renameClassesInThing( renames, entry.getValue() )
-                        ) );
-                }
-                map.clear();
-                for( Map.Entry<Key,Val> entry : entriesToAdd )
-                {
-                        map.put( entry.getKey(), entry.getValue() );
-                }
-        }
-        
-        @SuppressWarnings( "unchecked" )
-        private <T> T renameClassesInThing( Map<String,String> renames, T thing )
-        {
-                if( thing instanceof String )
-                {
-                        String stringEntry = (String)thing;
-                        if( renames.containsKey( stringEntry ) )
+                        String className = entry.getKey();
+                        String interfaceName = entry.getValue();
+                        if( interfaceName.equals( targetInterfaceName ) )
                         {
-                                return (T)renames.get( stringEntry );
+                                classNames.add( className );
+                                m_translationIndex.getSubclassNamesRecursively( classNames, className );
                         }
                 }
-                else if( thing instanceof ClassEntry )
-                {
-                        ClassEntry classEntry = (ClassEntry)thing;
-                        return (T)new ClassEntry( renameClassesInThing( renames, classEntry.getClassName() ) );
-                }
-                else if( thing instanceof FieldEntry )
-                {
-                        FieldEntry fieldEntry = (FieldEntry)thing;
-                        return (T)new FieldEntry(
-                                renameClassesInThing( renames, fieldEntry.getClassEntry() ),
-                                fieldEntry.getName()
-                        );
-                }
-                else if( thing instanceof ConstructorEntry )
-                {
-                        ConstructorEntry constructorEntry = (ConstructorEntry)thing;
-                        return (T)new ConstructorEntry(
-                                renameClassesInThing( renames, constructorEntry.getClassEntry() ),
-                                constructorEntry.getSignature()
-                        );
-                }
-                else if( thing instanceof MethodEntry )
-                {
-                        MethodEntry methodEntry = (MethodEntry)thing;
-                        return (T)new MethodEntry(
-                                renameClassesInThing( renames, methodEntry.getClassEntry() ),
-                                methodEntry.getName(),
-                                methodEntry.getSignature()
-                        );
-                }
-                else if( thing instanceof ArgumentEntry )
-                {
-                        ArgumentEntry argumentEntry = (ArgumentEntry)thing;
-                        return (T)new ArgumentEntry(
-                                renameClassesInThing( renames, argumentEntry.getMethodEntry() ),
-                                argumentEntry.getIndex(),
-                                argumentEntry.getName()
-                        );
-                }
-                else if( thing instanceof EntryReference )
-                {
-                        EntryReference<Entry,Entry> reference = (EntryReference<Entry,Entry>)thing;
-                        reference.entry = renameClassesInThing( renames, reference.entry );
-                        reference.context = renameClassesInThing( renames, reference.context );
-                        return thing;
-                }
-                else
-                {
-                        throw new Error( "Not an entry: " + thing );
-                }
-                
-                return thing;
+                return classNames;
         }
         
-        private <Key,Val> void renameMethodsInMultimap( Map<MethodEntry,MethodEntry> renames, Multimap<Key,Val> map )
+        public boolean isInterface( String className )
         {
-                // for each key/value pair...
-                Set<Map.Entry<Key,Val>> entriesToAdd = Sets.newHashSet();
-                Iterator<Map.Entry<Key,Val>> iter = map.entries().iterator();
-                while( iter.hasNext() )
-                {
-                        Map.Entry<Key,Val> entry = iter.next();
-                        iter.remove();
-                        entriesToAdd.add( new AbstractMap.SimpleEntry<Key,Val>(
-                                renameMethodsInThing( renames, entry.getKey() ),
-                                renameMethodsInThing( renames, entry.getValue() )
-                        ) );
-                }
-                for( Map.Entry<Key,Val> entry : entriesToAdd )
-                {
-                        map.put( entry.getKey(), entry.getValue() );
-                }
+                return m_interfaces.containsValue( className );
         }
         
-        @SuppressWarnings( "unchecked" )
-        private <T> T renameMethodsInThing( Map<MethodEntry,MethodEntry> renames, T thing )
+        public MethodEntry getBridgeMethod( MethodEntry methodEntry )
         {
-                if( thing instanceof MethodEntry )
-                {
-                        MethodEntry methodEntry = (MethodEntry)thing;
-                        MethodEntry newMethodEntry = renames.get( methodEntry );
-                        if( newMethodEntry != null )
-                        {
-                                return (T)new MethodEntry(
-                                        methodEntry.getClassEntry(),
-                                        newMethodEntry.getName(),
-                                        methodEntry.getSignature()
-                                );
-                        }
-                        return thing;
-                }
-                else if( thing instanceof ArgumentEntry )
-                {
-                        ArgumentEntry argumentEntry = (ArgumentEntry)thing;
-                        return (T)new ArgumentEntry(
-                                renameMethodsInThing( renames, argumentEntry.getMethodEntry() ),
-                                argumentEntry.getIndex(),
-                                argumentEntry.getName()
-                        );
-                }
-                else if( thing instanceof EntryReference )
-                {
-                        EntryReference<Entry,Entry> reference = (EntryReference<Entry,Entry>)thing;
-                        reference.entry = renameMethodsInThing( renames, reference.entry );
-                        reference.context = renameMethodsInThing( renames, reference.context );
-                        return thing;
-                }
-                return thing;
+                return m_bridgeMethods.get( methodEntry );
         }
-
+        
         public boolean containsObfClass( ClassEntry obfClassEntry )
         {
                 return m_obfClassEntries.contains( obfClassEntry );