Experiment with only recompiling changed classes

build.gradle

@@ -179,6 +179,11 @@ idea {
                     programParameters = "run --dist joined --neoforge net.neoforged:neoforge:21.0.0-beta:userdev --write-result=compiled:build/minecraft.jar --write-result=clientResources:build/client-extra.jar --write-result=sources:build/minecraft-sources.jar"
                     moduleRef(project, sourceSets.main)
                 }
+                "Experiments"(Application) {
+                    mainClass = mainClassName
+                    programParameters = "run --verbose --partial-recompile --dist joined --neoforge net.neoforged:neoforge:21.0.0-beta:userdev --interface-injection-data=interfaces_test.json --validated-access-transformer=test_at.cfg --write-result=node.sourcesWithNeoForge.output.output:build/sources-with-neoforge.jar"
+                    moduleRef(project, sourceSets.main)
+                }
                 "Run Neoform 1.21 (joined)"(Application) {
                     mainClass = mainClassName
                     programParameters = "run --dist joined --neoform net.neoforged:neoform:1.21-20240613.152323@zip --write-result=compiled:build/minecraft.jar --write-result=clientResources:build/client-extra.jar --write-result=sources:build/minecraft-sources.jar"

interfaces_test.json

@@ -0,0 +1,5 @@
+{
+  "net/minecraft/world/item/Item": [
+    "testproject/FunExtensions"
+  ]
+}

src/main/java/net/neoforged/neoform/runtime/actions/ApplySourceTransformAction.java

@@ -109,12 +109,12 @@ public void run(ProcessingEnvironment environment) throws IOException, Interrupt
 
             for (var path : validatedAccessTransformers) {
                 args.add("--access-transformer");
-                args.add(environment.getPathArgument(path));
+                args.add(environment.getPathArgument(path.toAbsolutePath()));
             }
 
             for (var path : additionalAccessTransformers) {
                 args.add("--access-transformer");
-                args.add(environment.getPathArgument(path));
+                args.add(environment.getPathArgument(path.toAbsolutePath()));
             }
         }
 

src/main/java/net/neoforged/neoform/runtime/actions/InjectModifiedClasses.java

@@ -0,0 +1,37 @@
+package net.neoforged.neoform.runtime.actions;
+
+import net.neoforged.neoform.runtime.engine.ProcessingEnvironment;
+
+import java.io.BufferedInputStream;
+import java.io.IOException;
+import java.nio.file.FileSystems;
+import java.nio.file.Files;
+import java.util.Map;
+import java.util.zip.ZipInputStream;
+
+public class InjectModifiedClasses extends BuiltInAction {
+    @Override
+    public void run(ProcessingEnvironment environment) throws IOException, InterruptedException {
+        var classesFile = environment.getRequiredInputPath("classes");
+        var extraClassesFile = environment.getRequiredInputPath("classes2");
+        var output = environment.getOutputPath("output");
+
+        // Copy the largest jar then use NIO to insert extra entries.
+        // This is faster than working with ZipFile streams which inflate and deflate all the entries.
+        Files.copy(classesFile, output);
+        try (var zfs = FileSystems.newFileSystem(output, Map.of("create", false))) {
+            var zfsRoot = zfs.getPath("/");
+
+            // Copy the extra classes to the output zip
+            try (var in = new ZipInputStream(new BufferedInputStream(Files.newInputStream(extraClassesFile)))) {
+                for (var entry = in.getNextEntry(); entry != null; entry = in.getNextEntry()) {
+                    if (entry.isDirectory()) {
+                        continue;
+                    }
+                    var targetPath = zfsRoot.resolve(entry.getName());
+                    Files.copy(in, targetPath);
+                }
+            }
+        }
+    }
+}

src/main/java/net/neoforged/neoform/runtime/actions/RecompileSourcesAction.java

@@ -1,6 +1,5 @@
 package net.neoforged.neoform.runtime.actions;
 
-import net.neoforged.neoform.runtime.artifacts.ClasspathItem;
 import net.neoforged.neoform.runtime.cache.CacheKeyBuilder;
 import net.neoforged.neoform.runtime.engine.ProcessingEnvironment;
 import net.neoforged.neoform.runtime.graph.ExecutionNodeAction;
@@ -62,4 +61,6 @@ public int getTargetJavaVersion() {
     public void setTargetJavaVersion(int targetJavaVersion) {
         this.targetJavaVersion = targetJavaVersion;
     }
+
+    public abstract RecompileSourcesAction copy();
 }

src/main/java/net/neoforged/neoform/runtime/actions/RecompileSourcesActionWithECJ.java

@@ -212,6 +212,17 @@ public void computeCacheKey(CacheKeyBuilder ck) {
         ck.add("compiler type", "eclipse");
     }
 
+    @Override
+    public RecompileSourcesAction copy() {
+        var ret = new RecompileSourcesActionWithECJ();
+        ret.getClasspath().setOverriddenClasspath(getClasspath().getOverriddenClasspath());
+        ret.getClasspath().setAdditionalClasspath(getClasspath().getAdditionalClasspath());
+        ret.getSourcepath().setOverriddenClasspath(getSourcepath().getOverriddenClasspath());
+        ret.getSourcepath().setAdditionalClasspath(getSourcepath().getAdditionalClasspath());
+        ret.setTargetJavaVersion(getTargetJavaVersion());
+        return ret;
+    }
+
     static class ECJFilesystem extends FileSystem {
         protected ECJFilesystem(Classpath[] paths, String[] initialFileNames, boolean annotationsFromClasspath) {
             super(paths, initialFileNames, annotationsFromClasspath);

src/main/java/net/neoforged/neoform/runtime/actions/RecompileSourcesActionWithJDK.java

@@ -108,4 +108,15 @@ private List<String> getCompilerOptions() {
                 "-implicit:none" // Prevents source files from the source-path from being emitted
         );
     }
+
+    @Override
+    public RecompileSourcesAction copy() {
+        var ret = new RecompileSourcesActionWithJDK();
+        ret.getClasspath().setOverriddenClasspath(getClasspath().getOverriddenClasspath());
+        ret.getClasspath().setAdditionalClasspath(getClasspath().getAdditionalClasspath());
+        ret.getSourcepath().setOverriddenClasspath(getSourcepath().getOverriddenClasspath());
+        ret.getSourcepath().setAdditionalClasspath(getSourcepath().getAdditionalClasspath());
+        ret.setTargetJavaVersion(getTargetJavaVersion());
+        return ret;
+    }
 }

src/main/java/net/neoforged/neoform/runtime/actions/SelectSourcesToRecompile.java

@@ -0,0 +1,100 @@
+package net.neoforged.neoform.runtime.actions;
+
+import net.neoforged.neoform.runtime.engine.ProcessingEnvironment;
+
+import java.io.IOException;
+import java.nio.file.FileSystems;
+import java.nio.file.Files;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipFile;
+import java.util.zip.ZipOutputStream;
+
+public class SelectSourcesToRecompile extends BuiltInAction {
+    @Override
+    public void run(ProcessingEnvironment environment) throws IOException, InterruptedException {
+        var originalSources = environment.getRequiredInputPath("originalSources");
+        var originalClasses = environment.getRequiredInputPath("originalClasses");
+        var transformedSources = environment.getRequiredInputPath("transformedSources");
+
+        var unchangedClasses = environment.getOutputPath("unchangedClasses");
+        var changedSourcesOnly = environment.getOutputPath("changedSourcesOnly");
+
+        // Read all original sources to memory
+        var originalSourcesContents = new HashMap<String, byte[]>();
+        try (var zf = new ZipFile(originalSources.toFile())) {
+            var entries = zf.entries();
+            while (entries.hasMoreElements()) {
+                var entry = entries.nextElement();
+                if (entry.isDirectory()) {
+                    continue;
+                }
+                try (var is = zf.getInputStream(entry)) {
+                    originalSourcesContents.put(entry.getName(), is.readAllBytes());
+                }
+            }
+        }
+
+        // Copy unchanged sources to the output
+        var changedSourcePaths = new HashSet<String>();
+        try (var os = Files.newOutputStream(changedSourcesOnly);
+             var zos = new ZipOutputStream(os)) {
+
+            try (var transformedZip = new ZipFile(transformedSources.toFile())) {
+                var entries = transformedZip.entries();
+                while (entries.hasMoreElements()) {
+                    var entry = entries.nextElement();
+                    if (entry.isDirectory()) {
+                        continue;
+                    }
+
+                    try (var is = transformedZip.getInputStream(entry)) {
+                        var data = is.readAllBytes();
+
+                        if (!Arrays.equals(data, originalSourcesContents.get(entry.getName()))) {
+                            changedSourcePaths.add(entry.getName());
+                            // Copy to output
+                            var copiedEntry = new ZipEntry(entry.getName());
+                            copiedEntry.setMethod(entry.getMethod());
+                            zos.putNextEntry(copiedEntry);
+                            zos.write(data);
+                            zos.closeEntry();
+                        }
+                    }
+                }
+            }
+        }
+
+        // Copy unchanged classes, by first copying the whole zip then deleting unwanted entries using NIO.
+        // This is faster than working with ZipFile streams which inflate and deflate all the entries.
+        Files.copy(originalClasses, unchangedClasses);
+        try (var zfs = FileSystems.newFileSystem(unchangedClasses, Map.of("create", "false"))) {
+            try (var originalZip = new ZipFile(originalClasses.toFile())) {
+                var entries = originalZip.entries();
+                while (entries.hasMoreElements()) {
+                    var entry = entries.nextElement();
+                    if (entry.isDirectory()) {
+                        continue;
+                    }
+                    // Remove trailing .class
+                    var sourceName = entry.getName();
+                    if (sourceName.endsWith(".class")) {
+                        sourceName = sourceName.substring(0, sourceName.length() - 6);
+                    }
+                    // Remove inner class suffix
+                    sourceName = sourceName.split("\\$")[0];
+                    // Add trailing .java
+                    sourceName += ".java";
+
+                    if (changedSourcePaths.contains(sourceName)) {
+                        // Delete!
+                        Files.delete(zfs.getPath(entry.getName()));
+                    }
+                }
+            }
+        }
+    }
+}

src/main/java/net/neoforged/neoform/runtime/cli/RunNeoFormCommand.java

@@ -5,12 +5,13 @@
 import net.neoforged.neoform.runtime.actions.InjectFromZipFileSource;
 import net.neoforged.neoform.runtime.actions.InjectZipContentAction;
 import net.neoforged.neoform.runtime.actions.MergeWithSourcesAction;
+import net.neoforged.neoform.runtime.actions.InjectModifiedClasses;
 import net.neoforged.neoform.runtime.actions.PatchActionFactory;
 import net.neoforged.neoform.runtime.actions.RecompileSourcesAction;
+import net.neoforged.neoform.runtime.actions.SelectSourcesToRecompile;
 import net.neoforged.neoform.runtime.actions.StripManifestDigestContentFilter;
 import net.neoforged.neoform.runtime.artifacts.ClasspathItem;
 import net.neoforged.neoform.runtime.config.neoforge.NeoForgeConfig;
-import net.neoforged.neoform.runtime.config.neoform.NeoFormDistConfig;
 import net.neoforged.neoform.runtime.engine.NeoFormEngine;
 import net.neoforged.neoform.runtime.graph.ExecutionGraph;
 import net.neoforged.neoform.runtime.graph.ExecutionNode;
@@ -82,6 +83,9 @@ static class SourceArtifacts {
         String neoforge;
     }
 
+    @CommandLine.Option(names = "--partial-recompile", description = "[EXPERIMENTAL] Enables partial recompilation for user-defined transforms")
+    boolean partialRecompile;
+
     @Override
     protected void runWithNeoFormEngine(NeoFormEngine engine, List<AutoCloseable> closables) throws IOException, InterruptedException {
         var artifactManager = engine.getArtifactManager();
@@ -218,8 +222,6 @@ protected void runWithNeoFormEngine(NeoFormEngine engine, List<AutoCloseable> cl
             engine.loadNeoFormData(neoFormDataPath, dist);
         }
 
-        applyAdditionalAccessTransformers(engine);
-
         if (parchmentData != null) {
             var parchmentDataFile = artifactManager.get(parchmentData);
             Consumer<ApplySourceTransformAction> jstConsumer = transformSources -> {
@@ -237,36 +239,102 @@ protected void runWithNeoFormEngine(NeoFormEngine engine, List<AutoCloseable> cl
             }
         }
 
-        if (!interfaceInjectionDataFiles.isEmpty()) {
-            var transformNode = getOrAddTransformSourcesNode(engine);
-            ((ApplySourceTransformAction) transformNode.action()).setInjectedInterfaces(interfaceInjectionDataFiles);
-
-            // Add the stub source jar to the recomp classpath
-            engine.applyTransform(new ModifyAction<>(
-                    "recompile",
-                    RecompileSourcesAction.class,
-                    action -> {
-                        action.getSourcepath().add(ClasspathItem.of(transformNode.getRequiredOutput("stubs")));
-                    }
-            ));
-        }
+        applyAdditionalTransforms(engine);
 
         execute(engine);
     }
 
     /**
-     * Configure the engine to apply additional user-supplied access transformers to the game sources.
+     * Configure the engine to apply additional user-supplied access transformers and interfaces to the game sources.
+     * This is done by re-transforming the sources a second time, and only recompiling changed sources.
      */
-    private void applyAdditionalAccessTransformers(NeoFormEngine engine) {
+    private void applyAdditionalTransforms(NeoFormEngine engine) {
+        if (additionalAccessTransformers.isEmpty() && validatedAccessTransformers.isEmpty() && interfaceInjectionDataFiles.isEmpty()) {
+            return;
+        }
+
+        var graph = engine.getGraph();
+
+        ExecutionNode sourceTransformNode;
+        ApplySourceTransformAction sourceTransformAction;
+        String recompilationNodeName;
+
+        if (partialRecompile) {
+            var recompileOutput = graph.getRequiredOutput("recompile", "output");
+            var recompileNode = recompileOutput.getNode();
+            var recompileAction = (RecompileSourcesAction) recompileNode.action();
+            var startingSources = recompileNode.getRequiredInput("sources").parentOutput();
+
+            // Second transform action for additional transforms
+            var applyAdditionalTransformsBuilder = graph.nodeBuilder("applyAdditionalTransforms");
+            sourceTransform(engine, action -> {})
+                    .make(applyAdditionalTransformsBuilder, startingSources);
+            sourceTransformNode = applyAdditionalTransformsBuilder.build();
+            sourceTransformAction = (ApplySourceTransformAction) sourceTransformNode.action();
+            var transformedSources = sourceTransformNode.getRequiredOutput("output");
+
+            // Split off sources that were modified by the second transform
+            var selectSourcesBuilder = graph.nodeBuilder("selectSourcesToRecompile");
+            selectSourcesBuilder.input("originalSources", startingSources.asInput());
+            selectSourcesBuilder.input("originalClasses", recompileOutput.asInput());
+            selectSourcesBuilder.input("transformedSources", transformedSources.asInput());
+            var unchangedClasses = selectSourcesBuilder.output("unchangedClasses", NodeOutputType.JAR, "Classes that were already compiled and whose corresponding sources did not change.");
+            var changedSourcesOnly = selectSourcesBuilder.output("changedSourcesOnly", NodeOutputType.JAR, "Sources that were changed and need to be recompiled.");
+            selectSourcesBuilder.action(new SelectSourcesToRecompile());
+            var selectSources = selectSourcesBuilder.build();
+
+            // Recompile modified sources
+            var recompileModifiedSources = graph.nodeBuilder("recompileModifiedSources");
+            recompileModifiedSources.input("sources", changedSourcesOnly.asInput());
+            recompileModifiedSources.input("versionManifest", recompileOutput.getNode().getRequiredInput("versionManifest"));
+            var modifiedClasses = recompileModifiedSources.output("output", NodeOutputType.JAR, "Compiled minecraft sources that were changed.");
+            RecompileSourcesAction recompileModifiedAction = recompileAction.copy();
+            recompileModifiedAction.getClasspath().add(ClasspathItem.of(unchangedClasses));
+            recompileModifiedSources.action(recompileModifiedAction);
+            recompileModifiedSources.build();
+
+            // Recombine with the unmodified classes that were already recompiled before
+            var injectBuilder = graph.nodeBuilder("injectModifiedClasses");
+            injectBuilder.input("classes", unchangedClasses.asInput());
+            injectBuilder.input("classes2", modifiedClasses.asInput());
+            var injectedOutput = injectBuilder.output("output", NodeOutputType.JAR, "Compiled Minecraft sources with additional changes merged in.");
+            injectBuilder.action(new InjectModifiedClasses());
+            injectBuilder.build();
+
+            // Since we replace one node by many, we cannot use the ReplaceNodeOutput with the usual factory,
+            // but rather directly call this helper
+            ReplaceNodeOutput.replaceOutput(graph, recompileOutput, injectedOutput, List.of(selectSources));
+            // We also need to replace usages of the sources
+            ReplaceNodeOutput.replaceOutput(graph, startingSources, transformedSources, List.of(recompileNode, sourceTransformNode, selectSources));
+
+            recompilationNodeName = "recompileModifiedSources";
+        } else {
+            sourceTransformNode = getOrAddTransformSourcesNode(engine);
+            sourceTransformAction = (ApplySourceTransformAction) sourceTransformNode.action();
+            recompilationNodeName = "recompile";
+        }
+
         if (!additionalAccessTransformers.isEmpty() || !validatedAccessTransformers.isEmpty()) {
-            var transformSources = getOrAddTransformSourcesAction(engine);
-            transformSources.setAdditionalAccessTransformers(additionalAccessTransformers.stream().map(Paths::get).toList());
-            transformSources.setValidatedAccessTransformers(validatedAccessTransformers.stream().map(Paths::get).toList());
+            sourceTransformAction.setAdditionalAccessTransformers(additionalAccessTransformers.stream().map(Paths::get).toList());
+            sourceTransformAction.setValidatedAccessTransformers(validatedAccessTransformers.stream().map(Paths::get).toList());
 
             if (validateAccessTransformers) {
-                transformSources.addArg("--access-transformer-validation=error");
+                sourceTransformAction.addArg("--access-transformer-validation=error");
             }
         }
+
+        if (!interfaceInjectionDataFiles.isEmpty()) {
+            sourceTransformAction.setInjectedInterfaces(interfaceInjectionDataFiles);
+
+            // Add the stub source jar to the recomp classpath
+            engine.applyTransform(new ModifyAction<>(
+                    recompilationNodeName,
+                    RecompileSourcesAction.class,
+                    action -> {
+                        action.getSourcepath().add(ClasspathItem.of(sourceTransformNode.getRequiredOutput("stubs")));
+                    }
+            ));
+        }
     }
 
     private static NodeOutput createCompiledWithNeoForge(NeoFormEngine engine, ZipFile neoforgeClassesZip) {

src/main/java/net/neoforged/neoform/runtime/graph/NodeInput.java

@@ -37,6 +37,8 @@ void setNode(ExecutionNode node) {
 
     public abstract <T> T getValue(ResultRepresentation<T> representation) throws IOException;
 
+    public abstract NodeOutput parentOutput();
+
     static final class NodeInputForOutput extends NodeInput {
         private NodeOutput output;
 
@@ -66,5 +68,10 @@ public void collectCacheKeyComponent(CacheKeyBuilder builder) {
         public <T> T getValue(ResultRepresentation<T> representation) throws IOException {
             return output.getResultRepresentation(representation);
         }
+
+        @Override
+        public NodeOutput parentOutput() {
+            return output;
+        }
     }
 }