5 files changed, 78 insertions, 29 deletions

diff --git a/src/main/java/cuchaz/enigma/analysis/index/BridgeMethodIndex.java b/src/main/java/cuchaz/enigma/analysis/index/BridgeMethodIndex.java
index de2ba1e..649ce25 100644
--- a/src/main/java/cuchaz/enigma/analysis/index/BridgeMethodIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/index/BridgeMethodIndex.java
@@ -2,7 +2,6 @@ package cuchaz.enigma.analysis.index;
 
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
-import cuchaz.enigma.translation.mapping.EntryResolver;
 import cuchaz.enigma.translation.representation.AccessFlags;
 import cuchaz.enigma.translation.representation.MethodDescriptor;
 import cuchaz.enigma.translation.representation.TypeDescriptor;
@@ -11,10 +10,7 @@ import cuchaz.enigma.translation.representation.entry.MethodDefEntry;
 import cuchaz.enigma.translation.representation.entry.MethodEntry;
 
 import javax.annotation.Nullable;
-import java.util.Collection;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 
 public class BridgeMethodIndex implements JarIndexer {
         private final EntryIndex entryIndex;
@@ -30,8 +26,7 @@ public class BridgeMethodIndex implements JarIndexer {
                 this.referenceIndex = referenceIndex;
         }
 
-        @Override
-        public void processIndex(EntryResolver resolver) {
+        public void findBridgeMethods() {
                 // look for access and bridged methods
                 for (MethodEntry methodEntry : entryIndex.getMethods()) {
                         MethodDefEntry methodDefEntry = (MethodDefEntry) methodEntry;
@@ -45,6 +40,23 @@ public class BridgeMethodIndex implements JarIndexer {
                 }
         }
 
+        @Override
+        public void processIndex(JarIndex index) {
+                Map<MethodEntry, MethodEntry> copiedAccessToBridge = new HashMap<>(accessedToBridge);
+
+                for (Map.Entry<MethodEntry, MethodEntry> entry : copiedAccessToBridge.entrySet()) {
+                        MethodEntry accessedEntry = entry.getKey();
+                        MethodEntry bridgeEntry = entry.getValue();
+                        if (bridgeEntry.getName().equals(accessedEntry.getName())) {
+                                continue;
+                        }
+
+                        MethodEntry renamedAccessedEntry = accessedEntry.withName(bridgeEntry.getName());
+                        bridgeMethods.add(renamedAccessedEntry);
+                        accessedToBridge.put(renamedAccessedEntry, accessedToBridge.get(accessedEntry));
+                }
+        }
+
         private void indexSyntheticMethod(MethodDefEntry syntheticMethod, AccessFlags access) {
                 MethodEntry accessedMethod = findAccessMethod(syntheticMethod);
                 if (accessedMethod == null) {
@@ -129,4 +141,8 @@ public class BridgeMethodIndex implements JarIndexer {
         public MethodEntry getBridgeFromAccessed(MethodEntry entry) {
                 return accessedToBridge.get(entry);
         }
+
+        public Map<MethodEntry, MethodEntry> getAccessedToBridge() {
+                return Collections.unmodifiableMap(accessedToBridge);
+        }
 }
diff --git a/src/main/java/cuchaz/enigma/analysis/index/EntryIndex.java b/src/main/java/cuchaz/enigma/analysis/index/EntryIndex.java
index 773eaf1..645110a 100644
--- a/src/main/java/cuchaz/enigma/analysis/index/EntryIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/index/EntryIndex.java
@@ -28,6 +28,18 @@ public class EntryIndex implements JarIndexer {
                 fields.put(fieldEntry, fieldEntry.getAccess());
         }
 
+        @Override
+        public void processIndex(JarIndex index) {
+                Map<MethodEntry, MethodEntry> accessedToBridge = index.getBridgeMethodIndex().getAccessedToBridge();
+                for (Map.Entry<MethodEntry, MethodEntry> entry : accessedToBridge.entrySet()) {
+                        MethodEntry accessedEntry = entry.getKey();
+                        MethodEntry bridgeEntry = entry.getValue();
+
+                        MethodEntry renamedAccessedEntry = accessedEntry.withName(bridgeEntry.getName());
+                        methods.put(renamedAccessedEntry, methods.remove(accessedEntry));
+                }
+        }
+
         public boolean hasClass(ClassEntry entry) {
                 return classes.containsKey(entry);
         }
diff --git a/src/main/java/cuchaz/enigma/analysis/index/JarIndex.java b/src/main/java/cuchaz/enigma/analysis/index/JarIndex.java
index cb58fce..a429ff6 100644
--- a/src/main/java/cuchaz/enigma/analysis/index/JarIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/index/JarIndex.java
@@ -53,19 +53,22 @@ public class JarIndex implements JarIndexer {
         }
 
         public void indexJar(ParsedJar jar, Consumer<String> progress) {
-                progress.accept("Indexing entries (1/3)");
+                progress.accept("Indexing entries (1/4)");
                 jar.visitReader(name -> new IndexClassVisitor(this, Opcodes.ASM5), ClassReader.SKIP_CODE);
 
-                progress.accept("Indexing entry references (2/3)");
+                progress.accept("Indexing entry references (2/4)");
                 jar.visitReader(name -> new IndexReferenceVisitor(this, Opcodes.ASM5), ClassReader.SKIP_FRAMES);
 
-                progress.accept("Processing index (3/3)");
-                processIndex(entryResolver);
+                progress.accept("Finding bridge methods (3/4)");
+                bridgeMethodIndex.findBridgeMethods();
+
+                progress.accept("Processing index (4/4)");
+                processIndex(this);
         }
 
         @Override
-        public void processIndex(EntryResolver resolver) {
-                indexers.forEach(indexer -> indexer.processIndex(entryResolver));
+        public void processIndex(JarIndex index) {
+                indexers.forEach(indexer -> indexer.processIndex(index));
         }
 
         @Override
diff --git a/src/main/java/cuchaz/enigma/analysis/index/JarIndexer.java b/src/main/java/cuchaz/enigma/analysis/index/JarIndexer.java
index a087e59..457c223 100644
--- a/src/main/java/cuchaz/enigma/analysis/index/JarIndexer.java
+++ b/src/main/java/cuchaz/enigma/analysis/index/JarIndexer.java
@@ -1,6 +1,5 @@
 package cuchaz.enigma.analysis.index;
 
-import cuchaz.enigma.translation.mapping.EntryResolver;
 import cuchaz.enigma.translation.representation.entry.*;
 
 public interface JarIndexer {
@@ -19,6 +18,6 @@ public interface JarIndexer {
         default void indexFieldReference(MethodDefEntry callerEntry, FieldEntry referencedEntry) {
         }
 
-        default void processIndex(EntryResolver resolver) {
+        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 ac11da4..ee28b3e 100644
--- a/src/main/java/cuchaz/enigma/analysis/index/ReferenceIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/index/ReferenceIndex.java
@@ -3,12 +3,12 @@ package cuchaz.enigma.analysis.index;
 import com.google.common.collect.HashMultimap;
 import com.google.common.collect.Multimap;
 import cuchaz.enigma.analysis.EntryReference;
-import cuchaz.enigma.translation.mapping.EntryResolver;
 import cuchaz.enigma.translation.mapping.ResolutionStrategy;
 import cuchaz.enigma.translation.representation.entry.*;
 
 import java.util.Collection;
 import java.util.Map;
+import java.util.Optional;
 
 public class ReferenceIndex implements JarIndexer {
         private Multimap<MethodEntry, MethodEntry> methodReferences = HashMultimap.create();
@@ -34,35 +34,54 @@ public class ReferenceIndex implements JarIndexer {
         }
 
         @Override
-        public void processIndex(EntryResolver resolver) {
-                methodReferences = resolveReferences(resolver, methodReferences);
-                referencesToMethods = resolveReferencesTo(resolver, referencesToMethods);
-                referencesToClasses = resolveReferencesTo(resolver, referencesToClasses);
-                referencesToFields = resolveReferencesTo(resolver, referencesToFields);
+        public void processIndex(JarIndex index) {
+                methodReferences = remapReferences(index, methodReferences);
+                referencesToMethods = remapReferencesTo(index, referencesToMethods);
+                referencesToClasses = remapReferencesTo(index, referencesToClasses);
+                referencesToFields = remapReferencesTo(index, referencesToFields);
         }
 
-        private <K extends Entry<?>, V extends Entry<?>> Multimap<K, V> resolveReferences(EntryResolver resolver, Multimap<K, V> multimap) {
+        private <K extends Entry<?>, V extends Entry<?>> Multimap<K, V> remapReferences(JarIndex index, Multimap<K, V> multimap) {
                 Multimap<K, V> resolved = HashMultimap.create();
                 for (Map.Entry<K, V> entry : multimap.entries()) {
-                        resolved.put(resolve(resolver, entry.getKey()), resolve(resolver, entry.getValue()));
+                        resolved.put(remap(index, entry.getKey()), remap(index, entry.getValue()));
                 }
                 return resolved;
         }
 
-        private <E extends Entry<?>, C extends Entry<?>> Multimap<E, EntryReference<E, C>> resolveReferencesTo(EntryResolver resolver, Multimap<E, EntryReference<E, C>> multimap) {
+        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();
                 for (Map.Entry<E, EntryReference<E, C>> entry : multimap.entries()) {
-                        resolved.put(resolve(resolver, entry.getKey()), resolve(resolver, entry.getValue()));
+                        resolved.put(remap(index, entry.getKey()), remap(index, entry.getValue()));
                 }
                 return resolved;
         }
 
-        private <E extends Entry<?>> E resolve(EntryResolver resolver, E entry) {
-                return resolver.resolveFirstEntry(entry, ResolutionStrategy.RESOLVE_CLOSEST);
+        private <E extends Entry<?>> E remap(JarIndex index, E entry) {
+                E resolvedEntry = index.getEntryResolver().resolveFirstEntry(entry, ResolutionStrategy.RESOLVE_CLOSEST);
+
+                Optional<E> remappedBridge = getRemappedBridge(index, resolvedEntry);
+                return remappedBridge.orElse(resolvedEntry);
+        }
+
+        private <E extends Entry<?>, C extends Entry<?>> EntryReference<E, C> remap(JarIndex index, EntryReference<E, C> reference) {
+                EntryReference<E, C> resolvedReference = index.getEntryResolver().resolveFirstReference(reference, ResolutionStrategy.RESOLVE_CLOSEST);
+
+                getRemappedBridge(index, resolvedReference.entry).ifPresent(e -> resolvedReference.entry = e);
+                getRemappedBridge(index, resolvedReference.context).ifPresent(e -> resolvedReference.context = e);
+
+                return resolvedReference;
         }
 
-        private <E extends Entry<?>, C extends Entry<?>> EntryReference<E, C> resolve(EntryResolver resolver, EntryReference<E, C> reference) {
-                return resolver.resolveFirstReference(reference, ResolutionStrategy.RESOLVE_CLOSEST);
+        @SuppressWarnings("unchecked")
+        private <E extends Entry<?>> Optional<E> getRemappedBridge(JarIndex index, E entry) {
+                if (entry instanceof MethodEntry) {
+                        MethodEntry bridgeEntry = index.getBridgeMethodIndex().getBridgeFromAccessed((MethodEntry) entry);
+                        if (bridgeEntry != null) {
+                                return Optional.of((E) entry.withName(bridgeEntry.getName()));
+                        }
+                }
+                return Optional.empty();
         }
 
         public Collection<MethodEntry> getMethodsReferencedBy(MethodEntry entry) {