Add API for reacting to collected replacements (#40)

Author: lukebemish

api/src/main/java/net/neoforged/jst/api/Replacements.java

@@ -4,10 +4,19 @@
 import com.intellij.psi.PsiElement;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 public final class Replacements {
-    private final List<Replacement> replacements = new ArrayList<>();
+    private final List<Replacement> replacements;
+
+    public Replacements(List<Replacement> replacements) {
+        this.replacements = replacements;
+    }
+    
+    public Replacements() {
+        this(new ArrayList<>());
+    }
 
     public boolean isEmpty() {
         return replacements.isEmpty();
@@ -81,5 +90,4 @@ public String apply(CharSequence originalContent) {
         writer.append(originalContent, replacements.get(replacements.size() - 1).range().getEndOffset(), originalContent.length());
         return writer.toString();
     }
-
 }

api/src/main/java/net/neoforged/jst/api/SourceTransformer.java

@@ -2,6 +2,8 @@
 
 import com.intellij.psi.PsiFile;
 
+import java.util.List;
+
 /**
  * Transformers are created through {@link SourceTransformerPlugin plugins}, and handle source replacements.
  * <p>
@@ -19,13 +21,25 @@ public interface SourceTransformer {
     default void beforeRun(TransformContext context) {
     }
 
+    /**
+     * Invoke after replacements are collected for a given file, but before they are applied.
+     * <p>
+     * Can be used to react to or verify the replacements that were collected.
+     * @param fileEntry the file entry being transformed
+     * @param replacements the replacements that were collected; read-only
+     * @return {@code true} if the transformation should continue, {@code false} if it should fail
+     */
+    default boolean beforeReplacement(FileEntry fileEntry, List<Replacement> replacements) {
+        return true;
+    }
+
     /**
      * Invoked after all source transformations are finished.
      * <p>
      * Can be used for post-transformation validation.
      *
      * @param context the transform context
-     * @return {@code true} if the transformation was successful, {@code false} otherwise
+     * @return {@code true} if the transformation was successful, {@code false} if it failed
      */
     default boolean afterRun(TransformContext context) {
         return true;

cli/src/main/java/net/neoforged/jst/cli/Replacement.java

@@ -6,6 +6,7 @@
 import net.neoforged.jst.api.FileSink;
 import net.neoforged.jst.api.FileSource;
 import net.neoforged.jst.api.Logger;
+import net.neoforged.jst.api.Replacement;
 import net.neoforged.jst.api.Replacements;
 import net.neoforged.jst.api.SourceTransformer;
 import net.neoforged.jst.api.TransformContext;
@@ -19,7 +20,9 @@
 import java.nio.file.attribute.FileTime;
 import java.time.Instant;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * Reference for out-of-IDE usage of the IntelliJ Java parser is from the Kotlin compiler
@@ -63,16 +66,22 @@ public boolean process(FileSource source, FileSink sink, List<SourceTransformer>
                 });
             }
         } else {
+            var success = new AtomicBoolean(true);
             try (var asyncOut = new OrderedParallelWorkQueue(sink, maxQueueDepth);
                  var stream = source.streamEntries()) {
                 stream.forEach(entry -> asyncOut.submitAsync(parallelSink -> {
                     try {
-                        processEntry(entry, sourceRoot, transformers, parallelSink);
+                        if (!processEntry(entry, sourceRoot, transformers, parallelSink)) {
+                            success.set(false);
+                        }
                     } catch (IOException e) {
                         throw new UncheckedIOException(e);
                     }
                 }));
             }
+            if (!success.get()) {
+                return false;
+            }
         }
 
         boolean isOk = true;
@@ -83,25 +92,31 @@ public boolean process(FileSource source, FileSink sink, List<SourceTransformer>
         return isOk;
     }
 
-    private void processEntry(FileEntry entry, VirtualFile sourceRoot, List<SourceTransformer> transformers, FileSink sink) throws IOException {
+    private boolean processEntry(FileEntry entry, VirtualFile sourceRoot, List<SourceTransformer> transformers, FileSink sink) throws IOException {
         if (entry.directory()) {
             sink.putDirectory(entry.relativePath());
-            return;
+            return true;
         }
+        
+        boolean[] success = {true};
 
         try (var in = entry.openInputStream()) {
             byte[] content = in.readAllBytes();
             var lastModified = entry.lastModified();
 
             if (!isIgnored(entry.relativePath()) && !transformers.isEmpty() && entry.hasExtension("java")) {
                 var orgContent = content;
-                content = transformSource(sourceRoot, entry.relativePath(), transformers, content);
+                content = transformSource(sourceRoot, entry, transformers, content, success);
+                if (!success[0]) {
+                    return false;
+                }
                 if (orgContent != content) {
                     lastModified = FileTime.from(Instant.now());
                 }
             }
             sink.putFile(entry.relativePath(), lastModified, content);
         }
+        return true;
     }
 
     private boolean isIgnored(String relativePath) {
@@ -113,11 +128,12 @@ private boolean isIgnored(String relativePath) {
         return false;
     }
 
-    byte[] transformSource(VirtualFile contentRoot, String path, List<SourceTransformer> transformers, byte[] originalContentBytes) {
+    private byte[] transformSource(VirtualFile contentRoot, FileEntry entry, List<SourceTransformer> transformers, byte[] originalContentBytes, boolean[] successOut) {
         // Instead of parsing the content we actually read from the file, we read the virtual file that is
         // visible to IntelliJ from adding the source jar. The reasoning is that IntelliJ will cache this internally
         // and reuse it when cross-referencing type-references. If we parsed from a String instead, it would parse
         // the same file twice.
+        var path = entry.relativePath();
         var sourceFile = contentRoot.findFileByRelativePath(path);
         if (sourceFile == null) {
             System.err.println("Can't transform " + path + " since IntelliJ doesn't see it in the source jar.");
@@ -130,14 +146,23 @@ byte[] transformSource(VirtualFile contentRoot, String path, List<SourceTransfor
         }
 
         // Gather replaced ranges in the source-file with their replacement
-        var replacements = new Replacements();
+        List<Replacement> replacementsList = new ArrayList<>();
+        var replacements = new Replacements(replacementsList);
 
         for (var transformer : transformers) {
             transformer.visitFile(psiFile, replacements);
         }
 
+        var readOnlyReplacements = Collections.unmodifiableList(replacementsList);
+        boolean success = true;
+        for (var transformer : transformers) {
+            success = success && transformer.beforeReplacement(entry, readOnlyReplacements);
+        }
+        
+        successOut[0] = success;
+
         // If no replacements were made, just stream the original content into the destination file
-        if (replacements.isEmpty()) {
+        if (!success || replacements.isEmpty()) {
             return originalContentBytes;
         }