Fixes for processing return array types and multi-dimensional array types in AstubxGenerator

jdk-annotations/astubx-generator/src/main/java/com/uber/nullaway/jdkannotations/AstubxGenerator.java

@@ -39,16 +39,17 @@
  */
 public class AstubxGenerator {
 
-  /** Used to strip only top-level nullness annotations from a parameter type signature */
-  private static final Pattern TOP_LEVEL_NULLNESS_ANNOTATION_PATTERN =
-      buildTopLevelNullnessAnnotationPattern();
+  /** Used to strip annotations at every depth from a type signature used as an astubx key. */
+  private static final Pattern TYPE_SIGNATURE_ANNOTATION_PATTERN =
+      buildTypeSignatureAnnotationPattern();
 
   /**
-   * Matches annotations immediately before the "[]" for array parameters. Does not properly handle
-   * explicit {@code @NonNull} annotations; see https://github.com/uber/NullAway/issues/1498
+   * Matches {@code @Nullable} on the root array type. javac renders the root as the first {@code
+   * []}, so the pattern must not cross an earlier bracket pair; for example, it matches {@code
+   * String @Nullable [][]} but not {@code String[] @Nullable []}.
    */
   private static final Pattern ARRAY_NULLNESS_ANNOTATION_PATTERN =
-      Pattern.compile("@[\\w.]+(?=\\s*\\[])");
+      Pattern.compile("^[^\\[]*?@(?:org\\.jspecify\\.annotations\\.)?Nullable(?=\\s*\\[])");
 
   /**
    * Matches annotations immediately before the "..." for varargs parameters Does not handle
@@ -248,14 +249,11 @@ private static void getMethodRecords(
       String methodName = method.name();
       // get return type nullness
       String returnType = removeGenericAnnotations(method.returnType());
-      ImmutableSet<String> returnTypeNullness = ImmutableSet.of();
-      // check if return type has Nullable annotation
-      if (returnType.contains("@org.jspecify.annotations.Nullable")) {
-        returnType = returnType.replace("@org.jspecify.annotations.Nullable ", "");
-        returnType = returnType.replaceAll("@Nullable\\s*", "");
-        returnType = returnType.replace(" []", "[]"); // remove whitespace in Array types
-        returnTypeNullness = ImmutableSet.of("Nullable");
-      }
+      ImmutableSet<String> returnTypeNullness =
+          hasTopLevelNullableAnnotation(returnType)
+              ? ImmutableSet.of("Nullable")
+              : ImmutableSet.of();
+      returnType = stripAnnotationsFromTypeSignature(returnType).replace(" []", "[]");
       ImmutableSet.Builder<Integer> nullableTypeParamBuilder = ImmutableSet.builder();
       for (int i = 0; i < method.typeParams().size(); i++) {
         TypeParamInfo typeParam = method.typeParams().get(i);
@@ -279,9 +277,9 @@ private static void getMethodRecords(
         if (hasTopLevelNullableAnnotation(typeSignature)) {
           argAnnotation.put(i, ImmutableSet.of("Nullable"));
         }
-        // Remove top-level annotations before writing the method signature key, while preserving
-        // the varargs ellipsis so the generated key still matches the erased bytecode signature.
-        argumentList[i] = stripTopLevelNullnessAnnotations(typeSignature).replace(" []", "[]");
+        // Remove annotations before writing the method signature key, while preserving the varargs
+        // ellipsis so the generated key still matches the erased bytecode signature.
+        argumentList[i] = stripAnnotationsFromTypeSignature(typeSignature).replace(" []", "[]");
       }
       ImmutableSetMultimap.Builder<Integer, NestedAnnotationInfo> nestedAnnotations =
           new ImmutableSetMultimap.Builder<>();
@@ -380,40 +378,39 @@ private static String removeGenericAnnotations(String typeSignature) {
   }
 
   /**
-   * Checks if the given parameter type has a top-level {@code @Nullable} annotation. Assumes there
-   * are no annotations on any generic type arguments in the type. We only handle JSpecify
+   * Checks if the given type has a top-level {@code @Nullable} annotation. Assumes there are no
+   * annotations on any generic type arguments in the type. We only handle JSpecify
    * {@code @Nullable} annotations for now, as those are the only type present in the JSpecify JDK.
    *
-   * @param parameterType the parameter type.
+   * @param type the type
    * @return true if the type has a top-level {@code @Nullable} annotation, false otherwise
    */
-  private static boolean hasTopLevelNullableAnnotation(String parameterType) {
-    if (!(parameterType.contains("@org.jspecify.annotations.Nullable")
-        || parameterType.contains("@Nullable"))) {
+  private static boolean hasTopLevelNullableAnnotation(String type) {
+    if (!(type.contains("@org.jspecify.annotations.Nullable") || type.contains("@Nullable"))) {
       return false;
     }
-    if (!parameterType.contains("...")) {
-      if (parameterType.contains("[")) {
+    if (!type.contains("...")) {
+      if (type.contains("[")) {
         // Arrays need special handling:
-        //   @Nullable String[]     -> nullable elements, not a nullable array parameter
-        //   String @Nullable []    -> nullable array parameter
+        //   @Nullable String[]     -> nullable elements, not a nullable array
+        //   String @Nullable []    -> nullable array
         // Only the latter is a top-level annotation
-        return ARRAY_NULLNESS_ANNOTATION_PATTERN.matcher(parameterType).find();
+        return ARRAY_NULLNESS_ANNOTATION_PATTERN.matcher(type).find();
       }
       return true;
     }
     // Varargs need special handling:
     //   @Nullable Object...      -> nullable elements, not a nullable array parameter
     //   Object @Nullable ...     -> nullable array parameter
     // Only the latter is a top-level annotation
-    return VARARGS_ARRAY_NULLNESS_ANNOTATION_PATTERN.matcher(parameterType).find();
+    return VARARGS_ARRAY_NULLNESS_ANNOTATION_PATTERN.matcher(type).find();
   }
 
-  private static String stripTopLevelNullnessAnnotations(String typeSignature) {
-    return TOP_LEVEL_NULLNESS_ANNOTATION_PATTERN.matcher(typeSignature).replaceAll("");
+  private static String stripAnnotationsFromTypeSignature(String typeSignature) {
+    return TYPE_SIGNATURE_ANNOTATION_PATTERN.matcher(typeSignature).replaceAll("");
   }
 
-  private static Pattern buildTopLevelNullnessAnnotationPattern() {
+  private static Pattern buildTypeSignatureAnnotationPattern() {
     String annotationWithSpace = "@[\\w.]+\\s";
     // top-level varargs array annotations (for the array itself) are rendered directly before the
     // ellipsis.

jdk-annotations/astubx-generator/src/test/java/com/uber/nullaway/jdkannotations/AstubxTest.java

@@ -322,6 +322,113 @@ public static void bothNullable(@Nullable String @Nullable [] arr) {}
     runTest(expectedMethodRecords, ImmutableMap.of(), ImmutableSet.of("ArrayParameters"));
   }
 
+  @Test
+  public void arrayReturnsNullableArrayVsElements() {
+    compilationHelper
+        .addSourceLines(
+            "ArrayReturns.java",
+            """
+            import org.jspecify.annotations.NullMarked;
+            import org.jspecify.annotations.Nullable;
+            @NullMarked
+            public class ArrayReturns {
+              public static String @Nullable [] arrayNullable() { return null; }
+              public static @Nullable String[] elementsNullable() { return null; }
+              public static @Nullable String @Nullable [] bothNullable() { return null; }
+            }
+            """)
+        .doTest();
+    ImmutableMap<String, MethodAnnotationsRecord> expectedMethodRecords =
+        ImmutableMap.of(
+            "ArrayReturns:java.lang.String[] arrayNullable()",
+            MethodAnnotationsRecord.create(
+                ImmutableSet.of("Nullable"),
+                ImmutableSet.of(),
+                ImmutableMap.of(),
+                ImmutableSetMultimap.of()),
+            "ArrayReturns:java.lang.String[] elementsNullable()",
+            MethodAnnotationsRecord.create(
+                ImmutableSet.of(),
+                ImmutableSet.of(),
+                ImmutableMap.of(),
+                ImmutableSetMultimap.of(
+                    -1,
+                    new NestedAnnotationInfo(
+                        Annotation.NULLABLE,
+                        ImmutableList.of(new TypePathEntry(Kind.ARRAY_ELEMENT, -1))))),
+            "ArrayReturns:java.lang.String[] bothNullable()",
+            MethodAnnotationsRecord.create(
+                ImmutableSet.of("Nullable"),
+                ImmutableSet.of(),
+                ImmutableMap.of(),
+                ImmutableSetMultimap.of(
+                    -1,
+                    new NestedAnnotationInfo(
+                        Annotation.NULLABLE,
+                        ImmutableList.of(new TypePathEntry(Kind.ARRAY_ELEMENT, -1))))));
+    runTest(expectedMethodRecords, ImmutableMap.of(), ImmutableSet.of("ArrayReturns"));
+  }
+
+  @Test
+  public void multidimensionalArrayReturnsNullableArrayVsComponents() {
+    compilationHelper
+        .addSourceLines(
+            "ArrayReturns.java",
+            """
+            import org.jspecify.annotations.NullMarked;
+            import org.jspecify.annotations.Nullable;
+            @NullMarked
+            public class ArrayReturns {
+              public static String @Nullable [][] arrayNullable() { return null; }
+              public static String[] @Nullable [] componentArrayNullable() { return null; }
+              public static String @Nullable [] @Nullable [] bothArraysNullable() { return null; }
+              public static @Nullable String[][] elementsNullable() { return null; }
+            }
+            """)
+        .doTest();
+    ImmutableMap<String, MethodAnnotationsRecord> expectedMethodRecords =
+        ImmutableMap.of(
+            "ArrayReturns:java.lang.String[][] arrayNullable()",
+            MethodAnnotationsRecord.create(
+                ImmutableSet.of("Nullable"),
+                ImmutableSet.of(),
+                ImmutableMap.of(),
+                ImmutableSetMultimap.of()),
+            "ArrayReturns:java.lang.String[][] componentArrayNullable()",
+            MethodAnnotationsRecord.create(
+                ImmutableSet.of(),
+                ImmutableSet.of(),
+                ImmutableMap.of(),
+                ImmutableSetMultimap.of(
+                    -1,
+                    new NestedAnnotationInfo(
+                        Annotation.NULLABLE,
+                        ImmutableList.of(new TypePathEntry(Kind.ARRAY_ELEMENT, -1))))),
+            "ArrayReturns:java.lang.String[][] bothArraysNullable()",
+            MethodAnnotationsRecord.create(
+                ImmutableSet.of("Nullable"),
+                ImmutableSet.of(),
+                ImmutableMap.of(),
+                ImmutableSetMultimap.of(
+                    -1,
+                    new NestedAnnotationInfo(
+                        Annotation.NULLABLE,
+                        ImmutableList.of(new TypePathEntry(Kind.ARRAY_ELEMENT, -1))))),
+            "ArrayReturns:java.lang.String[][] elementsNullable()",
+            MethodAnnotationsRecord.create(
+                ImmutableSet.of(),
+                ImmutableSet.of(),
+                ImmutableMap.of(),
+                ImmutableSetMultimap.of(
+                    -1,
+                    new NestedAnnotationInfo(
+                        Annotation.NULLABLE,
+                        ImmutableList.of(
+                            new TypePathEntry(Kind.ARRAY_ELEMENT, -1),
+                            new TypePathEntry(Kind.ARRAY_ELEMENT, -1))))));
+    runTest(expectedMethodRecords, ImmutableMap.of(), ImmutableSet.of("ArrayReturns"));
+  }
+
   @Test
   public void genericParameter() {
     compilationHelper

jdk-annotations/jdk-integration-test/src/test/java/com/uber/nullaway/jdkannotations/JDKIntegrationTest.java

@@ -313,6 +313,93 @@ void testCall() {
         .doTest();
   }
 
+  @Test
+  public void arrayReturnsNullableArrayVsElements() {
+    compilationHelper
+        .setArgs(
+            JSpecifyJavacConfig.withJSpecifyModeArgs(
+                Arrays.asList(
+                    "-d",
+                    temporaryFolder.getRoot().getAbsolutePath(),
+                    "-XepOpt:NullAway:AnnotatedPackages=com.uber",
+                    "-XepOpt:NullAway:JarInferEnabled=true")))
+        .addSourceLines(
+            "Test.java",
+            """
+            package com.uber;
+            import org.jspecify.annotations.Nullable;
+            import com.uber.nullaway.jdkannotations.ReturnAnnotation;
+            class Test {
+              void test() {
+                // BUG: Diagnostic contains: dereferenced expression 'ReturnAnnotation.returnNullableArray()' is @Nullable
+                int nullableArrayLength = ReturnAnnotation.returnNullableArray().length;
+                String @Nullable [] nullableArray = ReturnAnnotation.returnNullableArray();
+                if (nullableArray != null) {
+                  nullableArray[0].length();
+                }
+
+                int nullableElementsLength = ReturnAnnotation.returnNullableElements().length;
+                // BUG: Diagnostic contains: incompatible types
+                String[] nullableElements = ReturnAnnotation.returnNullableElements();
+
+                // BUG: Diagnostic contains: dereferenced expression 'ReturnAnnotation.returnNullableArrayAndElements()' is @Nullable
+                int nullableArrayAndElementsLength = ReturnAnnotation.returnNullableArrayAndElements().length;
+                // BUG: Diagnostic contains: incompatible types
+                String @Nullable [] nullableArrayAndElements =
+                    ReturnAnnotation.returnNullableArrayAndElements();
+              }
+            }
+            """)
+        .doTest();
+  }
+
+  @Test
+  public void multidimensionalArrayReturnsNullableArrays() {
+    compilationHelper
+        .setArgs(
+            JSpecifyJavacConfig.withJSpecifyModeArgs(
+                Arrays.asList(
+                    "-d",
+                    temporaryFolder.getRoot().getAbsolutePath(),
+                    "-XepOpt:NullAway:AnnotatedPackages=com.uber",
+                    "-XepOpt:NullAway:JarInferEnabled=true")))
+        .addSourceLines(
+            "Test.java",
+            """
+            package com.uber;
+            import org.jspecify.annotations.Nullable;
+            import com.uber.nullaway.jdkannotations.ReturnAnnotation;
+            class Test {
+              void test() {
+                // BUG: Diagnostic contains: dereferenced expression 'ReturnAnnotation.returnNullableOuterArray2D()' is @Nullable
+                int nullableOuterArrayLength = ReturnAnnotation.returnNullableOuterArray2D().length;
+                String @Nullable [][] nullableOuterArray =
+                    ReturnAnnotation.returnNullableOuterArray2D();
+                if (nullableOuterArray != null) {
+                  int componentArrayLength = nullableOuterArray[0].length;
+                }
+
+                int nullableComponentArraysLength =
+                    ReturnAnnotation.returnNullableComponentArrays2D().length;
+                // BUG: Diagnostic contains: incompatible types
+                String[][] nullableComponentArrays =
+                    ReturnAnnotation.returnNullableComponentArrays2D();
+                // TODO: This should report a nullable dereference. NullAway currently loses the
+                // nested library-model annotation on a directly indexed method return.
+                int componentArrayLength =
+                    ReturnAnnotation.returnNullableComponentArrays2D()[0].length;
+
+                // BUG: Diagnostic contains: dereferenced expression 'ReturnAnnotation.returnNullableOuterAndComponentArrays2D()' is @Nullable
+                int nullableOuterAndComponentArraysLength = ReturnAnnotation.returnNullableOuterAndComponentArrays2D().length;
+                // BUG: Diagnostic contains: incompatible types
+                String @Nullable [][] nullableOuterAndComponentArrays =
+                    ReturnAnnotation.returnNullableOuterAndComponentArrays2D();
+              }
+            }
+            """)
+        .doTest();
+  }
+
   @Test
   public void nullableGenericArrayTest() {
     compilationHelper

jdk-annotations/test-annotated/src/main/java/com/uber/nullaway/jdkannotations/ReturnAnnotation.java

@@ -101,6 +101,30 @@ public static Object getNewObjectIfNull(@Nullable Object object) {
     return new @Nullable String[] {"populated", "value", null};
   }
 
+  public static String @Nullable [] returnNullableArray() {
+    return null;
+  }
+
+  public static @Nullable String[] returnNullableElements() {
+    return new @Nullable String[] {"value", null};
+  }
+
+  public static @Nullable String @Nullable [] returnNullableArrayAndElements() {
+    return null;
+  }
+
+  public static String @Nullable [][] returnNullableOuterArray2D() {
+    return null;
+  }
+
+  public static String[] @Nullable [] returnNullableComponentArrays2D() {
+    return new String[1][];
+  }
+
+  public static String @Nullable [] @Nullable [] returnNullableOuterAndComponentArrays2D() {
+    return null;
+  }
+
   @SuppressWarnings({"unchecked", "rawtypes"})
   public static @Nullable List<@Nullable Integer>[] nestedAnnotMixed() {
     // inner list

jdk-annotations/test-unannotated/src/main/java/com/uber/nullaway/jdkannotations/ReturnAnnotation.java

@@ -91,6 +91,30 @@ public static String[] nestedAnnotArrayElement() {
     return new String[] {"populated", "value", null};
   }
 
+  public static String[] returnNullableArray() {
+    return null;
+  }
+
+  public static String[] returnNullableElements() {
+    return new String[] {"value", null};
+  }
+
+  public static String[] returnNullableArrayAndElements() {
+    return null;
+  }
+
+  public static String[][] returnNullableOuterArray2D() {
+    return null;
+  }
+
+  public static String[][] returnNullableComponentArrays2D() {
+    return new String[1][];
+  }
+
+  public static String[][] returnNullableOuterAndComponentArrays2D() {
+    return null;
+  }
+
   @SuppressWarnings({"unchecked", "rawtypes"})
   public static List<Integer>[] nestedAnnotMixed() {
     // inner list