Implement interface injections

README.md

@@ -33,7 +33,10 @@ Usage: jst [-hV] [--in-format=<inputFormat>] [--libraries-list=<librariesList>]
            [--enable-parchment --parchment-mappings=<mappingsPath> [--[no-]parchment-javadoc]
            [--parchment-conflict-prefix=<conflictPrefix>]] [--enable-accesstransformers
            --access-transformer=<atFiles> [--access-transformer=<atFiles>]...
-           [--access-transformer-validation=<validation>]] INPUT OUTPUT
+           [--access-transformer-validation=<validation>]] [--enable-interface-injection
+           --interface-injection-stub-location=<stubOut>
+           [--interface-injection-marker=<annotationMarker>]
+           [--interface-injection-data=<paths>]...] INPUT OUTPUT
       INPUT                Path to a single Java-file, a source-archive or a folder containing the
                              source to transform.
       OUTPUT               Path to where the resulting source should be placed.
@@ -74,6 +77,15 @@ Plugin - accesstransformers
                            The level of validation to use for ats
       --enable-accesstransformers
                            Enable accesstransformers
+Plugin - interface-injection
+      --enable-interface-injection
+                           Enable interface-injection
+      --interface-injection-data=<paths>
+                           The paths to read interface injection JSON files from
+      --interface-injection-marker=<annotationMarker>
+                           The name of an annotation to use as a marker for injected interfaces
+      --interface-injection-stub-location=<stubOut>
+                           The path to save interface stubs to
 ```
 
 ## Licenses

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

@@ -2,6 +2,7 @@
 
 import com.intellij.core.CoreApplicationEnvironment;
 import com.intellij.core.JavaCoreProjectEnvironment;
+import com.intellij.psi.JavaPsiFacade;
 import com.intellij.psi.PsiManager;
 
 public interface IntelliJEnvironment {
@@ -10,4 +11,6 @@ public interface IntelliJEnvironment {
     JavaCoreProjectEnvironment getProjectEnv();
 
     PsiManager getPsiManager();
+
+    JavaPsiFacade getPsiFacade();
 }

cli/build.gradle

@@ -20,6 +20,7 @@ dependencies {
     implementation "info.picocli:picocli:$picocli_version"
     implementation project(":parchment")
     implementation project(":accesstransformers")
+    implementation project(':interfaceinjection')
     implementation 'org.slf4j:slf4j-simple:2.0.13'
 
     testImplementation platform("org.junit:junit-bom:$junit_version")

cli/src/main/java/net/neoforged/jst/cli/intellij/IntelliJEnvironmentImpl.java

@@ -19,6 +19,7 @@
 import com.intellij.pom.java.InternalPersistentJavaLanguageLevelReaderService;
 import com.intellij.pom.java.LanguageLevel;
 import com.intellij.psi.JavaModuleSystem;
+import com.intellij.psi.JavaPsiFacade;
 import com.intellij.psi.PsiElementFinder;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiFileFactory;
@@ -50,6 +51,7 @@ public class IntelliJEnvironmentImpl implements IntelliJEnvironment, AutoCloseab
     private final MockProject project;
     private final JavaCoreProjectEnvironment javaEnv;
     private final PsiManager psiManager;
+    private final JavaPsiFacade psiFacade;
 
     public IntelliJEnvironmentImpl(Logger logger) throws IOException {
         this.logger = logger;
@@ -80,13 +82,19 @@ protected VirtualFileSystem createJrtFileSystem() {
         LanguageLevelProjectExtension.getInstance(project).setLanguageLevel(LanguageLevel.JDK_17);
 
         psiManager = PsiManager.getInstance(project);
+        psiFacade = JavaPsiFacade.getInstance(project);
     }
 
     @Override
     public PsiManager getPsiManager() {
         return psiManager;
     }
 
+    @Override
+    public JavaPsiFacade getPsiFacade() {
+        return psiFacade;
+    }
+
     @Override
     public CoreApplicationEnvironment getAppEnv() {
         return javaEnv.getEnvironment();

gradle.properties

@@ -4,3 +4,4 @@ jetbrains_annotations_version=24.1.0
 picocli_version=4.7.6
 junit_version=5.10.3
 assertj_version=3.26.0
+gson_version=2.10.1

interfaceinjection/build.gradle

@@ -0,0 +1,8 @@
+plugins {
+    id 'java-library'
+}
+
+dependencies {
+    implementation project(':api')
+    implementation "com.google.code.gson:gson:${project.gson_version}"
+}

interfaceinjection/src/main/java/net/neoforged/jst/interfaceinjection/InjectInterfacesVisitor.java

@@ -0,0 +1,84 @@
+package net.neoforged.jst.interfaceinjection;
+
+import com.intellij.psi.PsiClass;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
+import com.intellij.psi.PsiRecursiveElementVisitor;
+import com.intellij.psi.PsiWhiteSpace;
+import com.intellij.psi.util.ClassUtil;
+import com.intellij.util.containers.MultiMap;
+import net.neoforged.jst.api.Replacements;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.Collection;
+import java.util.stream.Collectors;
+
+class InjectInterfacesVisitor extends PsiRecursiveElementVisitor {
+    private final Replacements replacements;
+    private final MultiMap<String, String> interfaces;
+    private final StubStore stubs;
+
+    @Nullable
+    private final String marker;
+
+    InjectInterfacesVisitor(Replacements replacements, MultiMap<String, String> interfaces, StubStore stubs, String marker) {
+        this.replacements = replacements;
+        this.interfaces = interfaces;
+        this.stubs = stubs;
+        this.marker = marker;
+    }
+
+    @Override
+    public void visitElement(@NotNull PsiElement element) {
+        if (element instanceof PsiClass psiClass) {
+            if (psiClass.getQualifiedName() == null) return;
+            String className = ClassUtil.getJVMClassName(psiClass);
+            inject(psiClass, interfaces.get(className.replace('.', '/')));
+
+            for (PsiClass innerClass : psiClass.getInnerClasses()) {
+                visitElement(innerClass);
+            }
+        }
+    }
+
+    @Override
+    public void visitFile(@NotNull PsiFile file) {
+        file.acceptChildren(this);
+    }
+
+    private void inject(PsiClass psiClass, Collection<String> targets) {
+        // We cannot add implements clauses to anonymous or unnamed classes
+        if (targets.isEmpty() || psiClass.getImplementsList() == null) return;
+
+        var interfaceImplementation = targets.stream()
+                .distinct().map(stubs::createStub)
+                .map(this::decorate)
+                .collect(Collectors.joining(", "));
+
+        var implementsList = psiClass.isInterface() ? psiClass.getExtendsList() : psiClass.getImplementsList();
+        if (implementsList.getChildren().length == 0) {
+            StringBuilder text = new StringBuilder();
+
+            // `public class Cls{}` is valid, but we cannot inject the implements exactly next to the class name, so we need
+            // to make sure that we have spacing
+            if (!(psiClass.getLBrace().getPrevSibling() instanceof PsiWhiteSpace)) {
+                text.append(' ');
+            }
+            text.append(psiClass.isInterface() ? "extends" : "implements").append(' ');
+            text.append(interfaceImplementation);
+            text.append(' ');
+
+            replacements.insertBefore(psiClass.getLBrace(), text.toString());
+        } else {
+            replacements.insertAfter(implementsList.getLastChild(), ", " + interfaceImplementation);
+        }
+    }
+
+    private String decorate(String iface) {
+        if (marker == null) {
+            return iface;
+        }
+        return "@" + marker + " " + iface;
+    }
+}

interfaceinjection/src/main/java/net/neoforged/jst/interfaceinjection/InterfaceInjectionPlugin.java

@@ -0,0 +1,21 @@
+package net.neoforged.jst.interfaceinjection;
+
+import net.neoforged.jst.api.SourceTransformer;
+import net.neoforged.jst.api.SourceTransformerPlugin;
+
+/**
+ * Plugin that injects stub interfaces to classes.
+ * <p>
+ * Mods can use interface injection to have compile-time access to the interfaces they add to classes via Mixins.
+ */
+public class InterfaceInjectionPlugin implements SourceTransformerPlugin {
+    @Override
+    public String getName() {
+        return "interface-injection";
+    }
+
+    @Override
+    public SourceTransformer createTransformer() {
+        return new InterfaceInjectionTransformer();
+    }
+}

interfaceinjection/src/main/java/net/neoforged/jst/interfaceinjection/InterfaceInjectionTransformer.java

@@ -0,0 +1,77 @@
+package net.neoforged.jst.interfaceinjection;
+
+import com.google.gson.Gson;
+import com.google.gson.JsonObject;
+import com.intellij.psi.PsiFile;
+import com.intellij.util.containers.MultiMap;
+import net.neoforged.jst.api.Replacements;
+import net.neoforged.jst.api.SourceTransformer;
+import net.neoforged.jst.api.TransformContext;
+import picocli.CommandLine;
+
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+
+public class InterfaceInjectionTransformer implements SourceTransformer {
+    private static final Gson GSON = new Gson();
+
+    @CommandLine.Option(names = "--interface-injection-stub-location", required = true, description = "The path to save interface stubs to")
+    public Path stubOut;
+
+    @CommandLine.Option(names = "--interface-injection-marker", description = "The name of an annotation to use as a marker for injected interfaces")
+    public String annotationMarker;
+
+    @CommandLine.Option(names = "--interface-injection-data", description = "The paths to read interface injection JSON files from")
+    public List<Path> paths;
+
+    private MultiMap<String, String> interfaces;
+    private StubStore stubs;
+    private String marker;
+
+    @Override
+    public void beforeRun(TransformContext context) {
+        interfaces = new MultiMap<>();
+        stubs = new StubStore(context.logger(), context.environment().getPsiFacade());
+
+        if (annotationMarker != null) {
+            marker = annotationMarker.replace('/', '.').replace('$', '.');
+        }
+
+        for (Path path : paths) {
+            try {
+                var json = GSON.fromJson(Files.readString(path), JsonObject.class);
+                for (String clazz : json.keySet()) {
+                    var entry = json.get(clazz);
+                    if (entry.isJsonArray()) {
+                        entry.getAsJsonArray().forEach(el -> interfaces.putValue(clazz, el.getAsString()));
+                    } else {
+                        interfaces.putValue(clazz, entry.getAsString());
+                    }
+                }
+            } catch (IOException exception) {
+                context.logger().error("Failed to read interface injection data file: %s", exception.getMessage());
+                throw new UncheckedIOException(exception);
+            }
+        }
+    }
+
+    @Override
+    public boolean afterRun(TransformContext context) {
+        try {
+            stubs.save(stubOut);
+        } catch (IOException e) {
+            context.logger().error("Failed to save stubs: %s", e.getMessage());
+            throw new UncheckedIOException(e);
+        }
+
+        return true;
+    }
+
+    @Override
+    public void visitFile(PsiFile psiFile, Replacements replacements) {
+        new InjectInterfacesVisitor(replacements, interfaces, stubs, marker).visitFile(psiFile);
+    }
+}