Annotate nullable array type parameters (#779)

- openrewrite/rewrite-migrate-java#934

Author: timtebeek

src/main/java/org/openrewrite/staticanalysis/AnnotateNullableParameters.java

@@ -31,6 +31,8 @@
 
 import java.util.*;
 
+import static java.util.Collections.singletonList;
+import static java.util.Comparator.comparing;
 import static java.util.stream.Collectors.toList;
 
 @EqualsAndHashCode(callSuper = false)
@@ -110,7 +112,31 @@ public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration methodDecl
                                             String.format("package %s;public @interface %s {}", fullyQualifiedPackage, simpleName)))
                                     .build()
                                     .apply(new Cursor(getCursor(), vd),
-                                            vd.getCoordinates().addAnnotation(Comparator.comparing(J.Annotation::getSimpleName)));
+                                            vd.getCoordinates().addAnnotation(comparing(J.Annotation::getSimpleName)));
+
+                            // For array types, move annotation from leading annotations to array brackets
+                            if (annotated.getTypeExpression() instanceof J.ArrayType) {
+                                // Find the annotation we just added
+                                J.Annotation nullableAnnotation = null;
+                                for (J.Annotation ann : annotated.getLeadingAnnotations()) {
+                                    if (ann.getSimpleName().equals(simpleName)) {
+                                        nullableAnnotation = ann;
+                                        break;
+                                    }
+                                }
+                                if (nullableAnnotation != null) {
+                                    J.Annotation finalAnnotation = nullableAnnotation;
+                                    J.ArrayType arrayType = (J.ArrayType) annotated.getTypeExpression();
+                                    annotated = annotated.withLeadingAnnotations(ListUtils.map(annotated.getLeadingAnnotations(),
+                                            a -> a == finalAnnotation ? null : a));
+                                    arrayType = arrayType.withAnnotations(singletonList(finalAnnotation.withPrefix(Space.SINGLE_SPACE)));
+                                    if (annotated.getLeadingAnnotations().isEmpty()) {
+                                        arrayType = arrayType.withPrefix(Space.EMPTY);
+                                    }
+                                    annotated = annotated.withTypeExpression(arrayType);
+                                }
+                            }
+
                             doAfterVisit(ShortenFullyQualifiedTypeReferences.modifyOnly(annotated));
                             doAfterVisit(new MoveFieldAnnotationToType(fullyQualifiedName).getVisitor());
                             return annotated.withModifiers(ListUtils.mapFirst(annotated.getModifiers(), first -> first.withPrefix(Space.SINGLE_SPACE)));

src/test/java/org/openrewrite/staticanalysis/AnnotateNullableParametersTest.java

@@ -21,6 +21,7 @@
 import org.junit.jupiter.params.provider.CsvSource;
 import org.junit.jupiter.params.provider.ValueSource;
 import org.openrewrite.DocumentExample;
+import org.openrewrite.Issue;
 import org.openrewrite.java.JavaParser;
 import org.openrewrite.test.RecipeSpec;
 import org.openrewrite.test.RewriteTest;
@@ -256,6 +257,36 @@ public void bar(@Nullable final String name) {
               )
             );
         }
+
+        @Issue("https://github.com/openrewrite/rewrite-migrate-java/issues/934")
+        @Test
+        void arrayParameter() {
+            rewriteRun(
+              //language=java
+              java(
+                """
+                  public class ArrayProcessor {
+                      public void processItems(String[] items) {
+                          if (items == null) {
+                              return;
+                          }
+                      }
+                  }
+                  """,
+                """
+                  import org.jspecify.annotations.Nullable;
+
+                  public class ArrayProcessor {
+                      public void processItems(String @Nullable[] items) {
+                          if (items == null) {
+                              return;
+                          }
+                      }
+                  }
+                  """
+              )
+            );
+        }
     }
 
     @Nested