Allow plugins to be able to translate entries from deobfuscated to obfuscated as well

10 files changed, 198 insertions, 15 deletions

diff --git a/enigma-server/src/main/java/cuchaz/enigma/network/packet/EntryChangeC2SPacket.java b/enigma-server/src/main/java/cuchaz/enigma/network/packet/EntryChangeC2SPacket.java
index 6a7ffe99..f45ac997 100644
--- a/enigma-server/src/main/java/cuchaz/enigma/network/packet/EntryChangeC2SPacket.java
+++ b/enigma-server/src/main/java/cuchaz/enigma/network/packet/EntryChangeC2SPacket.java
@@ -39,7 +39,7 @@ public class EntryChangeC2SPacket implements Packet<ServerPacketHandler> {
                 boolean valid = handler.getServer().canModifyEntry(handler.getClient(), this.change.getTarget());
 
                 if (valid) {
-                        EntryUtil.applyChange(vc, handler.getServer().getMappings(), this.change);
+                        EntryUtil.applyChange(vc, null, handler.getServer().getMappings(), this.change);
                         valid = vc.canProceed();
                 }
 
diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
index 7691f39c..5ea9ab1c 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
@@ -583,7 +583,7 @@ public class GuiController implements ClientPacketHandler, GuiView, DataInvalida
 
                 Entry<?> target = change.getTarget();
                 EntryMapping prev = this.project.getMapper().getDeobfMapping(target);
-                EntryMapping mapping = EntryUtil.applyChange(vc, this.project.getMapper(), change);
+                EntryMapping mapping = EntryUtil.applyChange(vc, this.project, this.project.getMapper(), change);
 
                 boolean renamed = !change.getDeobfName().isUnchanged();
 
diff --git a/enigma/src/main/java/cuchaz/enigma/EnigmaProject.java b/enigma/src/main/java/cuchaz/enigma/EnigmaProject.java
index 4a528a7b..86a70372 100644
--- a/enigma/src/main/java/cuchaz/enigma/EnigmaProject.java
+++ b/enigma/src/main/java/cuchaz/enigma/EnigmaProject.java
@@ -38,9 +38,11 @@ import cuchaz.enigma.classprovider.ObfuscationFixClassProvider;
 import cuchaz.enigma.source.Decompiler;
 import cuchaz.enigma.source.DecompilerService;
 import cuchaz.enigma.source.SourceSettings;
+import cuchaz.enigma.translation.ObfuscatingTranslator;
 import cuchaz.enigma.translation.ProposingTranslator;
 import cuchaz.enigma.translation.Translatable;
 import cuchaz.enigma.translation.Translator;
+import cuchaz.enigma.translation.mapping.EntryChange;
 import cuchaz.enigma.translation.mapping.EntryMapping;
 import cuchaz.enigma.translation.mapping.EntryRemapper;
 import cuchaz.enigma.translation.mapping.MappingsChecker;
@@ -63,6 +65,9 @@ public class EnigmaProject implements ProjectView {
         private final Set<String> projectClasses;
 
         private EntryRemapper mapper;
+        private Translator proposingTranslator;
+        @Nullable
+        private ObfuscatingTranslator inverseTranslator;
 
         private final List<DataInvalidationListener> dataInvalidationListeners = new ArrayList<>();
 
@@ -78,7 +83,7 @@ public class EnigmaProject implements ProjectView {
                 this.jarChecksum = jarChecksum;
                 this.projectClasses = projectClasses;
 
-                this.mapper = EntryRemapper.empty(jarIndex);
+                setMappings(null);
         }
 
         public void setMappings(EntryTree<EntryMapping> mappings) {
@@ -87,6 +92,13 @@ public class EnigmaProject implements ProjectView {
                 } else {
                         mapper = EntryRemapper.empty(jarIndex);
                 }
+
+                NameProposalService[] nameProposalServices = enigma.getServices().get(NameProposalService.TYPE).toArray(new NameProposalService[0]);
+                proposingTranslator = nameProposalServices.length == 0 ? mapper.getDeobfuscator() : new ProposingTranslator(mapper, nameProposalServices);
+
+                if (inverseTranslator != null) {
+                        inverseTranslator.refreshAll(proposingTranslator);
+                }
         }
 
         public Enigma getEnigma() {
@@ -121,6 +133,10 @@ public class EnigmaProject implements ProjectView {
                 for (Entry<?> entry : dropped) {
                         mappings.trackChange(entry);
                 }
+
+                if (inverseTranslator != null) {
+                        inverseTranslator.refreshAll(proposingTranslator);
+                }
         }
 
         private Collection<Entry<?>> dropMappings(EntryTree<EntryMapping> mappings, ProgressListener progress) {
@@ -224,21 +240,18 @@ public class EnigmaProject implements ProjectView {
                 Collection<ClassEntry> classEntries = jarIndex.getEntryIndex().getClasses();
                 ClassProvider fixingClassProvider = new ObfuscationFixClassProvider(classProvider, jarIndex);
 
-                NameProposalService[] nameProposalServices = getEnigma().getServices().get(NameProposalService.TYPE).toArray(new NameProposalService[0]);
-                Translator deobfuscator = nameProposalServices.length == 0 ? mapper.getDeobfuscator() : new ProposingTranslator(mapper, nameProposalServices);
-
                 AtomicInteger count = new AtomicInteger();
                 progress.init(classEntries.size(), I18n.translate("progress.classes.deobfuscating"));
 
                 Map<String, ClassNode> compiled = classEntries.parallelStream().map(entry -> {
-                        ClassEntry translatedEntry = deobfuscator.translate(entry);
+                        ClassEntry translatedEntry = proposingTranslator.translate(entry);
                         progress.step(count.getAndIncrement(), translatedEntry.toString());
 
                         ClassNode node = fixingClassProvider.get(entry.getFullName());
 
                         if (node != null) {
                                 ClassNode translatedNode = new ClassNode();
-                                node.accept(new TranslationClassVisitor(deobfuscator, Enigma.ASM_VERSION, translatedNode));
+                                node.accept(new TranslationClassVisitor(proposingTranslator, Enigma.ASM_VERSION, translatedNode));
                                 return translatedNode;
                         }
 
@@ -348,7 +361,37 @@ public class EnigmaProject implements ProjectView {
         @Override
         @SuppressWarnings("unchecked")
         public <T extends EntryView> T deobfuscate(T entry) {
-                return (T) mapper.extendedDeobfuscate((Translatable) entry).getValue();
+                return (T) proposingTranslator.extendedTranslate((Translatable) entry).getValue();
+        }
+
+        @Override
+        @SuppressWarnings("unchecked")
+        public <T extends EntryView> T obfuscate(T entry) {
+                if (inverseTranslator == null) {
+                        throw new IllegalStateException("Must call registerForInverseMappings before calling obfuscate");
+                }
+
+                return (T) inverseTranslator.extendedTranslate((Entry<?>) entry).getValue();
+        }
+
+        @Override
+        public void registerForInverseMappings() {
+                if (inverseTranslator == null) {
+                        inverseTranslator = new ObfuscatingTranslator(jarIndex);
+                        inverseTranslator.refreshAll(proposingTranslator);
+                }
+        }
+
+        public void onEntryChange(EntryMapping prevMapping, EntryChange<?> change) {
+                if (inverseTranslator == null || change.getDeobfName().isUnchanged()) {
+                        return;
+                }
+
+                String newName = change.getDeobfName().isSet() ? change.getDeobfName().getNewValue() : proposingTranslator.extendedTranslate(change.getTarget()).getValue().getName();
+
+                for (Entry<?> equivalentEntry : mapper.getObfResolver().resolveEquivalentEntries(change.getTarget())) {
+                        inverseTranslator.refreshName(equivalentEntry, prevMapping.targetName(), newName);
+                }
         }
 
         @Override
diff --git a/enigma/src/main/java/cuchaz/enigma/api/view/ProjectView.java b/enigma/src/main/java/cuchaz/enigma/api/view/ProjectView.java
index b8760e29..8f53d118 100644
--- a/enigma/src/main/java/cuchaz/enigma/api/view/ProjectView.java
+++ b/enigma/src/main/java/cuchaz/enigma/api/view/ProjectView.java
@@ -13,6 +13,13 @@ import cuchaz.enigma.api.view.entry.EntryView;
 public interface ProjectView {
         <T extends EntryView> T deobfuscate(T entry);
 
+        /**
+         * Must call {@link #registerForInverseMappings()} before using this method.
+         */
+        <T extends EntryView> T obfuscate(T entry);
+
+        void registerForInverseMappings();
+
         Collection<String> getProjectClasses();
 
         @Nullable
diff --git a/enigma/src/main/java/cuchaz/enigma/translation/ObfuscatingTranslator.java b/enigma/src/main/java/cuchaz/enigma/translation/ObfuscatingTranslator.java
new file mode 100644
index 00000000..27220e0c
--- /dev/null
+++ b/enigma/src/main/java/cuchaz/enigma/translation/ObfuscatingTranslator.java
@@ -0,0 +1,115 @@
+package cuchaz.enigma.translation;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Set;
+
+import org.jetbrains.annotations.Nullable;
+
+import cuchaz.enigma.analysis.index.JarIndex;
+import cuchaz.enigma.translation.mapping.EntryMapping;
+import cuchaz.enigma.translation.mapping.EntryResolver;
+import cuchaz.enigma.translation.mapping.ResolutionStrategy;
+import cuchaz.enigma.translation.mapping.tree.EntryTree;
+import cuchaz.enigma.translation.mapping.tree.HashEntryTree;
+import cuchaz.enigma.translation.representation.entry.ClassEntry;
+import cuchaz.enigma.translation.representation.entry.Entry;
+import cuchaz.enigma.translation.representation.entry.FieldEntry;
+import cuchaz.enigma.translation.representation.entry.MethodEntry;
+import cuchaz.enigma.translation.representation.entry.ParentedEntry;
+
+public class ObfuscatingTranslator implements Translator {
+        private final JarIndex jarIndex;
+
+        private final EntryTree<EntryMapping> inverseMappings = new HashEntryTree<>();
+        private final EntryResolver resolver = new ObfuscatingResolver();
+
+        public ObfuscatingTranslator(JarIndex jarIndex) {
+                this.jarIndex = jarIndex;
+        }
+
+        @Override
+        @Nullable
+        @SuppressWarnings("unchecked")
+        public <T extends Translatable> TranslateResult<T> extendedTranslate(@Nullable T translatable) {
+                if (translatable == null) {
+                        return null;
+                }
+
+                if (translatable instanceof FieldEntry || translatable instanceof MethodEntry) {
+                        ParentedEntry<ClassEntry> key = obfOwnerAndDesc((ParentedEntry<ClassEntry>) translatable);
+                        EntryMapping mapping = inverseMappings.get(key);
+                        return mapping == null ? TranslateResult.obfuscated((T) key) : TranslateResult.deobfuscated((T) key.withName(mapping.targetName()));
+                }
+
+                return (TranslateResult<T>) translatable.extendedTranslate(this, resolver, inverseMappings);
+        }
+
+        public void refreshAll(Translator deobfuscator) {
+                inverseMappings.clear();
+
+                for (ClassEntry clazz : jarIndex.getEntryIndex().getClasses()) {
+                        inverseMappings.insert(deobfuscator.extendedTranslate(clazz).getValue(), new EntryMapping(clazz.getName()));
+                }
+
+                for (FieldEntry field : jarIndex.getEntryIndex().getFields()) {
+                        inverseMappings.insert(obfOwnerAndDesc(deobfuscator.extendedTranslate(field).getValue()), new EntryMapping(field.getName()));
+                }
+
+                for (MethodEntry method : jarIndex.getEntryIndex().getMethods()) {
+                        inverseMappings.insert(obfOwnerAndDesc(deobfuscator.extendedTranslate(method).getValue()), new EntryMapping(method.getName()));
+                }
+        }
+
+        public void refreshName(Entry<?> entry, String oldDeobfName, String newDeobfName) {
+                inverseMappings.remove(entry.withName(oldDeobfName));
+                inverseMappings.insert(entry.withName(newDeobfName), new EntryMapping(entry.getName()));
+        }
+
+        @SuppressWarnings("unchecked")
+        private <T extends ParentedEntry<ClassEntry>> T obfOwnerAndDesc(T translatable) {
+                if (translatable.getParent() != null) {
+                        translatable = (T) translatable.withParent(extendedTranslate(translatable.getParent()).getValue());
+                }
+
+                if (translatable instanceof FieldEntry field) {
+                        translatable = (T) field.withDesc(extendedTranslate(field.getDesc()).getValue());
+                } else if (translatable instanceof MethodEntry method) {
+                        translatable = (T) method.withDesc(extendedTranslate(method.getDesc()).getValue());
+                }
+
+                return translatable;
+        }
+
+        private class ObfuscatingResolver implements EntryResolver {
+                @Override
+                @SuppressWarnings("unchecked")
+                public <E extends Entry<?>> Collection<E> resolveEntry(E entry, ResolutionStrategy strategy) {
+                        ClassEntry containingClass;
+
+                        if (entry instanceof ClassEntry || (containingClass = entry.findAncestor(ClassEntry.class)) == null) {
+                                return List.of(entry);
+                        }
+
+                        List<E> result = new ArrayList<>();
+                        result.add(entry);
+
+                        for (ClassEntry parentClass : jarIndex.getInheritanceIndex().getAncestors(containingClass)) {
+                                result.add((E) entry.replaceAncestor(containingClass, parentClass));
+                        }
+
+                        return result;
+                }
+
+                @Override
+                public Set<Entry<?>> resolveEquivalentEntries(Entry<?> entry) {
+                        return Set.of(entry);
+                }
+
+                @Override
+                public Set<MethodEntry> resolveEquivalentMethods(MethodEntry methodEntry) {
+                        return Set.of(methodEntry);
+                }
+        }
+}
diff --git a/enigma/src/main/java/cuchaz/enigma/translation/Translator.java b/enigma/src/main/java/cuchaz/enigma/translation/Translator.java
index 56889908..593b4efa 100644
--- a/enigma/src/main/java/cuchaz/enigma/translation/Translator.java
+++ b/enigma/src/main/java/cuchaz/enigma/translation/Translator.java
@@ -18,10 +18,12 @@ import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 
+import org.jetbrains.annotations.Contract;
 import org.jetbrains.annotations.Nullable;
 
 public interface Translator {
         @Nullable
+        @Contract("null -> null; !null -> !null")
         <T extends Translatable> TranslateResult<T> extendedTranslate(@Nullable T translatable);
 
         @Deprecated
diff --git a/enigma/src/main/java/cuchaz/enigma/translation/mapping/EntryRemapper.java b/enigma/src/main/java/cuchaz/enigma/translation/mapping/EntryRemapper.java
index 12ed7b1d..dea2d151 100644
--- a/enigma/src/main/java/cuchaz/enigma/translation/mapping/EntryRemapper.java
+++ b/enigma/src/main/java/cuchaz/enigma/translation/mapping/EntryRemapper.java
@@ -54,11 +54,11 @@ public class EntryRemapper {
                 doPutMapping(vc, obfuscatedEntry, deobfMapping, true);
         }
 
-        public void putMapping(ValidationContext vc, Entry<?> obfuscatedEntry, @NotNull EntryMapping deobfMapping) {
-                doPutMapping(vc, obfuscatedEntry, deobfMapping, false);
+        public boolean putMapping(ValidationContext vc, Entry<?> obfuscatedEntry, @NotNull EntryMapping deobfMapping) {
+                return doPutMapping(vc, obfuscatedEntry, deobfMapping, false);
         }
 
-        private void doPutMapping(ValidationContext vc, Entry<?> obfuscatedEntry, @NotNull EntryMapping deobfMapping, boolean validateOnly) {
+        private boolean doPutMapping(ValidationContext vc, Entry<?> obfuscatedEntry, @NotNull EntryMapping deobfMapping, boolean validateOnly) {
                 if (obfuscatedEntry instanceof FieldEntry) {
                         FieldEntry fieldEntry = (FieldEntry) obfuscatedEntry;
                         ClassEntry classEntry = fieldEntry.getParent();
@@ -77,7 +77,7 @@ public class EntryRemapper {
                 }
 
                 if (validateOnly || !vc.canProceed()) {
-                        return;
+                        return false;
                 }
 
                 for (Entry<?> resolvedEntry : resolvedEntries) {
@@ -87,6 +87,8 @@ public class EntryRemapper {
                                 obfToDeobf.insert(resolvedEntry, deobfMapping);
                         }
                 }
+
+                return true;
         }
 
         // A little bit of a hack to also map the getter method for record fields.
diff --git a/enigma/src/main/java/cuchaz/enigma/translation/mapping/EntryUtil.java b/enigma/src/main/java/cuchaz/enigma/translation/mapping/EntryUtil.java
index fb8ab631..9e6641be 100644
--- a/enigma/src/main/java/cuchaz/enigma/translation/mapping/EntryUtil.java
+++ b/enigma/src/main/java/cuchaz/enigma/translation/mapping/EntryUtil.java
@@ -1,17 +1,23 @@
 package cuchaz.enigma.translation.mapping;
 
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
+import cuchaz.enigma.EnigmaProject;
 import cuchaz.enigma.translation.representation.entry.Entry;
 import cuchaz.enigma.utils.validation.ValidationContext;
 
 public class EntryUtil {
-        public static EntryMapping applyChange(ValidationContext vc, EntryRemapper remapper, EntryChange<?> change) {
+        public static EntryMapping applyChange(ValidationContext vc, @Nullable EnigmaProject project, EntryRemapper remapper, EntryChange<?> change) {
                 Entry<?> target = change.getTarget();
                 EntryMapping prev = remapper.getDeobfMapping(target);
                 EntryMapping mapping = EntryUtil.applyChange(prev, change);
 
-                remapper.putMapping(vc, target, mapping);
+                if (remapper.putMapping(vc, target, mapping)) {
+                        if (project != null) {
+                                project.onEntryChange(prev, change);
+                        }
+                }
 
                 return mapping;
         }
diff --git a/enigma/src/main/java/cuchaz/enigma/translation/representation/entry/FieldEntry.java b/enigma/src/main/java/cuchaz/enigma/translation/representation/entry/FieldEntry.java
index dba86448..90c256c1 100644
--- a/enigma/src/main/java/cuchaz/enigma/translation/representation/entry/FieldEntry.java
+++ b/enigma/src/main/java/cuchaz/enigma/translation/representation/entry/FieldEntry.java
@@ -58,6 +58,10 @@ public class FieldEntry extends ParentedEntry<ClassEntry> implements Comparable<
                 return new FieldEntry(parent, name, desc, null);
         }
 
+        public FieldEntry withDesc(TypeDescriptor desc) {
+                return new FieldEntry(parent, name, desc, null);
+        }
+
         @Override
         public FieldEntry withParent(ClassEntry parent) {
                 return new FieldEntry(parent, this.name, this.desc, null);
diff --git a/enigma/src/main/java/cuchaz/enigma/translation/representation/entry/MethodEntry.java b/enigma/src/main/java/cuchaz/enigma/translation/representation/entry/MethodEntry.java
index d86f1564..da6751e1 100644
--- a/enigma/src/main/java/cuchaz/enigma/translation/representation/entry/MethodEntry.java
+++ b/enigma/src/main/java/cuchaz/enigma/translation/representation/entry/MethodEntry.java
@@ -69,6 +69,10 @@ public class MethodEntry extends ParentedEntry<ClassEntry> implements Comparable
                 return new MethodEntry(parent, name, descriptor, javadocs);
         }
 
+        public MethodEntry withDesc(MethodDescriptor descriptor) {
+                return new MethodEntry(parent, name, descriptor, javadocs);
+        }
+
         @Override
         public MethodEntry withParent(ClassEntry parent) {
                 return new MethodEntry(new ClassEntry(parent.getFullName()), name, descriptor, javadocs);