Use name proposal service when exporting deobfuscated jar or decompiled sources (#159)

2 files changed, 61 insertions, 8 deletions

diff --git a/src/main/java/cuchaz/enigma/EnigmaProject.java b/src/main/java/cuchaz/enigma/EnigmaProject.java
index bb2f7faa..a589c638 100644
--- a/src/main/java/cuchaz/enigma/EnigmaProject.java
+++ b/src/main/java/cuchaz/enigma/EnigmaProject.java
@@ -8,12 +8,11 @@ import com.strobel.decompiler.languages.java.ast.CompilationUnit;
 import cuchaz.enigma.analysis.ClassCache;
 import cuchaz.enigma.analysis.EntryReference;
 import cuchaz.enigma.analysis.index.JarIndex;
+import cuchaz.enigma.api.service.NameProposalService;
 import cuchaz.enigma.bytecode.translators.SourceFixVisitor;
 import cuchaz.enigma.bytecode.translators.TranslationClassVisitor;
 import cuchaz.enigma.translation.Translator;
-import cuchaz.enigma.translation.mapping.EntryMapping;
-import cuchaz.enigma.translation.mapping.EntryRemapper;
-import cuchaz.enigma.translation.mapping.MappingsChecker;
+import cuchaz.enigma.translation.mapping.*;
 import cuchaz.enigma.translation.mapping.tree.DeltaTrackingTree;
 import cuchaz.enigma.translation.mapping.tree.EntryTree;
 import cuchaz.enigma.translation.representation.entry.ClassEntry;
@@ -26,16 +25,19 @@ import org.objectweb.asm.tree.ClassNode;
 
 import java.io.BufferedWriter;
 import java.io.IOException;
+import java.io.PrintWriter;
 import java.io.StringWriter;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Collection;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.jar.JarEntry;
 import java.util.jar.JarOutputStream;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 public class EnigmaProject {
         private final Enigma enigma;
@@ -141,7 +143,12 @@ public class EnigmaProject {
 
         public JarExport exportRemappedJar(ProgressListener progress) {
                 Collection<ClassEntry> classEntries = jarIndex.getEntryIndex().getClasses();
-                Translator deobfuscator = mapper.getDeobfuscator();
+
+                Translator deobfuscator = getEnigma()
+                                .getServices()
+                                .get(NameProposalService.TYPE)
+                                .map(nameProposalService -> (Translator) new ProposingTranslator(mapper, nameProposalService))
+                                .orElse(mapper.getDeobfuscator());
 
                 AtomicInteger count = new AtomicInteger();
                 progress.init(classEntries.size(), "Deobfuscating classes...");
@@ -233,11 +240,14 @@ public class EnigmaProject {
                 }
 
                 private String decompileClass(ClassNode translatedNode, SourceProvider sourceProvider) {
-                        CompilationUnit sourceTree = sourceProvider.getSources(translatedNode.name);
-
                         StringWriter writer = new StringWriter();
-                        sourceProvider.writeSource(writer, sourceTree);
-
+                        try {
+                                CompilationUnit sourceTree = sourceProvider.getSources(translatedNode.name);
+                                sourceProvider.writeSource(writer, sourceTree);
+                        } catch (Throwable t) {
+                                t.printStackTrace();
+                                t.printStackTrace(new PrintWriter(writer));
+                        }
                         return writer.toString();
                 }
         }
diff --git a/src/main/java/cuchaz/enigma/ProposingTranslator.java b/src/main/java/cuchaz/enigma/ProposingTranslator.java
new file mode 100644
index 00000000..0807880e
--- /dev/null
+++ b/src/main/java/cuchaz/enigma/ProposingTranslator.java
@@ -0,0 +1,43 @@
+package cuchaz.enigma;
+
+import cuchaz.enigma.api.service.NameProposalService;
+import cuchaz.enigma.translation.Translatable;
+import cuchaz.enigma.translation.Translator;
+import cuchaz.enigma.translation.mapping.*;
+import cuchaz.enigma.translation.representation.entry.Entry;
+
+import java.util.Optional;
+
+public class ProposingTranslator implements Translator {
+    private final EntryRemapper mapper;
+    private final NameProposalService nameProposalService;
+
+    public ProposingTranslator(EntryRemapper mapper, NameProposalService nameProposalService) {
+        this.mapper = mapper;
+        this.nameProposalService = nameProposalService;
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public <T extends Translatable> T translate(T translatable) {
+        if (translatable == null) {
+            return null;
+        }
+
+        T deobfuscated = mapper.deobfuscate(translatable);
+
+        if (translatable instanceof Entry && ((Entry) deobfuscated).getName().equals(((Entry<?>) translatable).getName())) {
+            return mapper.getObfResolver()
+                    .resolveEntry((Entry<?>) translatable, ResolutionStrategy.RESOLVE_ROOT)
+                    .stream()
+                    .map(e1 -> nameProposalService.proposeName(e1, mapper))
+                    .filter(Optional::isPresent)
+                    .map(Optional::get)
+                    .findFirst()
+                    .map(newName -> (T) ((Entry) deobfuscated).withName(newName))
+                    .orElse(deobfuscated);
+        }
+
+        return deobfuscated;
+    }
+}