Method type reference corrections (#142)

* Add more specific returns for translatables

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

4 files changed, 60 insertions, 31 deletions

diff --git a/src/main/java/cuchaz/enigma/analysis/index/IndexReferenceVisitor.java b/src/main/java/cuchaz/enigma/analysis/index/IndexReferenceVisitor.java
index ba5d3b6..b730a8a 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 ac907af..fd4e618 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 457c223..4f885b3 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 04306bd..f54c90d 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);