Move nullable annotations to array types

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

@@ -24,6 +24,7 @@
 import org.openrewrite.java.tree.Expression;
 import org.openrewrite.java.tree.J;
 import org.openrewrite.java.tree.JavaType;
+import org.openrewrite.java.tree.TypeTree;
 import org.openrewrite.staticanalysis.java.MoveFieldAnnotationToType;
 
 import java.util.Arrays;
@@ -80,8 +81,7 @@ public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration methodDecl
                         methodDeclaration.getMethodType() == null ||
                         methodDeclaration.getMethodType().getReturnType() instanceof JavaType.Primitive ||
                         service(AnnotationService.class).matches(getCursor(), new AnnotationMatcher("@" + fullyQualifiedName)) ||
-                        (methodDeclaration.getReturnTypeExpression() != null &&
-                                service(AnnotationService.class).matches(new Cursor(null, methodDeclaration.getReturnTypeExpression()), new AnnotationMatcher("@" + fullyQualifiedName)))) {
+                        hasNullableAnnotation(methodDeclaration.getReturnTypeExpression(), fullyQualifiedName)) {
                     return methodDeclaration;
                 }
 
@@ -100,6 +100,27 @@ public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration methodDecl
                 }
                 return md;
             }
+
+            private boolean hasNullableAnnotation(@Nullable TypeTree returnType, String annotationFqn) {
+                if (returnType == null) {
+                    return false;
+                }
+
+                // Check if the return type itself is annotated
+                if (service(AnnotationService.class).matches(new Cursor(null, returnType), new AnnotationMatcher("@" + annotationFqn))) {
+                    return true;
+                }
+
+                // For array types, check if the element type is annotated
+                if (returnType instanceof J.ArrayType) {
+                    J.ArrayType arrayType = (J.ArrayType) returnType;
+                    if (arrayType.getElementType() instanceof J.AnnotatedType) {
+                        return service(AnnotationService.class).matches(new Cursor(null, arrayType.getElementType()), new AnnotationMatcher("@" + annotationFqn));
+                    }
+                }
+
+                return false;
+            }
         };
         return Repeat.repeatUntilStable(javaIsoVisitor, 5);
     }

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

@@ -59,14 +59,23 @@ public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, Ex
                             J.MethodDeclaration m2 = m;
                             m2 = m2.withLeadingAnnotations(ListUtils.map(m2.getLeadingAnnotations(),
                                     a -> a == nullable.getTree() ? null : a));
-                            if (m2 != m) {
-                                m2 = m2.withReturnTypeExpression(new J.AnnotatedType(
-                                        Tree.randomId(),
-                                        Space.SINGLE_SPACE,
-                                        Markers.EMPTY,
-                                        singletonList(nullable.getTree().withPrefix(Space.EMPTY)),
-                                        m2.getReturnTypeExpression()
-                                ));
+                            if (m2 != m && m2.getReturnTypeExpression() != null) {
+                                // For array types, annotate the array brackets, not the element type
+                                if (m2.getReturnTypeExpression() instanceof J.ArrayType) {
+                                    // For type-use annotations on arrays (JSpecify style), the annotation should be on the outermost array dimension
+                                    // This creates: elementType @Nullable []
+                                    // We do this by wrapping the ENTIRE array type in an AnnotatedType, but with proper spacing
+                                    m2 = m2.withReturnTypeExpression(((J.ArrayType) m2.getReturnTypeExpression())
+                                            .withAnnotations(singletonList(nullable.getTree().withPrefix(Space.SINGLE_SPACE))));
+                                } else {
+                                    m2 = m2.withReturnTypeExpression(new J.AnnotatedType(
+                                            Tree.randomId(),
+                                            Space.SINGLE_SPACE,
+                                            Markers.EMPTY,
+                                            singletonList(nullable.getTree().withPrefix(Space.EMPTY)),
+                                            m2.getReturnTypeExpression()
+                                    ));
+                                }
                                 m2 = autoFormat(m2, m2.getReturnTypeExpression(), ctx, getCursor().getParentOrThrow());
                                 m2 = m2.withPrefix(m2.getPrefix().withWhitespace(m2.getPrefix().getWhitespace().replace("\n\n\n", "\n\n")));
                             }
@@ -75,6 +84,6 @@ public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, Ex
                         .orElse(m));
             }
         };
-        return Preconditions.check(new UsesType<>("*..Nullable", false),visitor);
+        return Preconditions.check(new UsesType<>("*..Nullable", false), visitor);
     }
 }

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

@@ -518,6 +518,47 @@ public class Test {
         );
     }
 
+    @Test
+    void methodReturnsNullableArray() {
+        rewriteRun(
+          //language=java
+          java(
+            """
+              public class Test {
+
+                  public String[] getArray() {
+                      return null;
+                  }
+
+                  public int[] getIntArray() {
+                      if (System.currentTimeMillis() % 2 == 0) {
+                          return new int[]{1, 2, 3};
+                      }
+                      return null;
+                  }
+              }
+              """,
+            """
+              import org.jspecify.annotations.Nullable;
+
+              public class Test {
+
+                  public String @Nullable[] getArray() {
+                      return null;
+                  }
+
+                  public int @Nullable[] getIntArray() {
+                      if (System.currentTimeMillis() % 2 == 0) {
+                          return new int[]{1, 2, 3};
+                      }
+                      return null;
+                  }
+              }
+              """
+          )
+        );
+    }
+
     @Test
     void typescriptCode() {
         rewriteRun(

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

@@ -131,4 +131,102 @@ String test() {
           )
         );
     }
+
+    @Test
+    void moveNullableToArrayType() {
+        rewriteRun(
+          //language=java
+          java(
+            """
+              import org.openrewrite.internal.lang.Nullable;
+              class Test {
+                  @Nullable
+                  public String[] test() {
+                      return null;
+                  }
+              }
+              """,
+            """
+              import org.openrewrite.internal.lang.Nullable;
+              class Test {
+
+                  public String @Nullable[] test() {
+                      return null;
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Test
+    void moveNullableToPrimitiveArrayType() {
+        rewriteRun(
+          //language=java
+          java(
+            """
+              import org.openrewrite.internal.lang.Nullable;
+              class Test {
+                  @Nullable
+                  public int[] test() {
+                      return null;
+                  }
+              }
+              """,
+            """
+              import org.openrewrite.internal.lang.Nullable;
+              class Test {
+
+                  public int @Nullable[] test() {
+                      return null;
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Test
+    void moveNullableToMultiDimensionalArray() {
+        rewriteRun(
+          //language=java
+          java(
+            """
+              import org.openrewrite.internal.lang.Nullable;
+              class Test {
+                  @Nullable
+                  public String[][] test() {
+                      return null;
+                  }
+              }
+              """,
+            """
+              import org.openrewrite.internal.lang.Nullable;
+              class Test {
+
+                  public String @Nullable[][] test() {
+                      return null;
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Test
+    void noChangeForNullableElements() {
+        rewriteRun(
+          //language=java
+          java(
+            """
+              import org.openrewrite.internal.lang.Nullable;
+              class Test {
+                  public @Nullable String[] test() {
+                      return null;
+                  }
+              }
+              """
+          )
+        );
+    }
 }