Flag to inherit access transformers on methods (#1)

Useful for modifying methods exposed by types with a lot of children
to avoid specifying them all.

Author: lynxplay

README.md

@@ -102,6 +102,9 @@ Plugin - parchment
 Plugin - accesstransformers
       --access-transformer=<atFiles>
 
+      --access-transformer-inherit-method
+                           Whether or not access transformers on methods should be inherited from
+                             parent types
       --access-transformer-validation=<validation>
                            The level of validation to use for ats
       --enable-accesstransformers

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

@@ -25,6 +25,12 @@ public class AccessTransformersTransformer implements SourceTransformer {
     @CommandLine.Option(names = "--access-transformer-validation", description = "The level of validation to use for ats")
     public AccessTransformerValidation validation = AccessTransformerValidation.LOG;
 
+    @CommandLine.Option(
+        names = "--access-transformer-inherit-method",
+        description = "Whether or not access transformers on methods should be inherited from parent types"
+    )
+    public boolean inheritMethodAccessTransformers = false;
+
     private AccessTransformerFiles ats;
     private Map<Target, Transformation> pendingATs;
     private Logger logger;
@@ -64,7 +70,13 @@ public boolean afterRun(TransformContext context) {
 
     @Override
     public void visitFile(PsiFile psiFile, Replacements replacements) {
-        var visitor = new ApplyATsVisitor(ats, replacements, pendingATs, logger);
+        var visitor = new ApplyATsVisitor(
+            ats,
+            replacements,
+            pendingATs,
+            this.inheritMethodAccessTransformers,
+            logger
+        );
         visitor.visitFile(psiFile);
         if (visitor.errored) {
             errored = true;

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

@@ -14,7 +14,6 @@
 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;
@@ -46,12 +45,20 @@ class ApplyATsVisitor extends PsiRecursiveElementVisitor {
     private final AccessTransformerFiles ats;
     private final Replacements replacements;
     private final Map<Target, Transformation> pendingATs;
+    private final boolean inheritMethodATs;
     private final Logger logger;
     boolean errored = false;
 
-    public ApplyATsVisitor(AccessTransformerFiles ats, Replacements replacements, Map<Target, Transformation> pendingATs, Logger logger) {
+    public ApplyATsVisitor(
+        AccessTransformerFiles ats,
+        Replacements replacements,
+        Map<Target, Transformation> pendingATs,
+        boolean inheritMethodATs,
+        Logger logger
+    ) {
         this.ats = ats;
         this.replacements = replacements;
+        this.inheritMethodATs = inheritMethodATs;
         this.logger = logger;
         this.pendingATs = pendingATs;
     }
@@ -61,7 +68,25 @@ public void visitElement(@NotNull PsiElement element) {
         if (element instanceof PsiClass psiClass) {
             if (psiClass.getQualifiedName() != null) {
                 String className = ClassUtil.getJVMClassName(psiClass);
-                if (!ats.containsClassTarget(className)) {
+                boolean shouldBeSkipped = !ats.containsClassTarget(className);
+
+                // Classes may need changing that are not in the source at set.
+                // Check if any of this classe's parents are in the source at set first to potentially inherit their ATs.
+                if (shouldBeSkipped && this.inheritMethodATs) {
+                    for (
+                        PsiClass parentType = psiClass.getSuperClass();
+                        parentType != null;
+                        parentType = parentType.getSuperClass()
+                    ) {
+                        if (this.ats.containsClassTarget(ClassUtil.getJVMClassName(parentType))) {
+                            // Process class as we inherit method ATs and we found *some* AT.
+                            // Heuristic isn't perfect, but the alternative would be scanning *all* method targets for the parent owning type.
+                            shouldBeSkipped = false;
+                            break;
+                        }
+                    }
+                }
+                if (shouldBeSkipped) {
                     // Skip this class and its children, but not the inner classes
                     for (PsiClass innerClass : psiClass.getInnerClasses()) {
                         visitElement(innerClass);
@@ -106,10 +131,20 @@ public void visitElement(@NotNull PsiElement element) {
                 apply(pendingATs.remove(new Target.FieldTarget(className, field.getName())), field, cls);
             }
         } else if (element instanceof PsiMethod method) {
-            final var cls = method.getContainingClass();
-            if (cls != null && cls.getQualifiedName() != null) {
-                String className = ClassUtil.getJVMClassName(cls);
-                apply(pendingATs.remove(method(className, method)), method, cls);
+            final var owningType = method.getContainingClass();
+            if (owningType != null) {
+                // Locate transformation by searching owning type and its parents.
+                // Remove if this is AT is defined for its immediate owning type.
+                Transformation foundTransformation = this.pendingATs.remove(method(ClassUtil.getJVMClassName(owningType), method));
+                PsiClass definedForSuperType = owningType;
+                for (;
+                     this.inheritMethodATs && foundTransformation == null && definedForSuperType != null && definedForSuperType.getQualifiedName() != null;
+                     definedForSuperType = definedForSuperType.getSuperClass()
+                ) {
+                    foundTransformation = this.ats.getAccessTransformers().get(method(ClassUtil.getJVMClassName(definedForSuperType), method));
+                }
+
+                apply(foundTransformation, method, owningType);
             }
         }
 

tests/data/accesstransformer/methods_inheritance/accesstransformer.cfg

@@ -0,0 +1 @@
+public Parent flag()Z

tests/data/accesstransformer/methods_inheritance/expected/Child.java

@@ -0,0 +1,6 @@
+class Child extends Parent {
+    @Override
+    public boolean flag() {
+        return true;
+    }
+}

tests/data/accesstransformer/methods_inheritance/expected/Grandchild.java

@@ -0,0 +1,6 @@
+class Grandchild extends Child {
+    @Override
+    public boolean flag() {
+        return false;
+    }
+}

tests/data/accesstransformer/methods_inheritance/expected/Parent.java

@@ -0,0 +1,5 @@
+class Parent {
+    public boolean flag() {
+        return false;
+    }
+}

tests/data/accesstransformer/methods_inheritance/source/Child.java

@@ -0,0 +1,6 @@
+class Child extends Parent {
+    @Override
+    protected boolean flag() {
+        return true;
+    }
+}

tests/data/accesstransformer/methods_inheritance/source/Grandchild.java

@@ -0,0 +1,6 @@
+class Grandchild extends Child {
+    @Override
+    protected boolean flag() {
+        return false;
+    }
+}

tests/data/accesstransformer/methods_inheritance/source/Parent.java

@@ -0,0 +1,5 @@
+class Parent {
+    protected boolean flag() {
+        return false;
+    }
+}