Don't remove referenced annotations from record parameters

[mvitz]

What's changed?

What's your motivation?

• This fixes RemoveUnusedImports removes used imports for annotated method parameters #4473

Anything in particular you'd like reviewers to focus on?

Anyone you would like to review specifically?

Have you considered any alternatives or workarounds?

Any additional context

Checklist

• I've added unit tests to cover both positive and negative cases
• I've read and applied the recipe conventions and best practices
• I've used the IntelliJ IDEA auto-formatter on affected files

[mvitz]

Currently, the test fails with:

    java.lang.AssertionError: Source file was parsed into an LST that contains non-whitespace characters in its whitespace. This is indicative of a bug in the parser.
    import java.beans.BeanProperty;
    record A(~~(non-whitespace)~~>@BeanProperty <~~String a) {}
        at org.openrewrite.test.RewriteTest.rewriteRun(RewriteTest.java:323)
        at org.openrewrite.test.RewriteTest.rewriteRun(RewriteTest.java:131)
        at org.openrewrite.test.RewriteTest.rewriteRun(RewriteTest.java:126)
        at org.openrewrite.java.RemoveUnusedImportsTest.doesNotRemoveReferencedAnnotationFromRecordParameter(RemoveUnusedImportsTest.java:1821)```

[timtebeek]

Thanks for getting this reduced and started @mvitz ! Looks like an oversight in the parser after the addition of records, which would then cause the annotation to not be parsed as such, but as "whitespace". The fix is then likely in ReloadableJava17ParserVisitor.

[mvitz]

Thanks for the hint.
When debugging I see that visitClass is called but there the record is already presented as a class with a constructor.
Unfortunately, at that point the annotation of the parameter is gone, and I didn't find any other visitor callback to see the original record.

[timtebeek]

Hmm; thanks for diving in! It looks like the annotation is still present at this point in the parsing:

rewrite/rewrite-java-17/src/main/java/org/openrewrite/java/isolated/ReloadableJava17Parser.java

Lines 193 to 194 in a98d83d

	JCTree.JCCompilationUnit jcCompilationUnit = compiler.parse(new ReloadableJava17ParserInputFileObject(input1, ctx));
	cus.put(input1, jcCompilationUnit);

But it's gone when we exit that org.openrewrite.java.isolated.ReloadableJava17Parser#parseInputsToCompilerAst 🤔

[mvitz]

Sorry, my fault. I added a more accurate testcase where the annotation contains PARAMETER and RECORD_COMPONENT as @Target and now the annotation is preserved on the generated constructor parameter and can be seen from ReloadableJava17ParserVisitor.

[mvitz]

Note to myself, if I change Lines 411 to 415 into

                if (member instanceof JCMethodDecl m) {
                    if (hasFlag(m.getModifiers(), Flags.GENERATEDCONSTR)) {
                        for (Tree param : m.getParameters()) {
                            stateVector.add(param);
                        }

The assertion failure is still there, but now marks another non-whitespace:

    java.lang.AssertionError: Source file was parsed into an LST that contains non-whitespace characters in its whitespace. This is indicative of a bug in the parser. 
    package b;
    import a.A;
    record B(~~(non-whitespace)~~>@A <~~~~(non-whitespace)~~>String <~~a) {}

Within private <J2 extends J> J2 convert(Tree t) { (Line 1667) the call to ((JCTree) t).getStartPosition() returns another value than cursor for this parameter and therefore the cursor is advanced, and the Annotation is skipped when the call to scan is done.
For Methods and their parameters (see visitMethod) the startPosition is equal to the cursor and therefore that annotations are preserved. However both are of type JCVariableDecl but the record constructor parameters contain an additional new line (\n) between the annotation and the type.

After thinking during the night and checking some examples today annotations on records opens a rabbit hole. All the following are valid annotations for a record component:

@Target({METHOD})
@Retention(RUNTIME)
public @interface MethodAnnotation {
}

@Target({PARAMETER})
@Retention(RUNTIME)
public @interface ParameterAnnotation {
}

@Target({FIELD})
@Retention(RUNTIME)
public @interface FieldAnnotation {
}

However, in the generated class (after compilation) they appear on different locations:

• MethodAnnotation appears at the generated accessor method.
• ParameterAnnotation appears at the corresponding argument within the generated record constructor.
• FieldAnnotation appears at the corresponding generated private final field.

Therefore, the first initial example with record B(@BeanProperty String a) is a valid case and the import for that annotation should not be removed. However, that would mean that this annotation needs to be detected by looking at the generated accessor method...

[timtebeek]

Thanks for the detailed look! Problems sure have a way of looking deceptively simple don't they? 🙃
Let me tag @knutwannheden here as he might have some ideas too, even if he's caught up with other work at the moment.

[Hillrunner2008]

///usr/bin/env jbang "$0" "$@" ; exit $?
//DEPS org.openrewrite:rewrite-java:8.44.1
//DEPS org.openrewrite:rewrite-java-21:8.44.1
//DEPS org.openrewrite:rewrite-test:8.44.1
//DEPS org.junit.jupiter:junit-jupiter:5.11.4
//DEPS org.junit.platform:junit-platform-launcher:1.11.4
//DEPS org.springframework:spring-web:6.1.3
//DEPS org.projectlombok:lombok:1.18.36

import org.openrewrite.java.tree.J;
import org.openrewrite.ExecutionContext;
import org.openrewrite.InMemoryExecutionContext;
import org.openrewrite.java.JavaParser;
import org.openrewrite.Tree;
import org.openrewrite.SourceFile;
import org.openrewrite.Recipe;
import org.openrewrite.TreeVisitor;
import org.openrewrite.java.RemoveUnusedImports;
import org.openrewrite.java.tree.JavaType;
import org.openrewrite.java.tree.TypeUtils;
import org.openrewrite.java.JavaIsoVisitor;
import org.openrewrite.java.tree.J;
import org.openrewrite.ExecutionContext;
import org.openrewrite.InMemoryExecutionContext;
import java.util.*;


public class DemoRemoveUnusedImports {
    public static void main(String[] args) {
        String sourceCode = """
                package com.example;
                
                import org.springframework.web.bind.annotation.RequestParam; // This should not be removed
                import java.util.List; // This should be removed
                import lombok.Builder;
                
                @Builder
                public record TestParams(
                        @RequestParam(required = false)
                        Integer limit
                ) {}
                """;

        JavaParser jp = JavaParser.fromJavaVersion()
                .classpath("spring-web", "lombok")
                .build();

       System.out.println("=== Testing RemoveUnusedImports ===");
        testRecipe(new RemoveUnusedImports(), jp, sourceCode);
       jp.reset();
        System.out.println("\n=== Testing PatchedRemoveUnusedImports ===");
        testRecipe(new PatchedRemoveUnusedImports(), jp, sourceCode);
    }
    
    private static void testRecipe(Recipe recipe, JavaParser jp, String sourceCode) {
        try {
            List<SourceFile> sourceFiles = jp.parse(sourceCode).toList();
            if (sourceFiles.isEmpty()) {
                System.out.println("No source files parsed.");
                return;
            }

            J.CompilationUnit cu = (J.CompilationUnit) sourceFiles.get(0);
            System.out.println("--- Original Source Code ---");
            System.out.println(sourceCode);
           System.out.println("---");

            Tree result = recipe.getVisitor().visit(cu, new InMemoryExecutionContext());
            
             System.out.println("--- Modified Source Code ---");
           System.out.println(((J.CompilationUnit) result).print());
            System.out.println("---");

        } catch (Exception e) {
            System.out.println("Error occurred:");
            e.printStackTrace(System.out);
        }
    }
}

class PatchedRemoveUnusedImports extends Recipe {
    @Override
    public String getDisplayName() {
        return "Remove unused imports (record-aware for older Rewrite)";
    }

    @Override
    public String getDescription() {
        return "Removes imports for types that are not referenced, scanning canonical constructor parameters for record annotations.";
    }

    @Override
    public TreeVisitor<?, ExecutionContext> getVisitor() {
        return new JavaIsoVisitor<ExecutionContext>() {
            private final Set<String> usedTypes = new HashSet<>();

            @Override
            public J.CompilationUnit visitCompilationUnit(J.CompilationUnit cu, ExecutionContext ctx) {
                // Normal AST traversal first
                super.visitCompilationUnit(cu, ctx);

                // Also gather from getTypesInUse()
                cu.getTypesInUse().getTypesInUse().forEach(t -> {
                    if (t instanceof JavaType.FullyQualified) {
                        usedTypes.add(((JavaType.FullyQualified) t).getFullyQualifiedName());
                    }
                });

                // Filter out unused imports
                List<J.Import> retainedImports = new ArrayList<>();
                for (J.Import anImport : cu.getImports()) {
                    JavaType importType = anImport.getQualid().getType();
                    if (importType instanceof JavaType.FullyQualified fq) {
                        if (usedTypes.contains(fq.getFullyQualifiedName())) {
                            retainedImports.add(anImport);
                        }
                    }
                }
                return cu.withImports(retainedImports);
            }

            @Override
            public J.ClassDeclaration visitClassDeclaration(J.ClassDeclaration classDecl, ExecutionContext ctx) {
                boolean isRecord = classDecl.getModifiers().stream()
                    .anyMatch(m -> 
                        m.getType() == J.Modifier.Type.LanguageExtension 
                        && "record".equalsIgnoreCase(m.getKeyword())
                    );

                if (isRecord) {
                    for (Object bodyStmt : classDecl.getBody().getStatements()) {
                        if (bodyStmt instanceof J.MethodDeclaration md) {
                            boolean constructorMatchesClassName =
                                md.isConstructor()
                                && md.getSimpleName().equals(classDecl.getSimpleName());
                            if (constructorMatchesClassName) {
                                for (Object param : md.getParameters()) {
                                    if (param instanceof J.VariableDeclarations varDec) {
                                        for (J.Annotation ann : varDec.getLeadingAnnotations()) {
                                            recordAnnotationUsage(ann);
                                        }
                                    }
                                }
                            }
                        }
                    }
                }

                for (J.Annotation ann : classDecl.getLeadingAnnotations()) {
                    recordAnnotationUsage(ann);
                }

                return super.visitClassDeclaration(classDecl, ctx);
            }

            @Override
            public J.Annotation visitAnnotation(J.Annotation annotation, ExecutionContext ctx) {
                recordAnnotationUsage(annotation);
                return super.visitAnnotation(annotation, ctx);
            }

            private void recordAnnotationUsage(J.Annotation ann) {
                JavaType.FullyQualified fq = TypeUtils.asFullyQualified(ann.getType());
                if (fq != null) {
                    usedTypes.add(fq.getFullyQualifiedName());
                }
            }
        };
    }
}
}

output

=== Testing RemoveUnusedImports ===
--- Original Source Code ---
package com.example;

import org.springframework.web.bind.annotation.RequestParam; // This should not be removed
import java.util.List; // This should be removed
import lombok.Builder;

@Builder
public record TestParams(
        @RequestParam(required = false)
        Integer limit
) {}

---
--- Modified Source Code ---
package com.example; // This should be removed
import lombok.Builder;

@Builder
public record TestParams(
        @RequestParam(required = false)
        Integer limit
) {}

---

=== Testing PatchedRemoveUnusedImports ===
--- Original Source Code ---
package com.example;

import org.springframework.web.bind.annotation.RequestParam; // This should not be removed
import java.util.List; // This should be removed
import lombok.Builder;

@Builder
public record TestParams(
        @RequestParam(required = false)
        Integer limit
) {}

---
--- Modified Source Code ---
package com.example; // This should be removed
import lombok.Builder;

@Builder
public record TestParams(
        @RequestParam(required = false)
        Integer limit
) {}

---

Long story short i didn't quite understand the best introspection approach in this context but if you want a reproducer this jbang does the job of before/after.