Generics widening

cli/build.gradle

@@ -35,6 +35,7 @@ dependencies {
     include project(":accesstransformers")
     include project(':interfaceinjection')
     include project(':unpick')
+    include project(':genericswidening')
 
     testImplementation platform("org.junit:junit-bom:$junit_version")
     testImplementation 'org.junit.jupiter:junit-jupiter'

genericswidening/build.gradle

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

genericswidening/src/main/java/net/neoforged/jst/genericswidening/GenericsTarget.java

@@ -0,0 +1,3 @@
+package net.neoforged.jst.genericswidening;
+
+record GenericsTarget(String methodName, String methodDescriptor, String genericParameterType, int genericParameterOrdinal) {}

genericswidening/src/main/java/net/neoforged/jst/genericswidening/GenericsWideningPlugin.java

@@ -0,0 +1,16 @@
+package net.neoforged.jst.genericswidening;
+
+import net.neoforged.jst.api.SourceTransformer;
+import net.neoforged.jst.api.SourceTransformerPlugin;
+
+public final class GenericsWideningPlugin implements SourceTransformerPlugin {
+    @Override
+    public String getName() {
+        return "generics-widening";
+    }
+
+    @Override
+    public SourceTransformer createTransformer() {
+        return new GenericsWideningTransformer();
+    }
+}

genericswidening/src/main/java/net/neoforged/jst/genericswidening/GenericsWideningTransformer.java

@@ -0,0 +1,88 @@
+package net.neoforged.jst.genericswidening;
+
+import com.google.gson.Gson;
+import com.google.gson.JsonElement;
+import com.google.gson.JsonObject;
+import com.intellij.psi.PsiFile;
+import com.intellij.util.Consumer;
+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.ArrayList;
+import java.util.List;
+
+final class GenericsWideningTransformer implements SourceTransformer {
+    private static final Gson GSON = new Gson();
+
+    @CommandLine.Option(names = "--generics-widening-data", description = "The paths to read generics widening JSON files from")
+    public List<Path> paths = new ArrayList<>();
+
+    private MultiMap<String, GenericsTarget> methods = MultiMap.empty();
+
+    @Override
+    public void beforeRun(TransformContext context) {
+        methods = MultiMap.createLinkedSet();
+        for (Path path : paths) {
+            try {
+                JsonObject json = GSON.fromJson(Files.readString(path), JsonObject.class);
+                for (String clazz : json.keySet()) {
+                    parseTarget(json.get(clazz), target -> methods.putValue(clazz, target));
+                }
+            } catch (IOException exception) {
+                context.logger().error("Failed to read generics widening data file: %s", exception.getMessage());
+                throw new UncheckedIOException(exception);
+            }
+        }
+    }
+
+    private static void parseTarget(JsonElement element, Consumer<GenericsTarget> appender) throws IOException {
+        if (element.isJsonArray()) {
+            for (JsonElement elem : element.getAsJsonArray()) {
+                parseTarget(elem, appender);
+            }
+        } else if (!element.isJsonObject()) {
+            throw new IOException("Invalid target entry: " + element);
+        }
+
+        JsonObject obj = element.getAsJsonObject();
+        String methodName = getRequiredString(obj, "methodName");
+        String methodDesc = getRequiredString(obj, "methodDescriptor");
+        String paramType = getRequiredString(obj, "parameterType");
+        int paramOrdinal = -1;
+        if (obj.has("parameterOrdinal")) {
+            JsonElement ordinalElem = obj.get("parameterOrdinal");
+            if (!ordinalElem.isJsonPrimitive()) {
+                throw new IOException("Invalid type for field: parameterOrdinal");
+            }
+            paramOrdinal = ordinalElem.getAsInt();
+        }
+        if (!methodDesc.contains(paramType)) {
+            throw new IOException("Descriptor '" + methodDesc + "' does not contain generic parameter '" + paramType + "'");
+        }
+        // TODO: throw for multiple occurrences when no ordinal is provided
+        appender.accept(new GenericsTarget(methodName, methodDesc, paramType, paramOrdinal));
+    }
+
+    private static String getRequiredString(JsonObject obj, String name) throws IOException {
+        if (!obj.has(name)) {
+            throw new IOException("Missing required field: " + name);
+        }
+        JsonElement element = obj.get(name);
+        if (!element.isJsonPrimitive()) {
+            throw new IOException("Invalid type for field: " + name);
+        }
+        return element.getAsString();
+    }
+
+    @Override
+    public void visitFile(PsiFile psiFile, Replacements replacements) {
+        new WidenGenericsVisitor(replacements, methods).visitFile(psiFile);
+    }
+}

genericswidening/src/main/java/net/neoforged/jst/genericswidening/WidenGenericsVisitor.java

@@ -0,0 +1,107 @@
+package net.neoforged.jst.genericswidening;
+
+import com.intellij.psi.PsiClass;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
+import com.intellij.psi.PsiJavaCodeReferenceElement;
+import com.intellij.psi.PsiMethod;
+import com.intellij.psi.PsiParameter;
+import com.intellij.psi.PsiParameterList;
+import com.intellij.psi.PsiRecursiveElementVisitor;
+import com.intellij.psi.PsiReferenceParameterList;
+import com.intellij.psi.PsiTypeElement;
+import com.intellij.psi.util.ClassUtil;
+import com.intellij.util.containers.MultiMap;
+import net.neoforged.jst.api.PsiHelper;
+import net.neoforged.jst.api.Replacements;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.Collection;
+
+final class WidenGenericsVisitor extends PsiRecursiveElementVisitor {
+    private final Replacements replacements;
+    private final MultiMap<String, GenericsTarget> methods;
+
+    WidenGenericsVisitor(Replacements replacements, MultiMap<String, GenericsTarget> methods) {
+        this.replacements = replacements;
+        this.methods = methods;
+    }
+
+    @Override
+    public void visitElement(@NotNull PsiElement element) {
+        if (element instanceof PsiClass psiClass) {
+            if (psiClass.getQualifiedName() == null) {
+                return;
+            }
+
+            String className = ClassUtil.getJVMClassName(psiClass);
+            widenGenerics(psiClass, methods.get(className.replace('.', '/')));
+
+            for (PsiClass innerClass : psiClass.getInnerClasses()) {
+                visitElement(innerClass);
+            }
+        }
+    }
+
+    @Override
+    public void visitFile(@NotNull PsiFile file) {
+        file.acceptChildren(this);
+    }
+
+    private void widenGenerics(PsiClass psiClass, Collection<GenericsTarget> targets) {
+        if (targets.isEmpty()) {
+            return;
+        }
+
+        for (GenericsTarget target : targets) {
+            for (PsiMethod method : psiClass.findMethodsByName(target.methodName(), false)) {
+                if (!PsiHelper.getBinaryMethodSignature(method).equals(target.methodDescriptor())) {
+                    continue;
+                }
+
+                PsiParameterList paramList = method.getParameterList();
+                PsiParameter param;
+                if (target.genericParameterOrdinal() >= 0) {
+                    param = paramList.getParameter(target.genericParameterOrdinal());
+                } else {
+                    param = findParamByType(paramList, target.genericParameterType());
+                }
+                if (param == null) {
+                    continue;
+                }
+                widenGeneric(param);
+            }
+        }
+    }
+
+    private static PsiParameter findParamByType(PsiParameterList paramList, String type) {
+        for (int i = 0; i < paramList.getParametersCount(); i++) {
+            PsiParameter paramCandidate = paramList.getParameter(i);
+            if (ClassUtil.getBinaryPresentation(paramCandidate.getType()).equals(type)) {
+                return paramCandidate;
+            }
+        }
+        return null;
+    }
+
+    private void widenGeneric(PsiParameter param) {
+        PsiTypeElement typeElement = findChildOfType(param, PsiTypeElement.class);
+        PsiJavaCodeReferenceElement refElement = findChildOfType(typeElement, PsiJavaCodeReferenceElement.class);
+        PsiReferenceParameterList refParamList = findChildOfType(refElement, PsiReferenceParameterList.class);
+        for (PsiElement child : refParamList.getChildren()) {
+            if (child instanceof PsiTypeElement paramTypeElement) {
+                replacements.insertBefore(paramTypeElement, "? extends ");
+            }
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    private static <T> T findChildOfType(PsiElement element, Class<T> type) {
+        for (PsiElement child : element.getChildren()) {
+            if (type.isInstance(child)) {
+                return (T) child;
+            }
+        }
+        return null;
+    }
+}

genericswidening/src/main/resources/META-INF/services/net.neoforged.jst.api.SourceTransformerPlugin

@@ -0,0 +1 @@
+net.neoforged.jst.genericswidening.GenericsWideningPlugin

settings.gradle

@@ -42,4 +42,5 @@ include 'tests'
 include 'accesstransformers'
 include 'interfaceinjection'
 include 'unpick'
+include 'genericswidening'
 include 'sourcetest'

tests/data/genericswidening/simple_widening/expected/net/me/Example.java

@@ -0,0 +1,9 @@
+package net.me;
+
+import java.util.Set;
+
+public class Example {
+    public static boolean contains(Set<? extends Object> set, Object value) {
+        return set.contains(value);
+    }
+}

tests/data/genericswidening/simple_widening/genericswideners.json

@@ -0,0 +1,7 @@
+{
+    "net/me/Example": {
+        "methodName": "contains",
+        "methodDescriptor": "(Ljava/util/Set;Ljava/lang/Object;)Z",
+        "parameterType": "Ljava/util/Set;"
+    }
+}

tests/data/genericswidening/simple_widening/source/net/me/Example.java

@@ -0,0 +1,9 @@
+package net.me;
+
+import java.util.Set;
+
+public class Example {
+    public static boolean contains(Set<Object> set, Object value) {
+        return set.contains(value);
+    }
+}

tests/src/test/java/net/neoforged/jst/tests/EmbeddedTest.java

@@ -424,6 +424,14 @@ void testStrings() throws Exception {
         }
     }
 
+    @Nested
+    class GenericsWidening {
+        @Test
+        void testSimpleWidening() throws Exception {
+            runGenericsWideningTest("simple_widening");
+        }
+    }
+
     protected final void runInterfaceInjectionTest(String testDirName, Path tempDir, String... additionalArgs) throws Exception {
         var stub = tempDir.resolve("jst-" + testDirName + "-stub.jar");
         testDirName = "interfaceinjection/" + testDirName;
@@ -451,6 +459,17 @@ protected final void runUnpickTest(String testDirName, String... additionalArgs)
         runTest(testDirName, UnaryOperator.identity(), args.toArray(String[]::new));
     }
 
+    protected final void runGenericsWideningTest(String testDirName, String... additionalArgs) throws Exception {
+        testDirName = "genericswidening/" + testDirName;
+        var testDir = testDataRoot.resolve(testDirName);
+        var inputPath = testDir.resolve("genericswideners.json");
+
+        var args = new ArrayList<>(Arrays.asList("--enable-generics-widening", "--generics-widening-data", inputPath.toString()));
+        args.addAll(Arrays.asList(additionalArgs));
+
+        runTest(testDirName, UnaryOperator.identity(), args.toArray(String[]::new));
+    }
+
     protected final void runATTest(String testDirName, final String... extraArgs) throws Exception {
         testDirName = "accesstransformer/" + testDirName;
         var atPath = testDataRoot.resolve(testDirName).resolve("accesstransformer.cfg");