Add a way to transform the decompiler input

6 files changed, 74 insertions, 7 deletions

diff --git a/enigma-cli/src/main/java/cuchaz/enigma/command/DecompileCommand.java b/enigma-cli/src/main/java/cuchaz/enigma/command/DecompileCommand.java
index 020bd979..baed29a6 100644
--- a/enigma-cli/src/main/java/cuchaz/enigma/command/DecompileCommand.java
+++ b/enigma-cli/src/main/java/cuchaz/enigma/command/DecompileCommand.java
@@ -47,7 +47,7 @@ public class DecompileCommand extends Command {
                 ProgressListener progress = new ConsoleProgressListener();
 
                 EnigmaProject.JarExport jar = project.exportRemappedJar(progress);
-                EnigmaProject.SourceExport source = jar.decompile(progress, decompilerService, DecompileErrorStrategy.TRACE_AS_SOURCE);
+                EnigmaProject.SourceExport source = jar.decompile(project, progress, decompilerService, DecompileErrorStrategy.TRACE_AS_SOURCE);
 
                 source.write(fileJarOut, progress);
         }
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 1c4549d6..7691f39c 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
@@ -311,7 +311,7 @@ public class GuiController implements ClientPacketHandler, GuiView, DataInvalida
 
                 return ProgressDialog.runOffThread(this.gui.getFrame(), progress -> {
                         EnigmaProject.JarExport jar = project.exportRemappedJar(progress);
-                        jar.decompileStream(progress, chp.getDecompilerService(), EnigmaProject.DecompileErrorStrategy.TRACE_AS_SOURCE).forEach(source -> {
+                        jar.decompileStream(project, progress, chp.getDecompilerService(), EnigmaProject.DecompileErrorStrategy.TRACE_AS_SOURCE).forEach(source -> {
                                 try {
                                         source.writeTo(source.resolvePath(path));
                                 } catch (IOException e) {
diff --git a/enigma/src/main/java/cuchaz/enigma/EnigmaProject.java b/enigma/src/main/java/cuchaz/enigma/EnigmaProject.java
index 3f4a0673..4a528a7b 100644
--- a/enigma/src/main/java/cuchaz/enigma/EnigmaProject.java
+++ b/enigma/src/main/java/cuchaz/enigma/EnigmaProject.java
@@ -27,6 +27,7 @@ import cuchaz.enigma.analysis.EntryReference;
 import cuchaz.enigma.analysis.index.JarIndex;
 import cuchaz.enigma.api.DataInvalidationEvent;
 import cuchaz.enigma.api.DataInvalidationListener;
+import cuchaz.enigma.api.service.DecompilerInputTransformerService;
 import cuchaz.enigma.api.service.NameProposalService;
 import cuchaz.enigma.api.service.ObfuscationTestService;
 import cuchaz.enigma.api.view.ProjectView;
@@ -277,13 +278,20 @@ public class EnigmaProject implements ProjectView {
                         }
                 }
 
-                public SourceExport decompile(ProgressListener progress, DecompilerService decompilerService, DecompileErrorStrategy errorStrategy) {
-                        List<ClassSource> decompiled = this.decompileStream(progress, decompilerService, errorStrategy).toList();
+                public SourceExport decompile(EnigmaProject project, ProgressListener progress, DecompilerService decompilerService, DecompileErrorStrategy errorStrategy) {
+                        List<ClassSource> decompiled = this.decompileStream(project, progress, decompilerService, errorStrategy).toList();
                         return new SourceExport(decompiled);
                 }
 
-                public Stream<ClassSource> decompileStream(ProgressListener progress, DecompilerService decompilerService, DecompileErrorStrategy errorStrategy) {
-                        Collection<ClassNode> classes = this.compiled.values().stream().filter(classNode -> classNode.name.indexOf('$') == -1).toList();
+                public Stream<ClassSource> decompileStream(EnigmaProject project, ProgressListener progress, DecompilerService decompilerService, DecompileErrorStrategy errorStrategy) {
+                        Collection<ClassNode> classes = this.compiled.values().stream()
+                                        .filter(classNode -> classNode.name.indexOf('$') == -1)
+                                        .peek(classNode -> {
+                                                for (DecompilerInputTransformerService transformer : project.enigma.getServices().get(DecompilerInputTransformerService.TYPE)) {
+                                                        transformer.transform(classNode);
+                                                }
+                                        })
+                                        .toList();
 
                         progress.init(classes.size(), I18n.translate("progress.classes.decompiling"));
 
diff --git a/enigma/src/main/java/cuchaz/enigma/api/service/DecompilerInputTransformerService.java b/enigma/src/main/java/cuchaz/enigma/api/service/DecompilerInputTransformerService.java
new file mode 100644
index 00000000..a560515a
--- /dev/null
+++ b/enigma/src/main/java/cuchaz/enigma/api/service/DecompilerInputTransformerService.java
@@ -0,0 +1,9 @@
+package cuchaz.enigma.api.service;
+
+import org.objectweb.asm.tree.ClassNode;
+
+public interface DecompilerInputTransformerService extends EnigmaService {
+        EnigmaServiceType<DecompilerInputTransformerService> TYPE = EnigmaServiceType.create("decompiler_input_transformer");
+
+        void transform(ClassNode classNode);
+}
diff --git a/enigma/src/main/java/cuchaz/enigma/classhandle/ClassHandleProvider.java b/enigma/src/main/java/cuchaz/enigma/classhandle/ClassHandleProvider.java
index cb80dda9..36c10cc1 100644
--- a/enigma/src/main/java/cuchaz/enigma/classhandle/ClassHandleProvider.java
+++ b/enigma/src/main/java/cuchaz/enigma/classhandle/ClassHandleProvider.java
@@ -24,6 +24,7 @@ import org.jetbrains.annotations.Nullable;
 
 import cuchaz.enigma.EnigmaProject;
 import cuchaz.enigma.classprovider.CachingClassProvider;
+import cuchaz.enigma.classprovider.DecompilerInputTransformingClassProvider;
 import cuchaz.enigma.classprovider.ObfuscationFixClassProvider;
 import cuchaz.enigma.events.ClassHandleListener;
 import cuchaz.enigma.events.ClassHandleListener.InvalidationType;
@@ -104,7 +105,13 @@ public final class ClassHandleProvider {
         }
 
         private Decompiler createDecompiler() {
-                return ds.create(new CachingClassProvider(new ObfuscationFixClassProvider(project.getClassProvider(), project.getJarIndex())), new SourceSettings(true, true));
+                return ds.create(
+                                new DecompilerInputTransformingClassProvider(
+                                                new CachingClassProvider(new ObfuscationFixClassProvider(project.getClassProvider(), project.getJarIndex())),
+                                                project.getEnigma().getServices()
+                                ),
+                                new SourceSettings(true, true)
+                );
         }
 
         /**
diff --git a/enigma/src/main/java/cuchaz/enigma/classprovider/DecompilerInputTransformingClassProvider.java b/enigma/src/main/java/cuchaz/enigma/classprovider/DecompilerInputTransformingClassProvider.java
new file mode 100644
index 00000000..960221dd
--- /dev/null
+++ b/enigma/src/main/java/cuchaz/enigma/classprovider/DecompilerInputTransformingClassProvider.java
@@ -0,0 +1,43 @@
+package cuchaz.enigma.classprovider;
+
+import java.util.Collection;
+
+import org.jetbrains.annotations.Nullable;
+import org.objectweb.asm.tree.ClassNode;
+
+import cuchaz.enigma.EnigmaServices;
+import cuchaz.enigma.api.service.DecompilerInputTransformerService;
+
+public class DecompilerInputTransformingClassProvider implements ClassProvider {
+        private final ClassProvider delegate;
+        private final EnigmaServices services;
+
+        public DecompilerInputTransformingClassProvider(ClassProvider delegate, EnigmaServices services) {
+                this.delegate = delegate;
+                this.services = services;
+        }
+
+        @Override
+        public Collection<String> getClassNames() {
+                return delegate.getClassNames();
+        }
+
+        @Override
+        @Nullable
+        public ClassNode get(String name) {
+                ClassNode classNode = delegate.get(name);
+
+                if (classNode == null) {
+                        return null;
+                }
+
+                // copy the class, so that the input class isn't modified (which could lead to the class being retransformed if
+                // it's cached)
+                ClassNode classCopy = new ClassNode();
+                classNode.accept(classCopy);
+
+                services.get(DecompilerInputTransformerService.TYPE).forEach(transformer -> transformer.transform(classCopy));
+
+                return classCopy;
+        }
+}