Fix ATs to classes with implicit constructors and add more validation (…

…#21)

---------

Co-authored-by: Sebastian Hartte <[email protected]>

accesstransformers/src/main/java/net/neoforged/jst/accesstransformers/ApplyATsVisitor.java

@@ -1,5 +1,6 @@
 package net.neoforged.jst.accesstransformers;
 
+import com.intellij.lang.jvm.JvmClassKind;
 import com.intellij.navigation.NavigationItem;
 import com.intellij.psi.PsiClass;
 import com.intellij.psi.PsiElement;
@@ -9,8 +10,11 @@
 import com.intellij.psi.PsiModifier;
 import com.intellij.psi.PsiModifierList;
 import com.intellij.psi.PsiModifierListOwner;
+import com.intellij.psi.PsiRecordComponent;
 import com.intellij.psi.PsiRecursiveElementVisitor;
+import com.intellij.psi.PsiWhiteSpace;
 import com.intellij.psi.util.ClassUtil;
+import com.intellij.psi.util.PsiClassUtil;
 import net.neoforged.accesstransformer.parser.AccessTransformerFiles;
 import net.neoforged.accesstransformer.parser.Target;
 import net.neoforged.accesstransformer.parser.Transformation;
@@ -22,9 +26,12 @@
 
 import java.util.Arrays;
 import java.util.EnumMap;
+import java.util.Locale;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 class ApplyATsVisitor extends PsiRecursiveElementVisitor {
     private static final Set<String> ACCESS_MODIFIERS = Set.of(PsiModifier.PUBLIC, PsiModifier.PRIVATE, PsiModifier.PROTECTED);
@@ -33,6 +40,8 @@ class ApplyATsVisitor extends PsiRecursiveElementVisitor {
     public static final EnumMap<Transformation.Modifier, String> MODIFIER_TO_STRING = new EnumMap<>(
             Map.of(Transformation.Modifier.PRIVATE, PsiModifier.PRIVATE, Transformation.Modifier.PUBLIC, PsiModifier.PUBLIC, Transformation.Modifier.PROTECTED, PsiModifier.PROTECTED)
     );
+    public static final Map<String, Transformation.Modifier> STRING_TO_MODIFIER = MODIFIER_TO_STRING.entrySet()
+            .stream().collect(Collectors.toMap(Map.Entry::getValue, Map.Entry::getKey));
 
     private final AccessTransformerFiles ats;
     private final Replacements replacements;
@@ -60,13 +69,16 @@ public void visitElement(@NotNull PsiElement element) {
                     return;
                 }
 
-                apply(pendingATs.remove(new Target.ClassTarget(className)), psiClass, psiClass);
+                var classAt = pendingATs.remove(new Target.ClassTarget(className));
+                apply(classAt, psiClass, psiClass);
                 // We also remove any possible inner class ATs declared for that class as all class targets targeting inner classes
                 // generate a InnerClassTarget AT
                 if (psiClass.getParent() instanceof PsiClass parent) {
                     pendingATs.remove(new Target.InnerClassTarget(ClassUtil.getJVMClassName(parent), className));
                 }
 
+                checkImplicitConstructor(psiClass, className, classAt);
+
                 var fieldWildcard = pendingATs.remove(new Target.WildcardFieldTarget(className));
                 if (fieldWildcard != null) {
                     for (PsiField field : psiClass.getFields()) {
@@ -117,7 +129,13 @@ public String toString() {
                 if (owner instanceof PsiClass cls) {
                     return ClassUtil.getJVMClassName(cls);
                 }
-                return ((NavigationItem) owner).getName() + " of " + ClassUtil.getJVMClassName(containingClass);
+                String memberName;
+                if (owner instanceof PsiMethod mtd && mtd.isConstructor()) {
+                    memberName = "constructor";
+                } else {
+                    memberName = ((NavigationItem) owner).getName();
+                }
+                return memberName + " of " + ClassUtil.getJVMClassName(containingClass);
             }
         };
         logger.debug("Applying AT %s to %s", at, targetInfo);
@@ -138,7 +156,10 @@ public String toString() {
                     }
                 }
             }
-        } else if (targetAcc == Transformation.Modifier.DEFAULT || !modifiers.hasModifierProperty(MODIFIER_TO_STRING.get(targetAcc))) {
+        } else if (containingClass.isEnum() && owner instanceof PsiMethod mtd && mtd.isConstructor() && at.modifier().ordinal() < Transformation.Modifier.DEFAULT.ordinal()) {
+            // Enum constructors can at best be package-private, any other attempt must be prevented
+            error("Access transformer %s targeting %s attempted to make an enum constructor %s", at, targetInfo, at.modifier());
+        } else if (targetAcc.ordinal() < detectModifier(modifiers, null).ordinal()) { // PUBLIC (0) < PROTECTED (1) < DEFAULT (2) < PRIVATE (3)
             modify(targetAcc, modifiers, Arrays.stream(modifiers.getChildren())
                     .filter(el -> el instanceof PsiKeyword)
                     .map(el -> (PsiKeyword) el)
@@ -211,6 +232,68 @@ private void modify(Transformation.Modifier targetAcc, PsiModifierList modifiers
         }
     }
 
+    /**
+     * Handle the different behavior between applying ATs to source vs. bytecode.
+     * In source, the implicit class constructor is not visible and its access level will automatically increase
+     * when we change that of the containing class, while that is not true in production, where ATs are applied
+     * to bytecode instead.
+     * It also handles additional validation for record constructors, which have special rules.
+     */
+    private void checkImplicitConstructor(PsiClass psiClass, String className, @Nullable Transformation classAt) {
+        if (psiClass.isRecord()) {
+            StringBuilder descriptor = new StringBuilder("(");
+            for (PsiRecordComponent recordComponent : psiClass.getRecordComponents()) {
+                descriptor.append(ClassUtil.getBinaryPresentation(recordComponent.getType()));
+            }
+            descriptor.append(")V");
+            var desc = descriptor.toString();
+
+            var implicitAT = pendingATs.remove(new Target.MethodTarget(className, "<init>", desc));
+            if (implicitAT != null && implicitAT.modifier() != detectModifier(psiClass.getModifierList(), classAt)) {
+                error("Access transformer %s targeting the implicit constructor of %s is not valid, as a record's constructor must have the same access level as the record class. Please AT the record too: \"%s\"", implicitAT, className,
+                        implicitAT.modifier().toString().toLowerCase(Locale.ROOT) + " " + className);
+                pendingATs.remove(new Target.MethodTarget(className, "<init>", desc));
+            } else if (classAt != null && detectModifier(psiClass.getModifierList(), null).ordinal() > classAt.modifier().ordinal() && implicitAT == null) {
+                error("Access transformer %s targeting record class %s attempts to widen its access without widening the constructor's access. You must AT the constructor too: \"%s\"", classAt, className,
+                        classAt.modifier().toString().toLowerCase(Locale.ROOT) + " " + className + " <init>" + desc);
+                pendingATs.remove(new Target.MethodTarget(className, "<init>", desc));
+            }
+        } else if (psiClass.getClassKind() == JvmClassKind.CLASS) {
+            // When widening the access of a class, we must take into consideration the fact that implicit constructors follow the access level of their owner
+            if (psiClass.getConstructors().length == 0) {
+                var constructorTarget = new Target.MethodTarget(className, "<init>", PsiHelper.getImplicitConstructorSignature(psiClass));
+                var constructorAt = pendingATs.remove(constructorTarget);
+
+                if (classAt != null && detectModifier(psiClass.getModifierList(), null).ordinal() > classAt.modifier().ordinal()) {
+                    // If we cannot find an implicit constructor, we need to inject it if the AT doesn't match the expected constructor access
+                    if (constructorAt == null || constructorAt.modifier() != classAt.modifier()) {
+                        var expectedModifier = detectModifier(psiClass.getModifierList(), constructorAt);
+                        injectConstructor(psiClass, className, expectedModifier);
+                    }
+                } else if (constructorAt != null && constructorAt.modifier().ordinal() < detectModifier(psiClass.getModifierList(), null).ordinal()) {
+                    // If we're trying to widen the access of an undeclared implicit constructor, we must inject it
+                    injectConstructor(psiClass, className, constructorAt.modifier());
+                }
+            }
+        }
+    }
+
+    private void injectConstructor(PsiClass psiClass, String className, Transformation.Modifier modifier) {
+        logger.debug("Injecting implicit constructor with access level %s into class %s", modifier, className);
+
+        // Add 4 spaces of indent to indent the constructor inside the class
+        int indent = 4;
+        // If the class is preceded by whitespace, use the last line of that whitespace as the base indent
+        if (psiClass.getPrevSibling() instanceof PsiWhiteSpace psiWhiteSpace) {
+            indent += PsiHelper.getLastLineLength(psiWhiteSpace);
+        }
+
+        final String modifierString = modifier == Transformation.Modifier.DEFAULT ? "" : (MODIFIER_TO_STRING.get(modifier) + " ");
+        // Inject the constructor after the opening brace, on a new line
+        replacements.insertAfter(Objects.requireNonNull(psiClass.getLBrace()), "\n" +
+                " ".repeat(indent) + modifierString + psiClass.getName() + "() {}");
+    }
+
     private void error(String message, Object... args) {
         logger.error(message, args);
         errored = true;
@@ -235,4 +318,17 @@ private static Transformation merge(Transformation a, @Nullable Transformation b
     private static Target.MethodTarget method(String owner, PsiMethod method) {
         return new Target.MethodTarget(owner, PsiHelper.getBinaryMethodName(method), PsiHelper.getBinaryMethodSignature(method));
     }
+
+    private static Transformation.Modifier detectModifier(PsiModifierList owner, @Nullable Transformation trans) {
+        if (trans != null) {
+            return trans.modifier();
+        }
+
+        for (String mod : ACCESS_MODIFIERS) {
+            if (owner.hasModifierProperty(mod)) {
+                return STRING_TO_MODIFIER.get(mod);
+            }
+        }
+        return Transformation.Modifier.DEFAULT;
+    }
 }

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

@@ -8,6 +8,7 @@
 import com.intellij.psi.PsiParameterListOwner;
 import com.intellij.psi.PsiPrimitiveType;
 import com.intellij.psi.PsiTypes;
+import com.intellij.psi.PsiWhiteSpace;
 import com.intellij.psi.SyntaxTraverser;
 import com.intellij.psi.util.CachedValueProvider;
 import com.intellij.psi.util.CachedValuesManager;
@@ -79,6 +80,20 @@ public String next() {
         };
     }
 
+    public static String getImplicitConstructorSignature(PsiClass psiClass) {
+        StringBuilder signature = new StringBuilder();
+        signature.append("(");
+        // Non-Static inner class constructors have the enclosing class as their first argument
+        if (isNonStaticInnerClass(psiClass)) {
+            var parent = Objects.requireNonNull(Objects.requireNonNull(psiClass.getContainingClass()));
+            signature.append("L");
+            getBinaryClassName(parent, signature);
+            signature.append(";");
+        }
+        signature.append(")V");
+        return signature.toString();
+    }
+
     public static String getBinaryMethodSignature(PsiMethod method) {
         StringBuilder signature = new StringBuilder();
         signature.append("(");
@@ -164,13 +179,15 @@ private static boolean isEnumConstructor(PsiMethod method) {
     private static boolean isNonStaticInnerClassConstructor(PsiMethod method) {
         if (method.isConstructor()) {
             var containingClass = method.getContainingClass();
-            return containingClass != null
-                    && containingClass.getContainingClass() != null
-                    && !containingClass.hasModifierProperty(PsiModifier.STATIC);
+            return containingClass != null && isNonStaticInnerClass(containingClass);
         }
         return false;
     }
 
+    private static boolean isNonStaticInnerClass(PsiClass psiClass) {
+        return psiClass.getContainingClass() != null && !psiClass.hasModifierProperty(PsiModifier.STATIC);
+    }
+
     /**
      * Gets the local variable table indices of the parameters for the given method
      * or lambda expression
@@ -214,4 +231,13 @@ public static boolean isRecordConstructor(PsiMethod psiMethod) {
         var containingClass = psiMethod.getContainingClass();
         return containingClass != null && containingClass.isRecord() && psiMethod.isConstructor();
     }
+
+    public static int getLastLineLength(PsiWhiteSpace psiWhiteSpace) {
+        var lastNewline = psiWhiteSpace.getText().lastIndexOf('\n');
+        if (lastNewline != -1) {
+            return psiWhiteSpace.getTextLength() - lastNewline - 1;
+        } else {
+            return psiWhiteSpace.getTextLength();
+        }
+    }
 }

parchment/src/main/java/net/neoforged/jst/parchment/JavadocHelper.java

@@ -7,6 +7,7 @@
 import com.intellij.psi.PsiWhiteSpace;
 import com.intellij.psi.javadoc.PsiDocComment;
 import com.intellij.psi.javadoc.PsiDocToken;
+import net.neoforged.jst.api.PsiHelper;
 import net.neoforged.jst.api.Replacement;
 import net.neoforged.jst.api.Replacements;
 import org.jetbrains.annotations.Nullable;
@@ -107,7 +108,7 @@ public static void enrichJavadoc(PsiJavaDocumentedElement psiElement,
             int indent = 0;
             // If the element is preceded by whitespace, use the last line of that whitespace as the indent
             if (psiElement.getPrevSibling() instanceof PsiWhiteSpace psiWhiteSpace) {
-                indent = JavadocHelper.getLastLineLength(psiWhiteSpace);
+                indent = PsiHelper.getLastLineLength(psiWhiteSpace);
             }
             replacements.insertBefore(
                     psiElement,
@@ -179,11 +180,11 @@ public static int getIndent(PsiDocCommentBase comment) {
             if (child instanceof PsiDocToken token && child.getPrevSibling() instanceof PsiWhiteSpace precedingWhitespace) {
                 var type = token.getTokenType();
                 if (type == JavaDocTokenType.DOC_COMMENT_START) {
-                    indentLength = Math.max(indentLength, getLastLineLength(precedingWhitespace));
+                    indentLength = Math.max(indentLength, PsiHelper.getLastLineLength(precedingWhitespace));
 
                 } else if (type == JavaDocTokenType.DOC_COMMENT_LEADING_ASTERISKS || type == JavaDocTokenType.DOC_COMMENT_END) {
                     // Strip one space since it includes the inner-javadoc indent
-                    indentLength = Math.max(indentLength, getLastLineLength(precedingWhitespace) - 1);
+                    indentLength = Math.max(indentLength, PsiHelper.getLastLineLength(precedingWhitespace) - 1);
                 }
             }
         }
@@ -302,15 +303,6 @@ private static StringBuilder endLine(StringBuilder sb) {
         return sb.append('\n');
     }
 
-    public static int getLastLineLength(PsiWhiteSpace psiWhiteSpace) {
-        var lastNewline = psiWhiteSpace.getText().lastIndexOf('\n');
-        if (lastNewline != -1) {
-            return psiWhiteSpace.getTextLength() - lastNewline - 1;
-        } else {
-            return psiWhiteSpace.getTextLength();
-        }
-    }
-
     /**
      * @param refersTo In case of param or throws, this is the parameter name or exception that the tag refers to.
      */

tests/data/accesstransformer/classes/expected/C1.java

@@ -1,3 +1,3 @@
 public class C1 {
-
+    C1() {}
 }

tests/data/accesstransformer/classes/source/C1.java

@@ -1,3 +1,2 @@
 class C1 {
-
 }

tests/data/accesstransformer/illegal/accesstransformer.cfg

@@ -0,0 +1,2 @@
+public AnEnum <init>(Ljava/lang/String;IZ)V # This is invalid, we cannot widen the access of an enum ctor past package
+default AnEnum <init>(Ljava/lang/String;II)V # Useless, but valid

tests/data/accesstransformer/illegal/expected.log

@@ -0,0 +1 @@
+Access transformer PUBLIC LEAVE {atpath}:1 targeting constructor of AnEnum attempted to make an enum constructor PUBLIC

tests/data/accesstransformer/illegal/expected/AnEnum.java

@@ -0,0 +1,11 @@
+public enum AnEnum {
+    ;
+
+    AnEnum(int i) {
+
+    }
+
+    AnEnum(boolean ok) {
+
+    }
+}

tests/data/accesstransformer/illegal/source/AnEnum.java

@@ -0,0 +1,11 @@
+public enum AnEnum {
+    ;
+
+    private AnEnum(int i) {
+
+    }
+
+    AnEnum(boolean ok) {
+
+    }
+}

tests/data/accesstransformer/implicit_constructors/accesstransformer.cfg

@@ -0,0 +1,15 @@
+protected PrivateClass <init>()V # Should create a protected constructor
+default PrivateClass$NestedProtected <init>()V # Should not cause any change, as default < protected
+public TestPrivateCtor # Should create a default constructor
+
+# Should only change the class to public since the implicit constructor is expected to have the same access level
+public PrivateClass$Both
+public PrivateClass$Both <init>()V
+
+public PrivateRecord # Error because the ctor needs to be AT'd too
+
+public PrivateRecord$Nested
+public PrivateRecord$Nested <init>(I)V
+
+protected PrivateClass$Inner # Make the class itself protected
+public PrivateClass$Inner <init>(LPrivateClass;)V # and the constructor public

tests/data/accesstransformer/implicit_constructors/expected.log

@@ -0,0 +1 @@
+Access transformer PUBLIC LEAVE {atpath}:9 targeting record class PrivateRecord attempts to widen its access without widening the constructor's access. You must AT the constructor too: "public PrivateRecord <init>(I)V"

tests/data/accesstransformer/implicit_constructors/expected/PrivateClass.java

@@ -0,0 +1,16 @@
+private class PrivateClass {
+    protected PrivateClass() {}
+
+    protected static class NestedProtected {
+
+    }
+
+    public static class Both {
+
+    }
+
+    protected class Inner {
+        public Inner() {}
+
+    }
+}

tests/data/accesstransformer/implicit_constructors/expected/PrivateRecord.java

@@ -0,0 +1,5 @@
+public record PrivateRecord(int a) {
+    public record Nested(int b) {
+
+    }
+}

tests/data/accesstransformer/implicit_constructors/expected/TestPrivateCtor.java

@@ -0,0 +1,4 @@
+public class TestPrivateCtor {
+    TestPrivateCtor() {}
+
+}

tests/data/accesstransformer/implicit_constructors/source/PrivateClass.java

@@ -0,0 +1,14 @@
+private class PrivateClass {
+
+    protected static class NestedProtected {
+
+    }
+
+    protected static class Both {
+
+    }
+
+    private class Inner {
+
+    }
+}

tests/data/accesstransformer/implicit_constructors/source/PrivateRecord.java

@@ -0,0 +1,5 @@
+private record PrivateRecord(int a) {
+    private record Nested(int b) {
+
+    }
+}

tests/data/accesstransformer/implicit_constructors/source/TestPrivateCtor.java

@@ -0,0 +1,3 @@
+class TestPrivateCtor {
+
+}