Fix reading of JSpecify @Nullable annotations from varargs parameter in bytecode

nullaway/src/main/java/com/uber/nullaway/NullabilityUtil.java

@@ -51,6 +51,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.BiPredicate;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import javax.lang.model.element.AnnotationMirror;
@@ -328,13 +329,13 @@ private static Stream<Attribute.TypeCompound> getTypeUseAnnotations(
       return rawTypeAttributes.filter(
           (t) ->
               t.position.type.equals(TargetType.METHOD_RETURN)
-                  && isDirectTypeUseAnnotation(t, symbol, config));
+                  && (!onlyDirect || isDirectTypeUseAnnotation(t, symbol, config)));
     } else {
       // filter for annotations directly on the type
       return rawTypeAttributes.filter(
           t ->
               targetTypeMatches(symbol, t.position)
-                  && (!onlyDirect || NullabilityUtil.isDirectTypeUseAnnotation(t, symbol, config)));
+                  && (!onlyDirect || isDirectTypeUseAnnotation(t, symbol, config)));
     }
   }
 
@@ -540,22 +541,11 @@ public static <T> T castToNonNull(@Nullable T obj) {
    *     otherwise
    */
   public static boolean isArrayElementNullable(Symbol arraySymbol, Config config) {
-    for (Attribute.TypeCompound t : arraySymbol.getRawTypeAttributes()) {
-      for (TypeAnnotationPosition.TypePathEntry entry : t.position.location) {
-        if (entry.tag == TypeAnnotationPosition.TypePathEntryKind.ARRAY) {
-          if (Nullness.isNullableAnnotation(t.type.toString(), config)) {
-            return true;
-          }
-        }
-      }
-    }
-    // For varargs symbols we also consider the elements to be @Nullable if there is a @Nullable
-    // declaration annotation on the parameter
-    // NOTE this flag check does not work for the varargs parameter of a method defined in bytecodes
-    if ((arraySymbol.flags() & Flags.VARARGS) != 0) {
-      return Nullness.hasNullableDeclarationAnnotation(arraySymbol, config);
-    }
-    return false;
+    return checkArrayElementAnnotations(
+        arraySymbol,
+        config,
+        Nullness::isNullableAnnotation,
+        Nullness::hasNullableDeclarationAnnotation);
   }
 
   /**
@@ -582,12 +572,50 @@ public static boolean nullableVarargsElementsForSourceOrBytecode(
    *     otherwise
    */
   public static boolean isArrayElementNonNull(Symbol arraySymbol, Config config) {
+    return checkArrayElementAnnotations(
+        arraySymbol,
+        config,
+        Nullness::isNonNullAnnotation,
+        Nullness::hasNonNullDeclarationAnnotation);
+  }
+
+  /**
+   * Checks if the given varargs symbol has a {@code @NonNull} annotation for its elements. Works
+   * for both source and bytecode.
+   *
+   * @param varargsSymbol the symbol of the varargs parameter
+   * @param config NullAway configuration
+   * @return true if the varargs symbol has a {@code @NonNull} annotation for its elements, false
+   *     otherwise
+   */
+  public static boolean nonnullVarargsElementsForSourceOrBytecode(
+      Symbol varargsSymbol, Config config) {
+    return isArrayElementNonNull(varargsSymbol, config)
+        || Nullness.hasNonNullDeclarationAnnotation(varargsSymbol, config);
+  }
+
+  /**
+   * Checks if the annotations on the elements of some array symbol satisfy some predicate.
+   *
+   * @param arraySymbol the array symbol
+   * @param config NullAway configuration
+   * @param typeUseCheck the predicate to check the type-use annotations
+   * @param declarationCheck the predicate to check the declaration annotations (applied only to
+   *     varargs symbols)
+   * @return true if the annotations on the elements of the array symbol satisfy the given
+   *     predicates, false otherwise
+   */
+  private static boolean checkArrayElementAnnotations(
+      Symbol arraySymbol,
+      Config config,
+      BiPredicate<String, Config> typeUseCheck,
+      BiPredicate<Symbol, Config> declarationCheck) {
     if (getTypeUseAnnotations(arraySymbol, config, /* onlyDirect= */ false)
         .anyMatch(
             t -> {
               for (TypeAnnotationPosition.TypePathEntry entry : t.position.location) {
                 if (entry.tag == TypeAnnotationPosition.TypePathEntryKind.ARRAY) {
-                  if (Nullness.isNonNullAnnotation(t.type.toString(), config)) {
+                  if (typeUseCheck.test(t.type.toString(), config)) {
                     return true;
                   }
                 }
@@ -596,30 +624,14 @@ public static boolean isArrayElementNonNull(Symbol arraySymbol, Config config) {
             })) {
       return true;
     }
-    // For varargs symbols we also consider the elements to be @NonNull if there is a @NonNull
-    // declaration annotation on the parameter
+    // For varargs symbols we also check for declaration annotations on the parameter
     // NOTE this flag check does not work for the varargs parameter of a method defined in bytecodes
     if ((arraySymbol.flags() & Flags.VARARGS) != 0) {
-      return Nullness.hasNonNullDeclarationAnnotation(arraySymbol, config);
+      return declarationCheck.test(arraySymbol, config);
     }
     return false;
   }
 
-  /**
-   * Checks if the given varargs symbol has a {@code @NonNull} annotation for its elements. Works
-   * for both source and bytecode.
-   *
-   * @param varargsSymbol the symbol of the varargs parameter
-   * @param config NullAway configuration
-   * @return true if the varargs symbol has a {@code @NonNull} annotation for its elements, false
-   *     otherwise
-   */
-  public static boolean nonnullVarargsElementsForSourceOrBytecode(
-      Symbol varargsSymbol, Config config) {
-    return isArrayElementNonNull(varargsSymbol, config)
-        || Nullness.hasNonNullDeclarationAnnotation(varargsSymbol, config);
-  }
-
   /**
    * Does the given symbol have a JetBrains @NotNull declaration annotation? Useful for workarounds
    * in light of https://github.com/uber/NullAway/issues/720

nullaway/src/test/java/com/uber/nullaway/VarargsTests.java

@@ -106,6 +106,22 @@ public void nullableTypeUseVarArgsFromBytecode() {
         .doTest();
   }
 
+  @Test
+  public void nullableVarArgsFromBytecodeJSpecify() {
+    defaultCompilationHelper
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import com.uber.lib.Varargs;",
+            "public class Test {",
+            "  public void testTypeUse() {",
+            "    String x = null;",
+            "    Varargs.typeUseNullableElementsJSpecify(x);",
+            "  }",
+            "}")
+        .doTest();
+  }
+
   @Test
   public void nullableTypeUseVarargs() {
     defaultCompilationHelper

test-java-lib/src/main/java/com/uber/lib/Varargs.java

@@ -7,4 +7,7 @@ public class Varargs {
   public Varargs(@Nullable String... args) {}
 
   public static void typeUse(String @org.jspecify.annotations.Nullable ... args) {}
+
+  public static void typeUseNullableElementsJSpecify(
+      @org.jspecify.annotations.Nullable String... args) {}
 }