Method type reference corrections (#142)

* Add more specific returns for translatables

* Only index method descriptors for implemented methods and methods in generated lambda classes

7 files changed, 167 insertions, 33 deletions

diff --git a/src/main/java/cuchaz/enigma/analysis/index/IndexReferenceVisitor.java b/src/main/java/cuchaz/enigma/analysis/index/IndexReferenceVisitor.java
index ba5d3b6a..b730a8a1 100644
--- a/src/main/java/cuchaz/enigma/analysis/index/IndexReferenceVisitor.java
+++ b/src/main/java/cuchaz/enigma/analysis/index/IndexReferenceVisitor.java
@@ -1,16 +1,11 @@
 package cuchaz.enigma.analysis.index;
 
 import cuchaz.enigma.translation.representation.AccessFlags;
+import cuchaz.enigma.translation.representation.Lambda;
 import cuchaz.enigma.translation.representation.MethodDescriptor;
 import cuchaz.enigma.translation.representation.Signature;
-import cuchaz.enigma.translation.representation.entry.ClassEntry;
-import cuchaz.enigma.translation.representation.entry.FieldEntry;
-import cuchaz.enigma.translation.representation.entry.MethodDefEntry;
-import cuchaz.enigma.translation.representation.entry.MethodEntry;
-import org.objectweb.asm.ClassVisitor;
-import org.objectweb.asm.Handle;
-import org.objectweb.asm.MethodVisitor;
-import org.objectweb.asm.Opcodes;
+import cuchaz.enigma.translation.representation.entry.*;
+import org.objectweb.asm.*;
 
 public class IndexReferenceVisitor extends ClassVisitor {
         private final JarIndexer indexer;
@@ -54,29 +49,37 @@ public class IndexReferenceVisitor extends ClassVisitor {
                         this.indexer.indexMethodReference(callerEntry, methodEntry);
                 }
 
+                private static ParentedEntry<?> getHandleEntry(Handle handle) {
+                        switch (handle.getTag()) {
+                                case Opcodes.H_GETFIELD:
+                                case Opcodes.H_GETSTATIC:
+                                case Opcodes.H_PUTFIELD:
+                                case Opcodes.H_PUTSTATIC:
+                                        return FieldEntry.parse(handle.getOwner(), handle.getName(), handle.getDesc());
+                                case Opcodes.H_INVOKEINTERFACE:
+                                case Opcodes.H_INVOKESPECIAL:
+                                case Opcodes.H_INVOKESTATIC:
+                                case Opcodes.H_INVOKEVIRTUAL:
+                                case Opcodes.H_NEWINVOKESPECIAL:
+                                        return MethodEntry.parse(handle.getOwner(), handle.getName(), handle.getDesc());
+                        }
+                        throw new RuntimeException("Invalid handle tag " + handle.getTag());
+                }
+
                 @Override
                 public void visitInvokeDynamicInsn(String name, String desc, Handle bsm, Object... bsmArgs) {
-                        for (Object bsmArg : bsmArgs) {
-                                if (bsmArg instanceof Handle) {
-                                        Handle handle = (Handle) bsmArg;
-                                        switch (handle.getTag()) {
-                                                case Opcodes.H_GETFIELD:
-                                                case Opcodes.H_GETSTATIC:
-                                                case Opcodes.H_PUTFIELD:
-                                                case Opcodes.H_PUTSTATIC:
-                                                        FieldEntry fieldEntry = FieldEntry.parse(handle.getOwner(), handle.getName(), handle.getDesc());
-                                                        this.indexer.indexFieldReference(callerEntry, fieldEntry);
-                                                        break;
-                                                case Opcodes.H_INVOKEINTERFACE:
-                                                case Opcodes.H_INVOKESPECIAL:
-                                                case Opcodes.H_INVOKESTATIC:
-                                                case Opcodes.H_INVOKEVIRTUAL:
-                                                case Opcodes.H_NEWINVOKESPECIAL:
-                                                        MethodEntry methodEntry = MethodEntry.parse(handle.getOwner(), handle.getName(), handle.getDesc());
-                                                        this.indexer.indexMethodReference(callerEntry, methodEntry);
-                                                        break;
-                                        }
-                                }
+                        if ("java/lang/invoke/LambdaMetafactory".equals(bsm.getOwner()) && "metafactory".equals(bsm.getName())) {
+                                Type samMethodType = (Type) bsmArgs[0];
+                                Handle implMethod = (Handle) bsmArgs[1];
+                                Type instantiatedMethodType = (Type) bsmArgs[2];
+
+                                this.indexer.indexLambda(callerEntry, new Lambda(
+                                        name,
+                                        new MethodDescriptor(desc),
+                                        new MethodDescriptor(samMethodType.getDescriptor()),
+                                        getHandleEntry(implMethod),
+                                        new MethodDescriptor(instantiatedMethodType.getDescriptor())
+                                ));
                         }
                 }
         }
diff --git a/src/main/java/cuchaz/enigma/analysis/index/JarIndex.java b/src/main/java/cuchaz/enigma/analysis/index/JarIndex.java
index ac907af0..fd4e618b 100644
--- a/src/main/java/cuchaz/enigma/analysis/index/JarIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/index/JarIndex.java
@@ -16,6 +16,7 @@ import com.google.common.collect.Multimap;
 import cuchaz.enigma.analysis.ParsedJar;
 import cuchaz.enigma.translation.mapping.EntryResolver;
 import cuchaz.enigma.translation.mapping.IndexEntryResolver;
+import cuchaz.enigma.translation.representation.Lambda;
 import cuchaz.enigma.translation.representation.entry.*;
 import org.objectweb.asm.ClassReader;
 import org.objectweb.asm.Opcodes;
@@ -129,6 +130,15 @@ public class JarIndex implements JarIndexer {
                 indexers.forEach(indexer -> indexer.indexFieldReference(callerEntry, referencedEntry));
         }
 
+        @Override
+        public void indexLambda(MethodDefEntry callerEntry, Lambda lambda) {
+                if (callerEntry.getParent().isJre()) {
+                        return;
+                }
+
+                indexers.forEach(indexer -> indexer.indexLambda(callerEntry, lambda));
+        }
+
         public EntryIndex getEntryIndex() {
                 return entryIndex;
         }
diff --git a/src/main/java/cuchaz/enigma/analysis/index/JarIndexer.java b/src/main/java/cuchaz/enigma/analysis/index/JarIndexer.java
index 457c2237..4f885b39 100644
--- a/src/main/java/cuchaz/enigma/analysis/index/JarIndexer.java
+++ b/src/main/java/cuchaz/enigma/analysis/index/JarIndexer.java
@@ -1,5 +1,6 @@
 package cuchaz.enigma.analysis.index;
 
+import cuchaz.enigma.translation.representation.Lambda;
 import cuchaz.enigma.translation.representation.entry.*;
 
 public interface JarIndexer {
@@ -18,6 +19,9 @@ public interface JarIndexer {
         default void indexFieldReference(MethodDefEntry callerEntry, FieldEntry referencedEntry) {
         }
 
+        default void indexLambda(MethodDefEntry callerEntry, Lambda lambda) {
+        }
+
         default void processIndex(JarIndex index) {
         }
 }
diff --git a/src/main/java/cuchaz/enigma/analysis/index/ReferenceIndex.java b/src/main/java/cuchaz/enigma/analysis/index/ReferenceIndex.java
index 04306bd9..f54c90dd 100644
--- a/src/main/java/cuchaz/enigma/analysis/index/ReferenceIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/index/ReferenceIndex.java
@@ -4,6 +4,7 @@ import com.google.common.collect.HashMultimap;
 import com.google.common.collect.Multimap;
 import cuchaz.enigma.analysis.EntryReference;
 import cuchaz.enigma.translation.mapping.ResolutionStrategy;
+import cuchaz.enigma.translation.representation.Lambda;
 import cuchaz.enigma.translation.representation.MethodDescriptor;
 import cuchaz.enigma.translation.representation.TypeDescriptor;
 import cuchaz.enigma.translation.representation.entry.*;
@@ -64,8 +65,6 @@ public class ReferenceIndex implements JarIndexer {
                         ClassEntry referencedClass = referencedEntry.getParent();
                         referencesToClasses.put(referencedClass, new EntryReference<>(referencedClass, referencedEntry.getName(), callerEntry));
                 }
-
-                indexMethodDescriptor(callerEntry, referencedEntry.getDesc());
         }
 
         @Override
@@ -74,6 +73,19 @@ public class ReferenceIndex implements JarIndexer {
         }
 
         @Override
+        public void indexLambda(MethodDefEntry callerEntry, Lambda lambda) {
+                if (lambda.getImplMethod() instanceof MethodEntry) {
+                        indexMethodReference(callerEntry, (MethodEntry) lambda.getImplMethod());
+                } else {
+                        indexFieldReference(callerEntry, (FieldEntry) lambda.getImplMethod());
+                }
+
+                indexMethodDescriptor(callerEntry, lambda.getInvokedType());
+                indexMethodDescriptor(callerEntry, lambda.getSamMethodType());
+                indexMethodDescriptor(callerEntry, lambda.getInstantiatedMethodType());
+        }
+
+        @Override
         public void processIndex(JarIndex index) {
                 methodReferences = remapReferences(index, methodReferences);
                 referencesToMethods = remapReferencesTo(index, referencesToMethods);
diff --git a/src/main/java/cuchaz/enigma/translation/representation/Lambda.java b/src/main/java/cuchaz/enigma/translation/representation/Lambda.java
new file mode 100644
index 00000000..63eb5630
--- /dev/null
+++ b/src/main/java/cuchaz/enigma/translation/representation/Lambda.java
@@ -0,0 +1,105 @@
+package cuchaz.enigma.translation.representation;

+

+import cuchaz.enigma.translation.Translatable;

+import cuchaz.enigma.translation.Translator;

+import cuchaz.enigma.translation.mapping.EntryMap;

+import cuchaz.enigma.translation.mapping.EntryMapping;

+import cuchaz.enigma.translation.mapping.EntryResolver;

+import cuchaz.enigma.translation.mapping.ResolutionStrategy;

+import cuchaz.enigma.translation.representation.entry.ClassEntry;

+import cuchaz.enigma.translation.representation.entry.MethodEntry;

+import cuchaz.enigma.translation.representation.entry.ParentedEntry;

+

+import java.util.Objects;

+

+public class Lambda implements Translatable {

+        private final String invokedName;

+        private final MethodDescriptor invokedType;

+        private final MethodDescriptor samMethodType;

+        private final ParentedEntry<?> implMethod;

+        private final MethodDescriptor instantiatedMethodType;

+

+        public Lambda(String invokedName, MethodDescriptor invokedType, MethodDescriptor samMethodType, ParentedEntry<?> implMethod, MethodDescriptor instantiatedMethodType) {

+                this.invokedName = invokedName;

+                this.invokedType = invokedType;

+                this.samMethodType = samMethodType;

+                this.implMethod = implMethod;

+                this.instantiatedMethodType = instantiatedMethodType;

+        }

+

+        @Override

+        public Lambda translate(Translator translator, EntryResolver resolver, EntryMap<EntryMapping> mappings) {

+                MethodEntry samMethod = new MethodEntry(getInterface(), invokedName, samMethodType);

+                EntryMapping samMethodMapping = resolveMapping(resolver, mappings, samMethod);

+

+                return new Lambda(

+                        samMethodMapping != null ? samMethodMapping.getTargetName() : invokedName,

+                        invokedType.translate(translator, resolver, mappings),

+                        samMethodType.translate(translator, resolver, mappings),

+                        implMethod.translate(translator, resolver, mappings),

+                        instantiatedMethodType.translate(translator, resolver, mappings)

+                );

+        }

+

+        private EntryMapping resolveMapping(EntryResolver resolver, EntryMap<EntryMapping> mappings, MethodEntry methodEntry) {

+                for (MethodEntry entry : resolver.resolveEntry(methodEntry, ResolutionStrategy.RESOLVE_ROOT)) {

+                        EntryMapping mapping = mappings.get(entry);

+                        if (mapping != null) {

+                                return mapping;

+                        }

+                }

+                return null;

+        }

+

+        public ClassEntry getInterface() {

+                return invokedType.getReturnDesc().getTypeEntry();

+        }

+

+        public String getInvokedName() {

+                return invokedName;

+        }

+

+        public MethodDescriptor getInvokedType() {

+                return invokedType;

+        }

+

+        public MethodDescriptor getSamMethodType() {

+                return samMethodType;

+        }

+

+        public ParentedEntry<?> getImplMethod() {

+                return implMethod;

+        }

+

+        public MethodDescriptor getInstantiatedMethodType() {

+                return instantiatedMethodType;

+        }

+

+        @Override

+        public boolean equals(Object o) {

+                if (this == o) return true;

+                if (o == null || getClass() != o.getClass()) return false;

+                Lambda lambda = (Lambda) o;

+                return Objects.equals(invokedName, lambda.invokedName) &&

+                        Objects.equals(invokedType, lambda.invokedType) &&

+                        Objects.equals(samMethodType, lambda.samMethodType) &&

+                        Objects.equals(implMethod, lambda.implMethod) &&

+                        Objects.equals(instantiatedMethodType, lambda.instantiatedMethodType);

+        }

+

+        @Override

+        public int hashCode() {

+                return Objects.hash(invokedName, invokedType, samMethodType, implMethod, instantiatedMethodType);

+        }

+

+        @Override

+        public String toString() {

+                return "Lambda{" +

+                        "invokedName='" + invokedName + '\'' +

+                        ", invokedType=" + invokedType +

+                        ", samMethodType=" + samMethodType +

+                        ", implMethod=" + implMethod +

+                        ", instantiatedMethodType=" + instantiatedMethodType +

+                        '}';

+        }

+}

diff --git a/src/main/java/cuchaz/enigma/translation/representation/MethodDescriptor.java b/src/main/java/cuchaz/enigma/translation/representation/MethodDescriptor.java
index c59751f9..37a70148 100644
--- a/src/main/java/cuchaz/enigma/translation/representation/MethodDescriptor.java
+++ b/src/main/java/cuchaz/enigma/translation/representation/MethodDescriptor.java
@@ -118,7 +118,7 @@ public class MethodDescriptor implements Translatable {
         }
 
         @Override
-        public Translatable translate(Translator translator, EntryResolver resolver, EntryMap<EntryMapping> mappings) {
+        public MethodDescriptor translate(Translator translator, EntryResolver resolver, EntryMap<EntryMapping> mappings) {
                 List<TypeDescriptor> translatedArguments = new ArrayList<>(argumentDescs.size());
                 for (TypeDescriptor argument : argumentDescs) {
                         translatedArguments.add(translator.translate(argument));
diff --git a/src/main/java/cuchaz/enigma/translation/representation/entry/ParentedEntry.java b/src/main/java/cuchaz/enigma/translation/representation/entry/ParentedEntry.java
index 834da8de..b753d3a3 100644
--- a/src/main/java/cuchaz/enigma/translation/representation/entry/ParentedEntry.java
+++ b/src/main/java/cuchaz/enigma/translation/representation/entry/ParentedEntry.java
@@ -52,7 +52,7 @@ public abstract class ParentedEntry<P extends Entry<?>> implements Entry<P> {
         }
 
         @Override
-        public Translatable translate(Translator translator, EntryResolver resolver, EntryMap<EntryMapping> mappings) {
+        public ParentedEntry<P> translate(Translator translator, EntryResolver resolver, EntryMap<EntryMapping> mappings) {
                 P parent = getParent();
                 EntryMapping mapping = resolveMapping(resolver, mappings);
                 if (parent == null) {