Fix method reference and bridge detection

7 files changed, 90 insertions, 49 deletions

diff --git a/src/main/java/cuchaz/enigma/analysis/EntryRenamer.java b/src/main/java/cuchaz/enigma/analysis/EntryRenamer.java
index e37c1d0c..9be8378e 100644
--- a/src/main/java/cuchaz/enigma/analysis/EntryRenamer.java
+++ b/src/main/java/cuchaz/enigma/analysis/EntryRenamer.java
@@ -138,6 +138,13 @@ public class EntryRenamer {
                                         renameClassesInThing(renames, methodEntry.getSignature()),
                                         methodEntry.getAccess()
                         );
+                } else if (thing instanceof MethodEntry) {
+                        MethodEntry methodEntry = (MethodEntry) thing;
+                        return (T) new MethodEntry(
+                                        renameClassesInThing(renames, methodEntry.getOwnerClassEntry()),
+                                        methodEntry.getName(),
+                                        renameClassesInThing(renames, methodEntry.getDesc())
+                        );
                 } else if (thing instanceof LocalVariableEntry) {
                         LocalVariableEntry argumentEntry = (LocalVariableEntry) thing;
                         return (T) new LocalVariableEntry(renameClassesInThing(renames, argumentEntry.getOwnerEntry()), argumentEntry.getIndex(), argumentEntry.getName());
diff --git a/src/main/java/cuchaz/enigma/analysis/FieldReferenceTreeNode.java b/src/main/java/cuchaz/enigma/analysis/FieldReferenceTreeNode.java
index 3d0e48b4..f63b779a 100644
--- a/src/main/java/cuchaz/enigma/analysis/FieldReferenceTreeNode.java
+++ b/src/main/java/cuchaz/enigma/analysis/FieldReferenceTreeNode.java
@@ -63,7 +63,7 @@ public class FieldReferenceTreeNode extends DefaultMutableTreeNode implements Re
                                 add(new FieldReferenceTreeNode(this.deobfuscatingTranslator, reference, index.getAccess(this.entry)));
                         }
                 } else {
-                        for (EntryReference<MethodEntry, MethodDefEntry> reference : index.getMethodReferences(this.reference.context)) {
+                        for (EntryReference<MethodEntry, MethodDefEntry> reference : index.getMethodsReferencing(this.reference.context)) {
                                 add(new MethodReferenceTreeNode(this.deobfuscatingTranslator, reference, index.getAccess(this.reference.context)));
                         }
                 }
diff --git a/src/main/java/cuchaz/enigma/analysis/JarIndex.java b/src/main/java/cuchaz/enigma/analysis/JarIndex.java
index 8172deaa..27a8e07b 100644
--- a/src/main/java/cuchaz/enigma/analysis/JarIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/JarIndex.java
@@ -29,11 +29,13 @@ public class JarIndex {
         private Multimap<ClassEntry, FieldDefEntry> fields;
         private Multimap<ClassEntry, MethodDefEntry> methods;
         private Multimap<String, MethodDefEntry> methodImplementations;
-        private Multimap<MethodEntry, EntryReference<MethodEntry, MethodDefEntry>> methodReferences;
+        private Multimap<MethodEntry, EntryReference<MethodEntry, MethodDefEntry>> methodsReferencing;
+        private Multimap<MethodEntry, MethodEntry> methodReferences;
         private Multimap<FieldEntry, EntryReference<FieldEntry, MethodDefEntry>> fieldReferences;
         private Multimap<ClassEntry, ClassEntry> innerClassesByOuter;
         private Map<ClassEntry, ClassEntry> outerClassesByInner;
         private Map<MethodEntry, MethodEntry> bridgedMethods;
+        private Map<MethodEntry, MethodEntry> accessMethods;
         private Set<MethodEntry> syntheticMethods;
 
         public JarIndex(ReferencedEntryPool entryPool) {
@@ -44,11 +46,13 @@ public class JarIndex {
                 this.fields = HashMultimap.create();
                 this.methods = HashMultimap.create();
                 this.methodImplementations = HashMultimap.create();
+                this.methodsReferencing = HashMultimap.create();
                 this.methodReferences = HashMultimap.create();
                 this.fieldReferences = HashMultimap.create();
                 this.innerClassesByOuter = HashMultimap.create();
                 this.outerClassesByInner = Maps.newHashMap();
                 this.bridgedMethods = Maps.newHashMap();
+                this.accessMethods = Maps.newHashMap();
                 this.syntheticMethods = Sets.newHashSet();
         }
 
@@ -63,12 +67,15 @@ public class JarIndex {
                 // step 3: index field, method, constructor references
                 jar.visit(node -> node.accept(new IndexReferenceVisitor(this, Opcodes.ASM5)));
 
-                // step 4: index bridged methods
+                // step 4: index access and bridged methods
                 for (MethodDefEntry methodEntry : methods.values()) {
-                        // look for bridge and bridged methods
-                        MethodEntry bridgedMethod = findBridgedMethod(methodEntry);
-                        if (bridgedMethod != null) {
-                                this.bridgedMethods.put(methodEntry, bridgedMethod);
+                        // look for access and bridged methods
+                        MethodEntry accessedMethod = findAccessMethod(methodEntry);
+                        if (accessedMethod != null) {
+                                this.accessMethods.put(methodEntry, accessedMethod);
+                                if (isBridgedMethod(accessedMethod, methodEntry)) {
+                                        this.bridgedMethods.put(methodEntry, accessedMethod);
+                                }
                         }
                 }
 
@@ -89,6 +96,7 @@ public class JarIndex {
                         EntryRenamer.renameClassesInSet(renames, this.obfClassEntries);
                         this.translationIndex.renameClasses(renames);
                         EntryRenamer.renameClassesInMultimap(renames, this.methodImplementations);
+                        EntryRenamer.renameClassesInMultimap(renames, this.methodsReferencing);
                         EntryRenamer.renameClassesInMultimap(renames, this.methodReferences);
                         EntryRenamer.renameClassesInMultimap(renames, this.fieldReferences);
                         EntryRenamer.renameClassesInMap(renames, this.access);
@@ -134,7 +142,8 @@ public class JarIndex {
                 if (resolvedClassEntry != null && !resolvedClassEntry.equals(referencedMethod.getOwnerClassEntry())) {
                         referencedMethod = referencedMethod.updateOwnership(resolvedClassEntry);
                 }
-                methodReferences.put(referencedMethod, new EntryReference<>(referencedMethod, referencedMethod.getName(), callerEntry));
+                methodsReferencing.put(referencedMethod, new EntryReference<>(referencedMethod, referencedMethod.getName(), callerEntry));
+                methodReferences.put(callerEntry, referencedMethod);
         }
 
         protected void indexFieldAccess(MethodDefEntry callerEntry, String owner, String name, String desc) {
@@ -151,26 +160,54 @@ public class JarIndex {
                 this.outerClassesByInner.putIfAbsent(innerEntry, outerEntry);
         }
 
-        private MethodEntry findBridgedMethod(MethodDefEntry method) {
+        private MethodEntry findAccessMethod(MethodDefEntry method) {
 
-                // bridge methods just call another method, cast it to the return desc, and return the result
-                // let's see if we can detect this scenario
+                // we want to find all compiler-added methods that directly call another with no processing
 
                 // skip non-synthetic methods
                 if (!method.getAccess().isSynthetic()) {
                         return null;
                 }
 
-                // get all the called methods
-                final Collection<EntryReference<MethodEntry, MethodDefEntry>> referencedMethods = methodReferences.get(method);
+                // get all the methods that we call
+                final Collection<MethodEntry> referencedMethods = methodReferences.get(method);
 
                 // is there just one?
                 if (referencedMethods.size() != 1) {
                         return null;
                 }
 
-                // we have a bridge method!
-                return referencedMethods.stream().map(ref -> ref.context).filter(Objects::nonNull).findFirst().orElse(null);
+                return referencedMethods.stream().findFirst().orElse(null);
+        }
+
+        private boolean isBridgedMethod(MethodEntry called, MethodEntry access) {
+                // Bridged methods will always have the same name as the method they are calling
+                // They will also have the same amount of parameters (though equal descriptors cannot be guaranteed)
+                if (!called.getName().equals(access.getName()) || called.getDesc().getArgumentDescs().size() != access.getDesc().getArgumentDescs().size()) {
+                        return false;
+                }
+
+                TypeDescriptor accessReturn = access.getDesc().getReturnDesc();
+                TypeDescriptor calledReturn = called.getDesc().getReturnDesc();
+                if (calledReturn.isVoid() || calledReturn.isPrimitive() || accessReturn.isVoid() || accessReturn.isPrimitive()) {
+                        return false;
+                }
+
+                // Bridged methods will never have the same type as what they are calling
+                if (accessReturn.equals(calledReturn)) {
+                        return false;
+                }
+
+                String accessType = accessReturn.toString();
+
+                // If we're casting down from generic type to type-erased Object we're a bridge method
+                if (accessType.equals("Ljava/lang/Object;")) {
+                        return true;
+                }
+
+                // Now we need to detect cases where we are being casted down to a higher type bound
+                List<ClassEntry> calledAncestry = translationIndex.getAncestry(calledReturn.getTypeEntry());
+                return calledAncestry.contains(accessReturn.getTypeEntry());
         }
 
         public Set<ClassEntry> getObfClassEntries() {
@@ -302,11 +339,11 @@ public class JarIndex {
                         methodEntries.add(methodEntry);
                 }
 
-                // look at bridged methods!
-                MethodEntry bridgedEntry = getBridgedMethod(methodEntry);
-                while (bridgedEntry != null) {
-                        methodEntries.addAll(getRelatedMethodImplementations(bridgedEntry));
-                        bridgedEntry = getBridgedMethod(bridgedEntry);
+                // look at access methods!
+                MethodEntry accessMethod = getAccessMethod(methodEntry);
+                while (accessMethod != null) {
+                        methodEntries.addAll(getRelatedMethodImplementations(accessMethod));
+                        accessMethod = getAccessMethod(accessMethod);
                 }
 
                 // look at interface methods too
@@ -327,11 +364,11 @@ public class JarIndex {
                         methodEntries.add(methodEntry);
                 }
 
-                // look at bridged methods!
-                MethodEntry bridgedEntry = getBridgedMethod(methodEntry);
-                while (bridgedEntry != null) {
-                        methodEntries.addAll(getRelatedMethodImplementations(bridgedEntry));
-                        bridgedEntry = getBridgedMethod(bridgedEntry);
+                // look at access methods!
+                MethodEntry accessMethod = getAccessMethod(methodEntry);
+                while (accessMethod != null) {
+                        methodEntries.addAll(getRelatedMethodImplementations(accessMethod));
+                        accessMethod = getAccessMethod(accessMethod);
                 }
 
                 // recurse
@@ -355,19 +392,12 @@ public class JarIndex {
                 return fieldEntries;
         }
 
-        public Collection<EntryReference<MethodEntry, MethodDefEntry>> getMethodReferences(MethodEntry methodEntry) {
-                return this.methodReferences.get(methodEntry);
+        public Collection<EntryReference<MethodEntry, MethodDefEntry>> getMethodsReferencing(MethodEntry methodEntry) {
+                return this.methodsReferencing.get(methodEntry);
         }
 
         public Collection<MethodEntry> getReferencedMethods(MethodDefEntry methodEntry) {
-                // linear search is fast enough for now
-                Set<MethodEntry> behaviorEntries = Sets.newHashSet();
-                for (EntryReference<MethodEntry, MethodDefEntry> reference : this.methodReferences.values()) {
-                        if (reference.context == methodEntry) {
-                                behaviorEntries.add(reference.entry);
-                        }
-                }
-                return behaviorEntries;
+                return this.methodReferences.get(methodEntry);
         }
 
         public Collection<ClassEntry> getInnerClasses(ClassEntry obfOuterClassEntry) {
@@ -461,6 +491,10 @@ public class JarIndex {
                 return this.bridgedMethods.get(bridgeMethodEntry);
         }
 
+        public MethodEntry getAccessMethod(MethodEntry bridgeMethodEntry) {
+                return this.accessMethods.get(bridgeMethodEntry);
+        }
+
         public List<ClassEntry> getObfClassChain(ClassEntry obfClassEntry) {
 
                 // build class chain in inner-to-outer order
diff --git a/src/main/java/cuchaz/enigma/analysis/MethodReferenceTreeNode.java b/src/main/java/cuchaz/enigma/analysis/MethodReferenceTreeNode.java
index 15ae5153..76c73c15 100644
--- a/src/main/java/cuchaz/enigma/analysis/MethodReferenceTreeNode.java
+++ b/src/main/java/cuchaz/enigma/analysis/MethodReferenceTreeNode.java
@@ -64,7 +64,7 @@ public class MethodReferenceTreeNode extends DefaultMutableTreeNode
 
         public void load(JarIndex index, boolean recurse) {
                 // get all the child nodes
-                for (EntryReference<MethodEntry, MethodDefEntry> reference : index.getMethodReferences(this.entry)) {
+                for (EntryReference<MethodEntry, MethodDefEntry> reference : index.getMethodsReferencing(this.entry)) {
                         add(new MethodReferenceTreeNode(this.deobfuscatingTranslator, reference, index.getAccess(this.entry)));
                 }
 
diff --git a/src/test/java/cuchaz/enigma/TestJarIndexConstructorReferences.java b/src/test/java/cuchaz/enigma/TestJarIndexConstructorReferences.java
index dd275b38..763639a4 100644
--- a/src/test/java/cuchaz/enigma/TestJarIndexConstructorReferences.java
+++ b/src/test/java/cuchaz/enigma/TestJarIndexConstructorReferences.java
@@ -55,7 +55,7 @@ public class TestJarIndexConstructorReferences {
         @SuppressWarnings("unchecked")
         public void baseDefault() {
                 MethodEntry source = newMethod(baseClass, "<init>", "()V");
-                Collection<EntryReference<MethodEntry, MethodDefEntry>> references = index.getMethodReferences(source);
+                Collection<EntryReference<MethodEntry, MethodDefEntry>> references = index.getMethodsReferencing(source);
                 assertThat(references, containsInAnyOrder(
                                 newBehaviorReferenceByMethod(source, callerClass.getName(), "a", "()V"),
                                 newBehaviorReferenceByMethod(source, subClass.getName(), "<init>", "()V"),
@@ -67,7 +67,7 @@ public class TestJarIndexConstructorReferences {
         @SuppressWarnings("unchecked")
         public void baseInt() {
                 MethodEntry source = newMethod(baseClass, "<init>", "(I)V");
-                assertThat(index.getMethodReferences(source), containsInAnyOrder(
+                assertThat(index.getMethodsReferencing(source), containsInAnyOrder(
                                 newBehaviorReferenceByMethod(source, callerClass.getName(), "b", "()V")
                 ));
         }
@@ -76,7 +76,7 @@ public class TestJarIndexConstructorReferences {
         @SuppressWarnings("unchecked")
         public void subDefault() {
                 MethodEntry source = newMethod(subClass, "<init>", "()V");
-                assertThat(index.getMethodReferences(source), containsInAnyOrder(
+                assertThat(index.getMethodsReferencing(source), containsInAnyOrder(
                                 newBehaviorReferenceByMethod(source, callerClass.getName(), "c", "()V"),
                                 newBehaviorReferenceByMethod(source, subClass.getName(), "<init>", "(I)V")
                 ));
@@ -86,7 +86,7 @@ public class TestJarIndexConstructorReferences {
         @SuppressWarnings("unchecked")
         public void subInt() {
                 MethodEntry source = newMethod(subClass, "<init>", "(I)V");
-                assertThat(index.getMethodReferences(source), containsInAnyOrder(
+                assertThat(index.getMethodsReferencing(source), containsInAnyOrder(
                                 newBehaviorReferenceByMethod(source, callerClass.getName(), "d", "()V"),
                                 newBehaviorReferenceByMethod(source, subClass.getName(), "<init>", "(II)V"),
                                 newBehaviorReferenceByMethod(source, subsubClass.getName(), "<init>", "(I)V")
@@ -97,7 +97,7 @@ public class TestJarIndexConstructorReferences {
         @SuppressWarnings("unchecked")
         public void subIntInt() {
                 MethodEntry source = newMethod(subClass, "<init>", "(II)V");
-                assertThat(index.getMethodReferences(source), containsInAnyOrder(
+                assertThat(index.getMethodsReferencing(source), containsInAnyOrder(
                                 newBehaviorReferenceByMethod(source, callerClass.getName(), "e", "()V")
                 ));
         }
@@ -105,14 +105,14 @@ public class TestJarIndexConstructorReferences {
         @Test
         public void subIntIntInt() {
                 MethodEntry source = newMethod(subClass, "<init>", "(III)V");
-                assertThat(index.getMethodReferences(source), is(empty()));
+                assertThat(index.getMethodsReferencing(source), is(empty()));
         }
 
         @Test
         @SuppressWarnings("unchecked")
         public void subsubInt() {
                 MethodEntry source = newMethod(subsubClass, "<init>", "(I)V");
-                assertThat(index.getMethodReferences(source), containsInAnyOrder(
+                assertThat(index.getMethodsReferencing(source), containsInAnyOrder(
                                 newBehaviorReferenceByMethod(source, callerClass.getName(), "f", "()V")
                 ));
         }
@@ -121,7 +121,7 @@ public class TestJarIndexConstructorReferences {
         @SuppressWarnings("unchecked")
         public void defaultConstructable() {
                 MethodEntry source = newMethod(defaultClass, "<init>", "()V");
-                assertThat(index.getMethodReferences(source), containsInAnyOrder(
+                assertThat(index.getMethodsReferencing(source), containsInAnyOrder(
                                 newBehaviorReferenceByMethod(source, callerClass.getName(), "g", "()V")
                 ));
         }
diff --git a/src/test/java/cuchaz/enigma/TestJarIndexInheritanceTree.java b/src/test/java/cuchaz/enigma/TestJarIndexInheritanceTree.java
index 5bef4e57..23df1a99 100644
--- a/src/test/java/cuchaz/enigma/TestJarIndexInheritanceTree.java
+++ b/src/test/java/cuchaz/enigma/TestJarIndexInheritanceTree.java
@@ -152,7 +152,7 @@ public class TestJarIndexInheritanceTree {
 
                 // baseClass constructor
                 source = newMethod(baseClass, "<init>", "(Ljava/lang/String;)V");
-                references = index.getMethodReferences(source);
+                references = index.getMethodsReferencing(source);
                 assertThat(references, containsInAnyOrder(
                                 newBehaviorReferenceByMethod(source, subClassA.getName(), "<init>", "(Ljava/lang/String;)V"),
                                 newBehaviorReferenceByMethod(source, subClassB.getName(), "<init>", "()V")
@@ -160,14 +160,14 @@ public class TestJarIndexInheritanceTree {
 
                 // subClassA constructor
                 source = newMethod(subClassA, "<init>", "(Ljava/lang/String;)V");
-                references = index.getMethodReferences(source);
+                references = index.getMethodsReferencing(source);
                 assertThat(references, containsInAnyOrder(
                                 newBehaviorReferenceByMethod(source, subClassAA.getName(), "<init>", "()V")
                 ));
 
                 // baseClass.getName()
                 source = newMethod(baseClass, "a", "()Ljava/lang/String;");
-                references = index.getMethodReferences(source);
+                references = index.getMethodsReferencing(source);
                 assertThat(references, containsInAnyOrder(
                                 newBehaviorReferenceByMethod(source, subClassAA.getName(), "a", "()Ljava/lang/String;"),
                                 newBehaviorReferenceByMethod(source, subClassB.getName(), "a", "()V")
@@ -175,7 +175,7 @@ public class TestJarIndexInheritanceTree {
 
                 // subclassAA.getName()
                 source = newMethod(subClassAA, "a", "()Ljava/lang/String;");
-                references = index.getMethodReferences(source);
+                references = index.getMethodsReferencing(source);
                 assertThat(references, containsInAnyOrder(
                                 newBehaviorReferenceByMethod(source, subClassAA.getName(), "a", "()V")
                 ));
diff --git a/src/test/java/cuchaz/enigma/TestJarIndexLoneClass.java b/src/test/java/cuchaz/enigma/TestJarIndexLoneClass.java
index b1c128c2..b4529ddc 100644
--- a/src/test/java/cuchaz/enigma/TestJarIndexLoneClass.java
+++ b/src/test/java/cuchaz/enigma/TestJarIndexLoneClass.java
@@ -110,7 +110,7 @@ public class TestJarIndexLoneClass {
 
         @Test
         public void behaviorReferences() {
-                assertThat(index.getMethodReferences(newMethod("a", "a", "()Ljava/lang/String;")), is(empty()));
+                assertThat(index.getMethodsReferencing(newMethod("a", "a", "()Ljava/lang/String;")), is(empty()));
         }
 
         @Test