4 files changed, 155 insertions, 90 deletions

diff --git a/src/cuchaz/enigma/analysis/ClassInheritanceTreeNode.java b/src/cuchaz/enigma/analysis/ClassInheritanceTreeNode.java
index b132305..3eaa391 100644
--- a/src/cuchaz/enigma/analysis/ClassInheritanceTreeNode.java
+++ b/src/cuchaz/enigma/analysis/ClassInheritanceTreeNode.java
@@ -51,8 +51,8 @@ public class ClassInheritanceTreeNode extends DefaultMutableTreeNode {
         public void load(TranslationIndex ancestries, boolean recurse) {
                 // get all the child nodes
                 List<ClassInheritanceTreeNode> nodes = Lists.newArrayList();
-                for (String subclassName : ancestries.getSubclassNames(m_obfClassName)) {
-                        nodes.add(new ClassInheritanceTreeNode(m_deobfuscatingTranslator, subclassName));
+                for (ClassEntry subclassEntry : ancestries.getSubclass(new ClassEntry(m_obfClassName))) {
+                        nodes.add(new ClassInheritanceTreeNode(m_deobfuscatingTranslator, subclassEntry.getName()));
                 }
                 
                 // add them to this node
diff --git a/src/cuchaz/enigma/analysis/JarIndex.java b/src/cuchaz/enigma/analysis/JarIndex.java
index 4b03a33..c96d3bc 100644
--- a/src/cuchaz/enigma/analysis/JarIndex.java
+++ b/src/cuchaz/enigma/analysis/JarIndex.java
@@ -114,8 +114,8 @@ public class JarIndex {
                 // step 3: index extends, implements, fields, and methods
                 for (CtClass c : JarClassIterator.classes(jar)) {
                         ClassRenamer.moveAllClassesOutOfDefaultPackage(c, Constants.NonePackage);
+                        m_translationIndex.indexClass(c);
                         String className = Descriptor.toJvmName(c.getName());
-                        m_translationIndex.addSuperclass(className, Descriptor.toJvmName(c.getClassFile().getSuperclass()));
                         for (String interfaceName : c.getClassFile().getInterfaces()) {
                                 className = Descriptor.toJvmName(className);
                                 interfaceName = Descriptor.toJvmName(interfaceName);
@@ -191,8 +191,6 @@ public class JarIndex {
                 String className = Descriptor.toJvmName(field.getDeclaringClass().getName());
                 FieldEntry fieldEntry = new FieldEntry(new ClassEntry(className), field.getName());
                 
-                m_translationIndex.addField(className, field.getName());
-                
                 // is the field a class type?
                 if (field.getSignature().startsWith("L")) {
                         ClassEntry fieldTypeEntry = new ClassEntry(field.getSignature().substring(1, field.getSignature().length() - 1));
@@ -230,13 +228,12 @@ public class JarIndex {
                         behavior.instrument(new ExprEditor() {
                                 @Override
                                 public void edit(MethodCall call) {
-                                        String className = Descriptor.toJvmName(call.getClassName());
                                         MethodEntry calledMethodEntry = new MethodEntry(
-                                                new ClassEntry(className),
+                                                new ClassEntry(Descriptor.toJvmName(call.getClassName())),
                                                 call.getMethodName(),
                                                 call.getSignature()
                                         );
-                                        ClassEntry resolvedClassEntry = resolveEntryClass(calledMethodEntry);
+                                        ClassEntry resolvedClassEntry = m_translationIndex.resolveEntryClass(calledMethodEntry);
                                         if (resolvedClassEntry != null && !resolvedClassEntry.equals(calledMethodEntry.getClassEntry())) {
                                                 calledMethodEntry = new MethodEntry(
                                                         resolvedClassEntry,
@@ -254,12 +251,11 @@ public class JarIndex {
                                 
                                 @Override
                                 public void edit(FieldAccess call) {
-                                        String className = Descriptor.toJvmName(call.getClassName());
                                         FieldEntry calledFieldEntry = new FieldEntry(
-                                                new ClassEntry(className),
+                                                new ClassEntry(Descriptor.toJvmName(call.getClassName())),
                                                 call.getFieldName()
                                         );
-                                        ClassEntry resolvedClassEntry = resolveEntryClass(calledFieldEntry);
+                                        ClassEntry resolvedClassEntry = m_translationIndex.resolveEntryClass(calledFieldEntry);
                                         if (resolvedClassEntry != null && !resolvedClassEntry.equals(calledFieldEntry.getClassEntry())) {
                                                 calledFieldEntry = new FieldEntry(resolvedClassEntry, call.getFieldName());
                                         }
@@ -273,9 +269,8 @@ public class JarIndex {
                                 
                                 @Override
                                 public void edit(ConstructorCall call) {
-                                        String className = Descriptor.toJvmName(call.getClassName());
                                         ConstructorEntry calledConstructorEntry = new ConstructorEntry(
-                                                new ClassEntry(className),
+                                                new ClassEntry(Descriptor.toJvmName(call.getClassName())),
                                                 call.getSignature()
                                         );
                                         EntryReference<BehaviorEntry,BehaviorEntry> reference = new EntryReference<BehaviorEntry,BehaviorEntry>(
@@ -288,9 +283,8 @@ public class JarIndex {
                                 
                                 @Override
                                 public void edit(NewExpr call) {
-                                        String className = Descriptor.toJvmName(call.getClassName());
                                         ConstructorEntry calledConstructorEntry = new ConstructorEntry(
-                                                new ClassEntry(className),
+                                                new ClassEntry(Descriptor.toJvmName(call.getClassName())),
                                                 call.getSignature()
                                         );
                                         EntryReference<BehaviorEntry,BehaviorEntry> reference = new EntryReference<BehaviorEntry,BehaviorEntry>(
@@ -306,25 +300,6 @@ public class JarIndex {
                 }
         }
         
-        public ClassEntry resolveEntryClass(Entry obfEntry) {
-                
-                // 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 (!containsObfEntry(obfEntry)) {
-                        // is there a parent class?
-                        String superclassName = m_translationIndex.getSuperclassName(obfEntry.getClassName());
-                        if (superclassName == null) {
-                                // this is probably a method from a class in a library
-                                // we can't trace the implementation up any higher unless we index the library
-                                return null;
-                        }
-                        
-                        // move up to the parent class
-                        obfEntry = obfEntry.cloneToNewClass(new ClassEntry(superclassName));
-                }
-                return obfEntry.getClassEntry();
-        }
-        
         private CtMethod getBridgedMethod(CtMethod method) {
                 
                 // bridge methods just call another method, cast it to the return type, and return the result
@@ -402,9 +377,9 @@ public class JarIndex {
                                 if (reference.entry instanceof ConstructorEntry && reference.context instanceof ConstructorEntry) {
                                         
                                         // is the entry a superclass of the context?
-                                        String calledClassName = reference.entry.getClassName();
-                                        String callerSuperclassName = m_translationIndex.getSuperclassName(reference.context.getClassName());
-                                        if (callerSuperclassName != null && callerSuperclassName.equals(calledClassName)) {
+                                        ClassEntry calledClassEntry = reference.entry.getClassEntry();
+                                        ClassEntry superclassEntry = m_translationIndex.getSuperclass(reference.context.getClassEntry());
+                                        if (superclassEntry != null && superclassEntry.equals(calledClassEntry)) {
                                                 // it's a super call, skip
                                                 continue;
                                         }
@@ -599,7 +574,9 @@ public class JarIndex {
                 // get the root node
                 List<String> ancestry = Lists.newArrayList();
                 ancestry.add(obfClassEntry.getName());
-                ancestry.addAll(m_translationIndex.getAncestry(obfClassEntry.getName()));
+                for (ClassEntry classEntry : m_translationIndex.getAncestry(obfClassEntry)) {
+                        ancestry.add(classEntry.getName());
+                }
                 ClassInheritanceTreeNode rootNode = new ClassInheritanceTreeNode(
                         deobfuscatingTranslator,
                         ancestry.get(ancestry.size() - 1)
@@ -625,21 +602,21 @@ public class JarIndex {
         public MethodInheritanceTreeNode getMethodInheritance(Translator deobfuscatingTranslator, MethodEntry obfMethodEntry) {
                 
                 // travel to the ancestor implementation
-                String baseImplementationClassName = obfMethodEntry.getClassName();
-                for (String ancestorClassName : m_translationIndex.getAncestry(obfMethodEntry.getClassName())) {
+                ClassEntry baseImplementationClassEntry = obfMethodEntry.getClassEntry();
+                for (ClassEntry ancestorClassEntry : m_translationIndex.getAncestry(obfMethodEntry.getClassEntry())) {
                         MethodEntry ancestorMethodEntry = new MethodEntry(
-                                new ClassEntry(ancestorClassName),
+                                new ClassEntry(ancestorClassEntry),
                                 obfMethodEntry.getName(),
                                 obfMethodEntry.getSignature()
                         );
                         if (containsObfBehavior(ancestorMethodEntry)) {
-                                baseImplementationClassName = ancestorClassName;
+                                baseImplementationClassEntry = ancestorClassEntry;
                         }
                 }
                 
                 // make a root node at the base
                 MethodEntry methodEntry = new MethodEntry(
-                        new ClassEntry(baseImplementationClassName),
+                        baseImplementationClassEntry,
                         obfMethodEntry.getName(),
                         obfMethodEntry.getSignature()
                 );
@@ -781,8 +758,8 @@ public class JarIndex {
         public Set<String> getInterfaces(String className) {
                 Set<String> interfaceNames = new HashSet<String>();
                 interfaceNames.addAll(m_interfaces.get(className));
-                for (String ancestor : m_translationIndex.getAncestry(className)) {
-                        interfaceNames.addAll(m_interfaces.get(ancestor));
+                for (ClassEntry ancestor : m_translationIndex.getAncestry(new ClassEntry(className))) {
+                        interfaceNames.addAll(m_interfaces.get(ancestor.getName()));
                 }
                 return interfaceNames;
         }
@@ -795,7 +772,7 @@ public class JarIndex {
                         String interfaceName = entry.getValue();
                         if (interfaceName.equals(targetInterfaceName)) {
                                 classNames.add(className);
-                                m_translationIndex.getSubclassNamesRecursively(classNames, className);
+                                m_translationIndex.getSubclassNamesRecursively(classNames, new ClassEntry(className));
                         }
                 }
                 return classNames;
diff --git a/src/cuchaz/enigma/analysis/MethodInheritanceTreeNode.java b/src/cuchaz/enigma/analysis/MethodInheritanceTreeNode.java
index eba8d87..8718220 100644
--- a/src/cuchaz/enigma/analysis/MethodInheritanceTreeNode.java
+++ b/src/cuchaz/enigma/analysis/MethodInheritanceTreeNode.java
@@ -71,9 +71,9 @@ public class MethodInheritanceTreeNode extends DefaultMutableTreeNode {
         public void load(JarIndex index, boolean recurse) {
                 // get all the child nodes
                 List<MethodInheritanceTreeNode> nodes = Lists.newArrayList();
-                for (String subclassName : index.getTranslationIndex().getSubclassNames(m_entry.getClassName())) {
+                for (ClassEntry subclassEntry : index.getTranslationIndex().getSubclass(m_entry.getClassEntry())) {
                         MethodEntry methodEntry = new MethodEntry(
-                                new ClassEntry(subclassName),
+                                subclassEntry,
                                 m_entry.getName(),
                                 m_entry.getSignature()
                         );
diff --git a/src/cuchaz/enigma/analysis/TranslationIndex.java b/src/cuchaz/enigma/analysis/TranslationIndex.java
index c14fd59..4a356eb 100644
--- a/src/cuchaz/enigma/analysis/TranslationIndex.java
+++ b/src/cuchaz/enigma/analysis/TranslationIndex.java
@@ -11,97 +11,185 @@
 package cuchaz.enigma.analysis;
 
 import java.io.Serializable;
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
-import javassist.bytecode.Descriptor;
+import javassist.CtBehavior;
+import javassist.CtClass;
+import javassist.CtField;
 
 import com.google.common.collect.HashMultimap;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Multimap;
 
+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.JavassistUtil;
+import cuchaz.enigma.mapping.Translator;
+
 public class TranslationIndex implements Serializable {
         
         private static final long serialVersionUID = 738687982126844179L;
         
-        private Map<String,String> m_superclasses;
-        private Multimap<String,String> m_fields;
+        private Map<ClassEntry,ClassEntry> m_superclasses;
+        private Multimap<ClassEntry,FieldEntry> m_fieldEntries;
+        private Multimap<ClassEntry,BehaviorEntry> m_behaviorEntries;
         
         public TranslationIndex() {
                 m_superclasses = Maps.newHashMap();
-                m_fields = HashMultimap.create();
-        }
-        
-        public TranslationIndex(TranslationIndex other) {
-                m_superclasses = Maps.newHashMap(other.m_superclasses);
-                m_fields = HashMultimap.create(other.m_fields);
+                m_fieldEntries = HashMultimap.create();
+                m_behaviorEntries = HashMultimap.create();
         }
         
-        public void addSuperclass(String className, String superclassName) {
-                className = Descriptor.toJvmName(className);
-                superclassName = Descriptor.toJvmName(superclassName);
+        public TranslationIndex(TranslationIndex other, Translator translator) {
                 
-                if (className.equals(superclassName)) {
-                        throw new IllegalArgumentException("Class cannot be its own superclass! " + className);
+                // translate the superclasses
+                m_superclasses = Maps.newHashMap();
+                for (Map.Entry<ClassEntry,ClassEntry> mapEntry : other.m_superclasses.entrySet()) {
+                        m_superclasses.put(
+                                translator.translateEntry(mapEntry.getKey()),
+                                translator.translateEntry(mapEntry.getValue())
+                        );
                 }
                 
-                if (!isJre(className) && !isJre(superclassName)) {
-                        m_superclasses.put(className, superclassName);
+                // translate the fields
+                m_fieldEntries = HashMultimap.create();
+                for (Map.Entry<ClassEntry,FieldEntry> mapEntry : other.m_fieldEntries.entries()) {
+                        m_fieldEntries.put(
+                                translator.translateEntry(mapEntry.getKey()),
+                                translator.translateEntry(mapEntry.getValue())
+                        );
+                }
+                
+                m_behaviorEntries = HashMultimap.create();
+                for (Map.Entry<ClassEntry,BehaviorEntry> mapEntry : other.m_behaviorEntries.entries()) {
+                        m_behaviorEntries.put(
+                                translator.translateEntry(mapEntry.getKey()),
+                                translator.translateEntry(mapEntry.getValue())
+                        );
                 }
         }
         
-        public void addField(String className, String fieldName) {
-                m_fields.put(className, fieldName);
+        public void indexClass(CtClass c) {
+                
+                ClassEntry classEntry = JavassistUtil.getClassEntry(c);
+                
+                // add the superclass
+                ClassEntry superclassEntry = JavassistUtil.getSuperclassEntry(c);
+                if (!isJre(classEntry) && !isJre(superclassEntry)) {
+                        m_superclasses.put(classEntry, superclassEntry);
+                }
+                
+                // add fields
+                for (CtField field : c.getDeclaredFields()) {
+                        FieldEntry fieldEntry = JavassistUtil.getFieldEntry(field);
+                        m_fieldEntries.put(fieldEntry.getClassEntry(), fieldEntry);
+                }
+                
+                // add behaviors
+                for (CtBehavior behavior : c.getDeclaredBehaviors()) {
+                        BehaviorEntry behaviorEntry = JavassistUtil.getBehaviorEntry(behavior);
+                        m_behaviorEntries.put(behaviorEntry.getClassEntry(), behaviorEntry);
+                }
         }
         
         public void renameClasses(Map<String,String> renames) {
                 EntryRenamer.renameClassesInMap(renames, m_superclasses);
-                EntryRenamer.renameClassesInMultimap(renames, m_fields);
+                EntryRenamer.renameClassesInMultimap(renames, m_fieldEntries);
+                EntryRenamer.renameClassesInMultimap(renames, m_behaviorEntries);
         }
         
-        public String getSuperclassName(String className) {
-                return m_superclasses.get(className);
+        public ClassEntry getSuperclass(ClassEntry classEntry) {
+                return m_superclasses.get(classEntry);
         }
         
-        public List<String> getAncestry(String className) {
-                List<String> ancestors = new ArrayList<String>();
-                while (className != null) {
-                        className = getSuperclassName(className);
-                        if (className != null) {
-                                ancestors.add(className);
+        public List<ClassEntry> getAncestry(ClassEntry classEntry) {
+                List<ClassEntry> ancestors = Lists.newArrayList();
+                while (classEntry != null) {
+                        classEntry = getSuperclass(classEntry);
+                        if (classEntry != null) {
+                                ancestors.add(classEntry);
                         }
                 }
                 return ancestors;
         }
         
-        public List<String> getSubclassNames(String className) {
+        public List<ClassEntry> getSubclass(ClassEntry classEntry) {
                 // linear search is fast enough for now
-                List<String> subclasses = Lists.newArrayList();
-                for (Map.Entry<String,String> entry : m_superclasses.entrySet()) {
-                        String subclass = entry.getKey();
-                        String superclass = entry.getValue();
-                        if (className.equals(superclass)) {
+                List<ClassEntry> subclasses = Lists.newArrayList();
+                for (Map.Entry<ClassEntry,ClassEntry> entry : m_superclasses.entrySet()) {
+                        ClassEntry subclass = entry.getKey();
+                        ClassEntry superclass = entry.getValue();
+                        if (classEntry.equals(superclass)) {
                                 subclasses.add(subclass);
                         }
                 }
                 return subclasses;
         }
         
-        public void getSubclassNamesRecursively(Set<String> out, String className) {
-                for (String subclassName : getSubclassNames(className)) {
-                        out.add(subclassName);
-                        getSubclassNamesRecursively(out, subclassName);
+        public void getSubclassesRecursively(Set<ClassEntry> out, ClassEntry classEntry) {
+                for (ClassEntry subclassEntry : getSubclass(classEntry)) {
+                        out.add(subclassEntry);
+                        getSubclassesRecursively(out, subclassEntry);
+                }
+        }
+        
+        public void getSubclassNamesRecursively(Set<String> out, ClassEntry classEntry) {
+                for (ClassEntry subclassEntry : getSubclass(classEntry)) {
+                        out.add(subclassEntry.getName());
+                        getSubclassNamesRecursively(out, subclassEntry);
                 }
         }
         
-        public boolean containsField(String className, String fieldName) {
-                return m_fields.containsEntry(className, fieldName);
+        public boolean entryExists(Entry entry) {
+                if (entry instanceof FieldEntry) {
+                        return fieldExists((FieldEntry)entry);
+                } else if (entry instanceof BehaviorEntry) {
+                        return behaviorExists((BehaviorEntry)entry);
+                } else if (entry instanceof ArgumentEntry) {
+                        return behaviorExists(((ArgumentEntry)entry).getBehaviorEntry());
+                }
+                throw new IllegalArgumentException("Cannot check existence for " + entry.getClass());
+        }
+        
+        public boolean fieldExists(FieldEntry fieldEntry) {
+                return m_fieldEntries.containsEntry(fieldEntry.getClassEntry(), fieldEntry);
+        }
+        
+        public boolean behaviorExists(BehaviorEntry behaviorEntry) {
+                return m_behaviorEntries.containsEntry(behaviorEntry.getClassEntry(), behaviorEntry);
+        }
+        
+        public ClassEntry resolveEntryClass(Entry entry) {
+                
+                if (entry instanceof ClassEntry) {
+                        return (ClassEntry)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)) {
+                        
+                        // is there a parent class?
+                        ClassEntry superclassEntry = getSuperclass(entry.getClassEntry());
+                        if (superclassEntry == null) {
+                                // this is probably a method from a class in a library
+                                // we can't trace the implementation up any higher unless we index the library
+                                return null;
+                        }
+                        
+                        // move up to the parent class
+                        entry = entry.cloneToNewClass(superclassEntry);
+                }
+                return entry.getClassEntry();
         }
         
-        private boolean isJre(String className) {
-                return className.startsWith("java/") || className.startsWith("javax/");
+        private boolean isJre(ClassEntry classEntry) {
+                return classEntry.getPackageName().startsWith("java") || classEntry.getPackageName().startsWith("javax");
         }
 }