Correctly decompile bridges, and add command to add bridges to mappings (#180)

5 files changed, 85 insertions, 7 deletions

diff --git a/src/main/java/cuchaz/enigma/CommandMain.java b/src/main/java/cuchaz/enigma/CommandMain.java
index 5e7a1af6..93b013d5 100644
--- a/src/main/java/cuchaz/enigma/CommandMain.java
+++ b/src/main/java/cuchaz/enigma/CommandMain.java
@@ -86,6 +86,7 @@ public class CommandMain {
                 register(new ComposeMappingsCommand());
                 register(new InvertMappingsCommand());
                 register(new CheckMappingsCommand());
+                register(new MapSpecializedMethodsCommand());
         }
 
         private static final class CommandHelpException extends IllegalArgumentException {
diff --git a/src/main/java/cuchaz/enigma/EnigmaProject.java b/src/main/java/cuchaz/enigma/EnigmaProject.java
index 47a37907..2a7ca988 100644
--- a/src/main/java/cuchaz/enigma/EnigmaProject.java
+++ b/src/main/java/cuchaz/enigma/EnigmaProject.java
@@ -158,7 +158,7 @@ public class EnigmaProject {
                                         ClassNode node = classCache.getClassNode(entry.getFullName());
                                         if (node != null) {
                                                 ClassNode translatedNode = new ClassNode();
-                                                node.accept(new TranslationClassVisitor(deobfuscator, Opcodes.ASM5, translatedNode));
+                                                node.accept(new TranslationClassVisitor(deobfuscator, Opcodes.ASM5, new SourceFixVisitor(Opcodes.ASM5, translatedNode, jarIndex)));
                                                 return translatedNode;
                                         }
 
@@ -209,7 +209,6 @@ public class EnigmaProject {
 
                         //create a common instance outside the loop as mappings shouldn't be changing while this is happening
                         CompiledSourceTypeLoader typeLoader = new CompiledSourceTypeLoader(this.compiled::get);
-                        typeLoader.addVisitor(visitor -> new SourceFixVisitor(Opcodes.ASM5, visitor, jarIndex));
 
                         //synchronized to make sure the parallelStream doesn't CME with the cache
                         ITypeLoader synchronizedTypeLoader = new SynchronizedTypeLoader(typeLoader);
diff --git a/src/main/java/cuchaz/enigma/analysis/index/BridgeMethodIndex.java b/src/main/java/cuchaz/enigma/analysis/index/BridgeMethodIndex.java
index 8a9b4593..a4b1aac9 100644
--- a/src/main/java/cuchaz/enigma/analysis/index/BridgeMethodIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/index/BridgeMethodIndex.java
@@ -1,7 +1,6 @@
 package cuchaz.enigma.analysis.index;
 
 import com.google.common.collect.Maps;
-import com.google.common.collect.Sets;
 import cuchaz.enigma.translation.representation.AccessFlags;
 import cuchaz.enigma.translation.representation.MethodDescriptor;
 import cuchaz.enigma.translation.representation.TypeDescriptor;
@@ -17,7 +16,7 @@ public class BridgeMethodIndex implements JarIndexer {
         private final InheritanceIndex inheritanceIndex;
         private final ReferenceIndex referenceIndex;
 
-        private final Set<MethodEntry> bridgeToSpecialized = Sets.newHashSet();
+        private final Map<MethodEntry, MethodEntry> bridgeToSpecialized = Maps.newHashMap();
         private final Map<MethodEntry, MethodEntry> specializedToBridge = Maps.newHashMap();
 
         public BridgeMethodIndex(EntryIndex entryIndex, InheritanceIndex inheritanceIndex, ReferenceIndex referenceIndex) {
@@ -52,7 +51,6 @@ public class BridgeMethodIndex implements JarIndexer {
                         }
 
                         MethodEntry renamedSpecializedEntry = specializedEntry.withName(bridgeEntry.getName());
-                        bridgeToSpecialized.add(renamedSpecializedEntry);
                         specializedToBridge.put(renamedSpecializedEntry, specializedToBridge.get(specializedEntry));
                 }
         }
@@ -64,7 +62,7 @@ public class BridgeMethodIndex implements JarIndexer {
                 }
 
                 if (access.isBridge() || isPotentialBridge(syntheticMethod, specializedMethod)) {
-                        bridgeToSpecialized.add(syntheticMethod);
+                        bridgeToSpecialized.put(syntheticMethod, specializedMethod);
                         specializedToBridge.put(specializedMethod, syntheticMethod);
                 }
         }
@@ -130,7 +128,7 @@ public class BridgeMethodIndex implements JarIndexer {
         }
 
         public boolean isBridgeMethod(MethodEntry entry) {
-                return bridgeToSpecialized.contains(entry);
+                return bridgeToSpecialized.containsKey(entry);
         }
 
         public boolean isSpecializedMethod(MethodEntry entry) {
@@ -142,7 +140,17 @@ public class BridgeMethodIndex implements JarIndexer {
                 return specializedToBridge.get(specialized);
         }
 
+        public MethodEntry getSpecializedFromBridge(MethodEntry bridge) {
+                return bridgeToSpecialized.get(bridge);
+        }
+
+        /** Includes "renamed specialized -> bridge" entries. */
         public Map<MethodEntry, MethodEntry> getSpecializedToBridge() {
                 return Collections.unmodifiableMap(specializedToBridge);
         }
+
+        /** Only "bridge -> original name" entries. **/
+        public Map<MethodEntry, MethodEntry> getBridgeToSpecialized() {
+                return Collections.unmodifiableMap(bridgeToSpecialized);
+        }
 }
diff --git a/src/main/java/cuchaz/enigma/bytecode/translators/SourceFixVisitor.java b/src/main/java/cuchaz/enigma/bytecode/translators/SourceFixVisitor.java
index 735a65f2..d672622c 100644
--- a/src/main/java/cuchaz/enigma/bytecode/translators/SourceFixVisitor.java
+++ b/src/main/java/cuchaz/enigma/bytecode/translators/SourceFixVisitor.java
@@ -4,6 +4,7 @@ import cuchaz.enigma.analysis.index.BridgeMethodIndex;
 import cuchaz.enigma.analysis.index.JarIndex;
 import cuchaz.enigma.translation.representation.entry.ClassDefEntry;
 import cuchaz.enigma.translation.representation.entry.MethodDefEntry;
+import cuchaz.enigma.translation.representation.entry.MethodEntry;
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.MethodVisitor;
 import org.objectweb.asm.Opcodes;
@@ -30,6 +31,8 @@ public class SourceFixVisitor extends ClassVisitor {
                 BridgeMethodIndex bridgeIndex = index.getBridgeMethodIndex();
                 if (bridgeIndex.isBridgeMethod(methodEntry)) {
                         access |= Opcodes.ACC_BRIDGE;
+                } else if (bridgeIndex.isSpecializedMethod(methodEntry)) {
+                        name = bridgeIndex.getBridgeFromSpecialized(methodEntry).getName();
                 }
 
                 return super.visitMethod(access, name, descriptor, signature, exceptions);
diff --git a/src/main/java/cuchaz/enigma/command/MapSpecializedMethodsCommand.java b/src/main/java/cuchaz/enigma/command/MapSpecializedMethodsCommand.java
new file mode 100644
index 00000000..7696f0f1
--- /dev/null
+++ b/src/main/java/cuchaz/enigma/command/MapSpecializedMethodsCommand.java
@@ -0,0 +1,67 @@
+package cuchaz.enigma.command;
+
+import cuchaz.enigma.ProgressListener;
+import cuchaz.enigma.analysis.ClassCache;
+import cuchaz.enigma.analysis.index.BridgeMethodIndex;
+import cuchaz.enigma.analysis.index.JarIndex;
+import cuchaz.enigma.throwables.MappingParseException;
+import cuchaz.enigma.translation.MappingTranslator;
+import cuchaz.enigma.translation.Translator;
+import cuchaz.enigma.translation.mapping.*;
+import cuchaz.enigma.translation.mapping.tree.EntryTree;
+import cuchaz.enigma.translation.mapping.tree.EntryTreeNode;
+import cuchaz.enigma.translation.mapping.tree.HashEntryTree;
+import cuchaz.enigma.translation.representation.entry.MethodEntry;
+import cuchaz.enigma.utils.Utils;
+
+import java.io.IOException;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.Map;
+
+public class MapSpecializedMethodsCommand extends Command {
+    public MapSpecializedMethodsCommand() {
+        super("map-specialized-methods");
+    }
+
+    @Override
+    public String getUsage() {
+        return "<jar> <source-format> <source> <result-format> <result>";
+    }
+
+    @Override
+    public boolean isValidArgument(int length) {
+        return length == 5;
+    }
+
+    @Override
+    public void run(String... args) throws IOException, MappingParseException {
+        MappingSaveParameters saveParameters = new MappingSaveParameters(MappingFileNameFormat.BY_DEOBF);
+        EntryTree<EntryMapping> source = MappingCommandsUtil.read(args[1], Paths.get(args[2]), saveParameters);
+        EntryTree<EntryMapping> result = new HashEntryTree<>();
+        Path jar = Paths.get(args[0]);
+        ClassCache classCache = ClassCache.of(jar);
+        JarIndex jarIndex = classCache.index(ProgressListener.none());
+        BridgeMethodIndex bridgeMethodIndex = jarIndex.getBridgeMethodIndex();
+        Translator translator = new MappingTranslator(source, jarIndex.getEntryResolver());
+
+        // Copy all non-specialized methods
+        for (EntryTreeNode<EntryMapping> node : source) {
+            if (!(node.getEntry() instanceof MethodEntry) || !bridgeMethodIndex.isSpecializedMethod((MethodEntry) node.getEntry())) {
+                result.insert(node.getEntry(), node.getValue());
+            }
+        }
+
+        // Add correct mappings for specialized methods
+        for (Map.Entry<MethodEntry, MethodEntry> entry : bridgeMethodIndex.getBridgeToSpecialized().entrySet()) {
+            MethodEntry bridge = entry.getKey();
+            MethodEntry specialized = entry.getValue();
+            String name = translator.translate(bridge).getName();
+            result.insert(specialized, new EntryMapping(name));
+        }
+
+        Path output = Paths.get(args[4]);
+        Utils.delete(output);
+        MappingCommandsUtil.write(result, args[3], output, saveParameters);
+    }
+}