6 files changed, 254 insertions, 34 deletions

diff --git a/src/main/java/cuchaz/enigma/analysis/index/EntryIndex.java b/src/main/java/cuchaz/enigma/analysis/index/EntryIndex.java
index 773eaf1..31c6f54 100644
--- a/src/main/java/cuchaz/enigma/analysis/index/EntryIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/index/EntryIndex.java
@@ -65,6 +65,11 @@ public class EntryIndex implements JarIndexer {
         }
 
         @Nullable
+        public AccessFlags getClassAccess(ClassEntry entry) {
+                return classes.get(entry);
+        }
+
+        @Nullable
         public AccessFlags getEntryAccess(Entry<?> entry) {
                 if (entry instanceof MethodEntry) {
                         return getMethodAccess((MethodEntry) entry);
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 a429ff6..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;
@@ -29,18 +30,20 @@ public class JarIndex implements JarIndexer {
         private final InheritanceIndex inheritanceIndex;
         private final ReferenceIndex referenceIndex;
         private final BridgeMethodIndex bridgeMethodIndex;
+        private final PackageVisibilityIndex packageVisibilityIndex;
         private final EntryResolver entryResolver;
 
         private final Collection<JarIndexer> indexers;
 
         private final Multimap<String, MethodDefEntry> methodImplementations = HashMultimap.create();
 
-        public JarIndex(EntryIndex entryIndex, InheritanceIndex inheritanceIndex, ReferenceIndex referenceIndex, BridgeMethodIndex bridgeMethodIndex) {
+        public JarIndex(EntryIndex entryIndex, InheritanceIndex inheritanceIndex, ReferenceIndex referenceIndex, BridgeMethodIndex bridgeMethodIndex, PackageVisibilityIndex packageVisibilityIndex) {
                 this.entryIndex = entryIndex;
                 this.inheritanceIndex = inheritanceIndex;
                 this.referenceIndex = referenceIndex;
                 this.bridgeMethodIndex = bridgeMethodIndex;
-                this.indexers = Arrays.asList(entryIndex, inheritanceIndex, referenceIndex, bridgeMethodIndex);
+                this.packageVisibilityIndex = packageVisibilityIndex;
+                this.indexers = Arrays.asList(entryIndex, inheritanceIndex, referenceIndex, bridgeMethodIndex, packageVisibilityIndex);
                 this.entryResolver = new IndexEntryResolver(this);
         }
 
@@ -49,7 +52,8 @@ public class JarIndex implements JarIndexer {
                 InheritanceIndex inheritanceIndex = new InheritanceIndex(entryIndex);
                 ReferenceIndex referenceIndex = new ReferenceIndex();
                 BridgeMethodIndex bridgeMethodIndex = new BridgeMethodIndex(entryIndex, inheritanceIndex, referenceIndex);
-                return new JarIndex(entryIndex, inheritanceIndex, referenceIndex, bridgeMethodIndex);
+                PackageVisibilityIndex packageVisibilityIndex = new PackageVisibilityIndex();
+                return new JarIndex(entryIndex, inheritanceIndex, referenceIndex, bridgeMethodIndex, packageVisibilityIndex);
         }
 
         public void indexJar(ParsedJar jar, Consumer<String> progress) {
@@ -126,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;
         }
@@ -142,6 +155,10 @@ public class JarIndex implements JarIndexer {
                 return bridgeMethodIndex;
         }
 
+        public PackageVisibilityIndex getPackageVisibilityIndex() {
+                return packageVisibilityIndex;
+        }
+
         public EntryResolver getEntryResolver() {
                 return entryResolver;
         }
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/PackageVisibilityIndex.java b/src/main/java/cuchaz/enigma/analysis/index/PackageVisibilityIndex.java
new file mode 100644
index 0000000..da28ac4
--- /dev/null
+++ b/src/main/java/cuchaz/enigma/analysis/index/PackageVisibilityIndex.java
@@ -0,0 +1,127 @@
+package cuchaz.enigma.analysis.index;

+

+import com.google.common.collect.HashMultimap;

+import com.google.common.collect.Lists;

+import com.google.common.collect.Maps;

+import com.google.common.collect.Sets;

+import cuchaz.enigma.analysis.EntryReference;

+import cuchaz.enigma.translation.representation.AccessFlags;

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

+

+import java.util.*;

+

+public class PackageVisibilityIndex implements JarIndexer {

+        private static boolean requiresSamePackage(AccessFlags entryAcc, EntryReference ref, InheritanceIndex inheritanceIndex) {

+                if (entryAcc.isPublic()) return false;

+                if (entryAcc.isProtected()) {

+                        Set<ClassEntry> callerAncestors = inheritanceIndex.getAncestors(ref.context.getContainingClass());

+                        return !callerAncestors.contains(ref.entry.getContainingClass());

+                }

+                return !entryAcc.isPrivate(); // if isPrivate is false, it must be package-private

+        }

+

+        private final HashMultimap<ClassEntry, ClassEntry> connections = HashMultimap.create();

+        private final List<Set<ClassEntry>> partitions = Lists.newArrayList();

+        private final Map<ClassEntry, Set<ClassEntry>> classPartitions = Maps.newHashMap();

+

+        private void addConnection(ClassEntry classA, ClassEntry classB) {

+                connections.put(classA, classB);

+                connections.put(classB, classA);

+        }

+

+        private void buildPartition(Set<ClassEntry> unassignedClasses, Set<ClassEntry> partition, ClassEntry member) {

+                for (ClassEntry connected : connections.get(member)) {

+                        if (unassignedClasses.remove(connected)) {

+                                partition.add(connected);

+                                buildPartition(unassignedClasses, partition, connected);

+                        }

+                }

+        }

+

+        private void addConnections(EntryIndex entryIndex, ReferenceIndex referenceIndex, InheritanceIndex inheritanceIndex) {

+                for (FieldEntry entry : entryIndex.getFields()) {

+                        AccessFlags entryAcc = entryIndex.getFieldAccess(entry);

+                        if (!entryAcc.isPublic() && !entryAcc.isPrivate()) {

+                                for (EntryReference<FieldEntry, MethodDefEntry> ref : referenceIndex.getReferencesToField(entry)) {

+                                        if (requiresSamePackage(entryAcc, ref, inheritanceIndex)) {

+                                                addConnection(ref.entry.getContainingClass(), ref.context.getContainingClass());

+                                        }

+                                }

+                        }

+                }

+

+                for (MethodEntry entry : entryIndex.getMethods()) {

+                        AccessFlags entryAcc = entryIndex.getMethodAccess(entry);

+                        if (!entryAcc.isPublic() && !entryAcc.isPrivate()) {

+                                for (EntryReference<MethodEntry, MethodDefEntry> ref : referenceIndex.getReferencesToMethod(entry)) {

+                                        if (requiresSamePackage(entryAcc, ref, inheritanceIndex)) {

+                                                addConnection(ref.entry.getContainingClass(), ref.context.getContainingClass());

+                                        }

+                                }

+                        }

+                }

+

+                for (ClassEntry entry : entryIndex.getClasses()) {

+                        AccessFlags entryAcc = entryIndex.getClassAccess(entry);

+                        if (!entryAcc.isPublic() && !entryAcc.isPrivate()) {

+                                for (EntryReference<ClassEntry, FieldDefEntry> ref : referenceIndex.getFieldTypeReferencesToClass(entry)) {

+                                        if (requiresSamePackage(entryAcc, ref, inheritanceIndex)) {

+                                                addConnection(ref.entry.getContainingClass(), ref.context.getContainingClass());

+                                        }

+                                }

+

+                                for (EntryReference<ClassEntry, MethodDefEntry> ref : referenceIndex.getMethodTypeReferencesToClass(entry)) {

+                                        if (requiresSamePackage(entryAcc, ref, inheritanceIndex)) {

+                                                addConnection(ref.entry.getContainingClass(), ref.context.getContainingClass());

+                                        }

+                                }

+                        }

+

+                        for (ClassEntry parent : inheritanceIndex.getParents(entry)) {

+                                AccessFlags parentAcc = entryIndex.getClassAccess(parent);

+                                if (parentAcc != null && !parentAcc.isPublic() && !parentAcc.isPrivate()) {

+                                        addConnection(entry, parent);

+                                }

+                        }

+

+                        ClassEntry outerClass = entry.getOuterClass();

+                        if (outerClass != null) {

+                                addConnection(entry, outerClass);

+                        }

+                }

+        }

+

+        private void addPartitions(EntryIndex entryIndex) {

+                Set<ClassEntry> unassignedClasses = Sets.newHashSet(entryIndex.getClasses());

+                while (!unassignedClasses.isEmpty()) {

+                        Iterator<ClassEntry> iterator = unassignedClasses.iterator();

+                        ClassEntry initialEntry = iterator.next();

+                        iterator.remove();

+

+                        HashSet<ClassEntry> partition = Sets.newHashSet();

+                        partition.add(initialEntry);

+                        buildPartition(unassignedClasses, partition, initialEntry);

+                        partitions.add(partition);

+                        for (ClassEntry entry : partition) {

+                                classPartitions.put(entry, partition);

+                        }

+                }

+        }

+

+        public Collection<Set<ClassEntry>> getPartitions() {

+                return partitions;

+        }

+

+        public Set<ClassEntry> getPartition(ClassEntry classEntry) {

+                return classPartitions.get(classEntry);

+        }

+

+        @Override

+        public void processIndex(JarIndex index) {

+                EntryIndex entryIndex = index.getEntryIndex();

+                ReferenceIndex referenceIndex = index.getReferenceIndex();

+                InheritanceIndex inheritanceIndex = index.getInheritanceIndex();

+                addConnections(entryIndex, referenceIndex, inheritanceIndex);

+                addPartitions(entryIndex);

+        }

+}

diff --git a/src/main/java/cuchaz/enigma/analysis/index/ReferenceIndex.java b/src/main/java/cuchaz/enigma/analysis/index/ReferenceIndex.java
index 2b63c5d..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,9 @@ 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.*;
 
 import java.util.Collection;
@@ -15,6 +18,43 @@ public class ReferenceIndex implements JarIndexer {
         private Multimap<MethodEntry, EntryReference<MethodEntry, MethodDefEntry>> referencesToMethods = HashMultimap.create();
         private Multimap<ClassEntry, EntryReference<ClassEntry, MethodDefEntry>> referencesToClasses = HashMultimap.create();
         private Multimap<FieldEntry, EntryReference<FieldEntry, MethodDefEntry>> referencesToFields = HashMultimap.create();
+        private Multimap<ClassEntry, EntryReference<ClassEntry, FieldDefEntry>> fieldTypeReferences = HashMultimap.create();
+        private Multimap<ClassEntry, EntryReference<ClassEntry, MethodDefEntry>> methodTypeReferences = HashMultimap.create();
+
+        @Override
+        public void indexMethod(MethodDefEntry methodEntry) {
+            indexMethodDescriptor(methodEntry, methodEntry.getDesc());
+        }
+
+        private void indexMethodDescriptor(MethodDefEntry entry, MethodDescriptor descriptor) {
+                for (TypeDescriptor typeDescriptor : descriptor.getArgumentDescs()) {
+                        indexMethodTypeDescriptor(entry, typeDescriptor);
+                }
+                indexMethodTypeDescriptor(entry, descriptor.getReturnDesc());
+        }
+
+        private void indexMethodTypeDescriptor(MethodDefEntry method, TypeDescriptor typeDescriptor) {
+                if (typeDescriptor.isType()) {
+                        ClassEntry referencedClass = typeDescriptor.getTypeEntry();
+                        methodTypeReferences.put(referencedClass, new EntryReference<>(referencedClass, referencedClass.getName(), method));
+                } else if (typeDescriptor.isArray()) {
+                        indexMethodTypeDescriptor(method, typeDescriptor.getArrayType());
+                }
+        }
+
+        @Override
+        public void indexField(FieldDefEntry fieldEntry) {
+            indexFieldTypeDescriptor(fieldEntry, fieldEntry.getDesc());
+        }
+
+        private void indexFieldTypeDescriptor(FieldDefEntry field, TypeDescriptor typeDescriptor) {
+                if (typeDescriptor.isType()) {
+                        ClassEntry referencedClass = typeDescriptor.getTypeEntry();
+                        fieldTypeReferences.put(referencedClass, new EntryReference<>(referencedClass, referencedClass.getName(), field));
+                } else if (typeDescriptor.isArray()) {
+                    indexFieldTypeDescriptor(field, typeDescriptor.getArrayType());
+                }
+        }
 
         @Override
         public void indexMethodReference(MethodDefEntry callerEntry, MethodEntry referencedEntry) {
@@ -33,15 +73,30 @@ 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);
                 referencesToClasses = remapReferencesTo(index, referencesToClasses);
                 referencesToFields = remapReferencesTo(index, referencesToFields);
+                fieldTypeReferences = remapReferencesTo(index, fieldTypeReferences);
+                methodTypeReferences = remapReferencesTo(index, methodTypeReferences);
         }
 
         private <K extends Entry<?>, V extends Entry<?>> Multimap<K, V> remapReferences(JarIndex index, Multimap<K, V> multimap) {
-                Multimap<K, V> resolved = HashMultimap.create();
+                Multimap<K, V> resolved = HashMultimap.create(multimap.keySet().size(), multimap.size() / multimap.keySet().size());
                 for (Map.Entry<K, V> entry : multimap.entries()) {
                         resolved.put(remap(index, entry.getKey()), remap(index, entry.getValue()));
                 }
@@ -49,7 +104,8 @@ public class ReferenceIndex implements JarIndexer {
         }
 
         private <E extends Entry<?>, C extends Entry<?>> Multimap<E, EntryReference<E, C>> remapReferencesTo(JarIndex index, Multimap<E, EntryReference<E, C>> multimap) {
-                Multimap<E, EntryReference<E, C>> resolved = HashMultimap.create();
+                final int keySetSize = multimap.keySet().size();
+                Multimap<E, EntryReference<E, C>> resolved = HashMultimap.create(keySetSize, keySetSize == 0 ? 0 : multimap.size() / keySetSize);
                 for (Map.Entry<E, EntryReference<E, C>> entry : multimap.entries()) {
                         resolved.put(remap(index, entry.getKey()), remap(index, entry.getValue()));
                 }
@@ -79,4 +135,12 @@ public class ReferenceIndex implements JarIndexer {
         public Collection<EntryReference<MethodEntry, MethodDefEntry>> getReferencesToMethod(MethodEntry entry) {
                 return referencesToMethods.get(entry);
         }
+
+        public Collection<EntryReference<ClassEntry, FieldDefEntry>> getFieldTypeReferencesToClass(ClassEntry entry) {
+                return fieldTypeReferences.get(entry);
+        }
+
+        public Collection<EntryReference<ClassEntry, MethodDefEntry>> getMethodTypeReferencesToClass(ClassEntry entry) {
+                return methodTypeReferences.get(entry);
+        }
 }