resolve methods using interfaces as well as superclasses

2 files changed, 84 insertions, 21 deletions

diff --git a/src/cuchaz/enigma/analysis/JarIndex.java b/src/cuchaz/enigma/analysis/JarIndex.java
index e255468a..6b3cf677 100644
--- a/src/cuchaz/enigma/analysis/JarIndex.java
+++ b/src/cuchaz/enigma/analysis/JarIndex.java
@@ -59,7 +59,6 @@ public class JarIndex {
         
         private Set<ClassEntry> m_obfClassEntries;
         private TranslationIndex m_translationIndex;
-        private Multimap<String,String> m_interfaces;
         private Map<Entry,Access> m_access;
         private Multimap<ClassEntry,FieldEntry> m_fields;
         private Multimap<ClassEntry,BehaviorEntry> m_behaviors;
@@ -74,7 +73,6 @@ public class JarIndex {
         public JarIndex() {
                 m_obfClassEntries = Sets.newHashSet();
                 m_translationIndex = new TranslationIndex();
-                m_interfaces = HashMultimap.create();
                 m_access = Maps.newHashMap();
                 m_fields = HashMultimap.create();
                 m_behaviors = HashMultimap.create();
@@ -124,7 +122,6 @@ public class JarIndex {
                                 if (className.equals(interfaceName)) {
                                         throw new IllegalArgumentException("Class cannot be its own interface! " + className);
                                 }
-                                m_interfaces.put(className, interfaceName);
                         }
                         for (CtBehavior behavior : c.getDeclaredBehaviors()) {
                                 indexBehavior(behavior);
@@ -176,7 +173,6 @@ public class JarIndex {
                         }
                         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);
@@ -637,11 +633,11 @@ public class JarIndex {
                         interfaceMethodEntries.add(obfMethodEntry);
                 } else {
                         // get the interface class
-                        for (String interfaceName : getInterfaces(obfMethodEntry.getClassName())) {
+                        for (ClassEntry interfaceEntry : getInterfaces(obfMethodEntry.getClassName())) {
                                 
                                 // is this method defined in this interface?
                                 MethodEntry methodInterface = new MethodEntry(
-                                        new ClassEntry(interfaceName),
+                                        interfaceEntry,
                                         obfMethodEntry.getName(),
                                         obfMethodEntry.getSignature()
                                 );
@@ -745,31 +741,33 @@ public class JarIndex {
                 return m_anonymousClasses.get(obfInnerClassName);
         }
         
-        public Set<String> getInterfaces(String className) {
-                Set<String> interfaceNames = new HashSet<String>();
-                interfaceNames.addAll(m_interfaces.get(className));
-                for (ClassEntry ancestor : m_translationIndex.getAncestry(new ClassEntry(className))) {
-                        interfaceNames.addAll(m_interfaces.get(ancestor.getName()));
+        public Set<ClassEntry> getInterfaces(String className) {
+                ClassEntry classEntry = new ClassEntry(className);
+                Set<ClassEntry> interfaces = new HashSet<ClassEntry>();
+                interfaces.addAll(m_translationIndex.getInterfaces(classEntry));
+                for (ClassEntry ancestor : m_translationIndex.getAncestry(classEntry)) {
+                        interfaces.addAll(m_translationIndex.getInterfaces(ancestor));
                 }
-                return interfaceNames;
+                return interfaces;
         }
         
         public Set<String> getImplementingClasses(String targetInterfaceName) {
+                
                 // linear search is fast enough for now
                 Set<String> classNames = Sets.newHashSet();
-                for (Map.Entry<String,String> entry : m_interfaces.entries()) {
-                        String className = entry.getKey();
-                        String interfaceName = entry.getValue();
-                        if (interfaceName.equals(targetInterfaceName)) {
-                                classNames.add(className);
-                                m_translationIndex.getSubclassNamesRecursively(classNames, new ClassEntry(className));
+                for (Map.Entry<ClassEntry,ClassEntry> entry : m_translationIndex.getClassInterfaces()) {
+                        ClassEntry classEntry = entry.getKey();
+                        ClassEntry interfaceEntry = entry.getValue();
+                        if (interfaceEntry.getName().equals(targetInterfaceName)) {
+                                classNames.add(classEntry.getClassName());
+                                m_translationIndex.getSubclassNamesRecursively(classNames, classEntry);
                         }
                 }
                 return classNames;
         }
         
         public boolean isInterface(String className) {
-                return m_interfaces.containsValue(className);
+                return m_translationIndex.isInterface(new ClassEntry(className));
         }
         
         public boolean containsObfClass(ClassEntry obfClassEntry) {
diff --git a/src/cuchaz/enigma/analysis/TranslationIndex.java b/src/cuchaz/enigma/analysis/TranslationIndex.java
index bd77344c..e0e66bf7 100644
--- a/src/cuchaz/enigma/analysis/TranslationIndex.java
+++ b/src/cuchaz/enigma/analysis/TranslationIndex.java
@@ -16,6 +16,7 @@ import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
 import java.io.OutputStream;
 import java.io.Serializable;
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -26,6 +27,7 @@ import java.util.zip.GZIPOutputStream;
 import javassist.CtBehavior;
 import javassist.CtClass;
 import javassist.CtField;
+import javassist.bytecode.Descriptor;
 
 import com.google.common.collect.HashMultimap;
 import com.google.common.collect.Lists;
@@ -36,8 +38,8 @@ import cuchaz.enigma.mapping.ArgumentEntry;
 import cuchaz.enigma.mapping.BehaviorEntry;
 import cuchaz.enigma.mapping.ClassEntry;
 import cuchaz.enigma.mapping.Entry;
-import cuchaz.enigma.mapping.FieldEntry;
 import cuchaz.enigma.mapping.EntryFactory;
+import cuchaz.enigma.mapping.FieldEntry;
 import cuchaz.enigma.mapping.Translator;
 
 public class TranslationIndex implements Serializable {
@@ -47,11 +49,13 @@ public class TranslationIndex implements Serializable {
         private Map<ClassEntry,ClassEntry> m_superclasses;
         private Multimap<ClassEntry,FieldEntry> m_fieldEntries;
         private Multimap<ClassEntry,BehaviorEntry> m_behaviorEntries;
+        private Multimap<ClassEntry,ClassEntry> m_interfaces;
         
         public TranslationIndex() {
                 m_superclasses = Maps.newHashMap();
                 m_fieldEntries = HashMultimap.create();
                 m_behaviorEntries = HashMultimap.create();
+                m_interfaces = HashMultimap.create();
         }
         
         public TranslationIndex(TranslationIndex other, Translator translator) {
@@ -65,6 +69,15 @@ public class TranslationIndex implements Serializable {
                         );
                 }
                 
+                // translate the interfaces
+                m_interfaces = HashMultimap.create();
+                for (Map.Entry<ClassEntry,ClassEntry> mapEntry : other.m_interfaces.entries()) {
+                        m_interfaces.put(
+                                translator.translateEntry(mapEntry.getKey()),
+                                translator.translateEntry(mapEntry.getValue())
+                        );
+                }
+                
                 // translate the fields
                 m_fieldEntries = HashMultimap.create();
                 for (Map.Entry<ClassEntry,FieldEntry> mapEntry : other.m_fieldEntries.entries()) {
@@ -90,13 +103,24 @@ public class TranslationIndex implements Serializable {
         public void indexClass(CtClass c, boolean indexMembers) {
                 
                 ClassEntry classEntry = EntryFactory.getClassEntry(c);
+                if (isJre(classEntry)) {
+                        return;
+                }
                 
                 // add the superclass
                 ClassEntry superclassEntry = EntryFactory.getSuperclassEntry(c);
-                if (!isJre(classEntry) && superclassEntry != null && !isJre(superclassEntry)) {
+                if (superclassEntry != null && !isJre(superclassEntry)) {
                         m_superclasses.put(classEntry, superclassEntry);
                 }
                 
+                // add the interfaces
+                for (String interfaceClassName : c.getClassFile().getInterfaces()) {
+                        ClassEntry interfaceClassEntry = new ClassEntry(Descriptor.toJvmName(interfaceClassName));
+                        if (!isJre(interfaceClassEntry)) {
+                                m_interfaces.put(classEntry, interfaceClassEntry);
+                        }
+                }
+                
                 if (indexMembers) {
                         // add fields
                         for (CtField field : c.getDeclaredFields()) {
@@ -134,6 +158,7 @@ public class TranslationIndex implements Serializable {
         }
         
         public List<ClassEntry> getSubclass(ClassEntry classEntry) {
+                
                 // linear search is fast enough for now
                 List<ClassEntry> subclasses = Lists.newArrayList();
                 for (Map.Entry<ClassEntry,ClassEntry> entry : m_superclasses.entrySet()) {
@@ -160,6 +185,18 @@ public class TranslationIndex implements Serializable {
                 }
         }
         
+        public Collection<Map.Entry<ClassEntry,ClassEntry>> getClassInterfaces() {
+                return m_interfaces.entries();
+        }
+        
+        public Collection<ClassEntry> getInterfaces(ClassEntry classEntry) {
+                return m_interfaces.get(classEntry);
+        }
+        
+        public boolean isInterface(ClassEntry classEntry) {
+                return m_interfaces.containsValue(classEntry);
+        }
+        
         public boolean entryExists(Entry entry) {
                 if (entry instanceof FieldEntry) {
                         return fieldExists((FieldEntry)entry);
@@ -185,6 +222,21 @@ public class TranslationIndex implements Serializable {
                         return (ClassEntry)entry;
                 }
                 
+                ClassEntry superclassEntry = resolveSuperclass(entry);
+                if (superclassEntry != null) {
+                        return superclassEntry;
+                }
+                
+                ClassEntry interfaceEntry = resolveInterface(entry);
+                if (interfaceEntry != null) {
+                        return interfaceEntry;
+                }
+                
+                return null;
+        }
+        
+        public ClassEntry resolveSuperclass(Entry entry) {
+                
                 // this entry could refer to a method on a class where the method is not actually implemented
                 // travel up the inheritance tree to find the closest implementation
                 while (!entryExists(entry)) {
@@ -203,6 +255,19 @@ public class TranslationIndex implements Serializable {
                 return entry.getClassEntry();
         }
         
+        public ClassEntry resolveInterface(Entry entry) {
+                
+                // the interfaces for any class is a forest
+                // so let's look at all the trees
+                for (ClassEntry interfaceEntry : m_interfaces.get(entry.getClassEntry())) {
+                        ClassEntry resolvedClassEntry = resolveSuperclass(entry.cloneToNewClass(interfaceEntry));
+                        if (resolvedClassEntry != null) {
+                                return resolvedClassEntry;
+                        }
+                }
+                return null;
+        }
+        
         private boolean isJre(ClassEntry classEntry) {
                 String packageName = classEntry.getPackageName();
                 return packageName != null && (packageName.startsWith("java") || packageName.startsWith("javax"));