Fix bugs in reading varargs annotations from bytecodes

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

@@ -55,6 +55,7 @@
 import java.util.stream.Stream;
 import javax.lang.model.element.AnnotationMirror;
 import javax.lang.model.element.AnnotationValue;
+import javax.lang.model.element.ElementKind;
 import javax.lang.model.element.ExecutableElement;
 import org.checkerframework.nullaway.javacutil.AnnotationUtils;
 import org.jspecify.annotations.Nullable;
@@ -301,9 +302,27 @@
    * Gets the type use annotations on a symbol, ignoring annotations on components of the type (type
    * arguments, wildcards, etc.)
    */
-  public static Stream<? extends AnnotationMirror> getTypeUseAnnotations(
-      Symbol symbol, Config config) {
-    Stream<Attribute.TypeCompound> rawTypeAttributes = symbol.getRawTypeAttributes().stream();
+  public static Stream<Attribute.TypeCompound> getTypeUseAnnotations(Symbol symbol, Config config) {
+    return getTypeUseAnnotations(symbol, config, /* onlyDirect= */ true);
+  }
+
+  /**
+   * Gets the type use annotations on a symbol
+   *
+   * @param symbol the symbol
+   * @param config NullAway configuration
+   * @param onlyDirect if true, only return annotations that are directly on the type, not on
+   *     components of the type (type arguments, wildcards, array contents, etc.)
+   * @return the type use annotations on the symbol
+   */
+  private static Stream<Attribute.TypeCompound> getTypeUseAnnotations(
+      Symbol symbol, Config config, boolean onlyDirect) {
+    // Adapted from Error Prone's MoreAnnotations class:
+    // https://github.com/google/error-prone/blob/5f71110374e63f3c35b661f538295fa15b5c1db2/check_api/src/main/java/com/google/errorprone/util/MoreAnnotations.java#L84-L91
+    Symbol typeAnnotationOwner =
+        symbol.getKind().equals(ElementKind.PARAMETER) ? symbol.owner : symbol;
+    Stream<Attribute.TypeCompound> rawTypeAttributes =
+        typeAnnotationOwner.getRawTypeAttributes().stream();
     if (symbol instanceof Symbol.MethodSymbol) {
       // for methods, we want annotations on the return type
       return rawTypeAttributes.filter(
@@ -313,7 +332,45 @@
     } else {
       // filter for annotations directly on the type
       return rawTypeAttributes.filter(
-          t -> NullabilityUtil.isDirectTypeUseAnnotation(t, symbol, config));
+          t ->
+              targetTypeMatches(symbol, t.position)
+                  && (!onlyDirect || NullabilityUtil.isDirectTypeUseAnnotation(t, symbol, config)));
+    }
+  }
+
+  // Adapted from Error Prone MoreAnnotations:
+  // https://github.com/google/error-prone/blob/5f71110374e63f3c35b661f538295fa15b5c1db2/check_api/src/main/java/com/google/errorprone/util/MoreAnnotations.java#L128
+  private static boolean targetTypeMatches(Symbol sym, TypeAnnotationPosition position) {
+    switch (sym.getKind()) {
+      case LOCAL_VARIABLE:
+        return position.type == TargetType.LOCAL_VARIABLE;
+      case FIELD:
+      // treated like a field
+      case ENUM_CONSTANT:
+        return position.type == TargetType.FIELD;
+      case CONSTRUCTOR:
+      case METHOD:
+        return position.type == TargetType.METHOD_RETURN;
+      case PARAMETER:
+        if (position.type.equals(TargetType.METHOD_FORMAL_PARAMETER)) {
+          int parameterIndex = position.parameter_index;
+          if (position.onLambda != null) {
+            com.sun.tools.javac.util.List<JCTree.JCVariableDecl> lambdaParams =
+                position.onLambda.params;
+            return parameterIndex < lambdaParams.size()
+                && lambdaParams.get(parameterIndex).sym.equals(sym);
+          } else {
+            return ((Symbol.MethodSymbol) sym.owner).getParameters().indexOf(sym) == parameterIndex;
+          }
+        } else {
+          return false;
+        }
+      case CLASS:
+        // There are no type annotations on the top-level type of the class being declared, only
+        // on other types in the signature (e.g. `class Foo extends Bar<@A Baz> {}`).
+        return false;
+      default:
+        throw new AssertionError("unsupported element kind " + sym.getKind() + " symbol " + sym);
     }
   }
 
@@ -488,12 +545,28 @@
     }
     // 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;
   }
 
+  /**
+   * Checks if the given varargs symbol has a {@code @Nullable} 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 @Nullable} annotation for its elements, false
+   *     otherwise
+   */
+  public static boolean nullableVarargsElementsForSourceOrBytecode(
+      Symbol varargsSymbol, Config config) {
+    return isArrayElementNullable(varargsSymbol, config)
+        || Nullness.hasNullableDeclarationAnnotation(varargsSymbol, config);
+  }
+
   /**
    * Checks if the given array symbol has a {@code @NonNull} annotation for its elements.
    *
@@ -503,23 +576,44 @@
    *     otherwise
    */
   public static boolean isArrayElementNonNull(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.isNonNullAnnotation(t.type.toString(), config)) {
-            return true;
-          }
-        }
-      }
+    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)) {
+                    return true;
+                  }
+                }
+              }
+              return false;
+            })) {
+      return true;
     }
     // For varargs symbols we also consider the elements to be @NonNull if there is a @NonNull
     // 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.hasNonNullDeclarationAnnotation(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/main/java/com/uber/nullaway/Nullness.java

@@ -227,9 +227,8 @@ public static boolean paramHasNullableAnnotation(
     if (symbol.isVarArgs()
         && paramInd == symbol.getParameters().size() - 1
         && !config.isLegacyAnnotationLocation()) {
-      // individual arguments passed in the varargs positions can be @Nullable if the array element
-      // type of the parameter is @Nullable
-      return NullabilityUtil.isArrayElementNullable(symbol.getParameters().get(paramInd), config);
+      return NullabilityUtil.nullableVarargsElementsForSourceOrBytecode(
+          symbol.getParameters().get(paramInd), config);
     } else {
       return hasNullableAnnotation(
           NullabilityUtil.getAllAnnotationsForParameter(symbol, paramInd, config), config);
@@ -269,9 +268,8 @@ public static boolean paramHasNonNullAnnotation(
     if (symbol.isVarArgs()
         && paramInd == symbol.getParameters().size() - 1
         && !config.isLegacyAnnotationLocation()) {
-      // individual arguments passed in the varargs positions must be @NonNull if the array element
-      // type of the parameter is @NonNull
-      return NullabilityUtil.isArrayElementNonNull(symbol.getParameters().get(paramInd), config);
+      return NullabilityUtil.nonnullVarargsElementsForSourceOrBytecode(
+          symbol.getParameters().get(paramInd), config);
     } else {
       return hasNonNullAnnotation(
           NullabilityUtil.getAllAnnotationsForParameter(symbol, paramInd, config), config);

nullaway/src/main/java/com/uber/nullaway/dataflow/CoreNullnessStoreInitializer.java

@@ -67,6 +67,8 @@ private static NullnessStore methodInitialStore(
     for (LocalVariableNode param : parameters) {
       Symbol paramSymbol = (Symbol) param.getElement();
       Nullness assumed;
+      // Using this flag to check for a varargs parameter works since we know paramSymbol represents
+      // a parameter defined in source code
       if ((paramSymbol.flags() & Flags.VARARGS) != 0) {
         assumed = Nullness.varargsArrayIsNullable(paramSymbol, config) ? NULLABLE : NONNULL;
       } else {

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

@@ -47,4 +47,25 @@ public void methodReferenceOnNullableVariable() {
             "}")
         .doTest();
   }
+
+  /** test that we can properly read an explicit type-use annotation on a lambda parameter */
+  @Test
+  public void testNullableLambdaParamTypeUse() {
+    defaultCompilationHelper
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import org.jspecify.annotations.Nullable;",
+            "class Test {",
+            "    @FunctionalInterface",
+            "    interface NullableParamFunctionTypeUse<T, U> {",
+            "      U takeVal(@Nullable T x);",
+            "    }",
+            "    static void testParamTypeUse() {",
+            "      NullableParamFunctionTypeUse n3 = (@Nullable Object x) -> (x == null) ? \"null\" : x.toString();",
+            "      NullableParamFunctionTypeUse n4 = (x) -> (x == null) ? \"null\" : x.toString();",
+            "    }",
+            "}")
+        .doTest();
+  }
 }

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

@@ -71,6 +71,41 @@ public void testNullableVarargs() {
         .doTest();
   }
 
+  /** Test for a @Nullable declaration annotation on a varargs parameter defined in bytecode */
+  @Test
+  public void nullableDeclarationVarArgsFromBytecode() {
+    defaultCompilationHelper
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import com.uber.lib.Varargs;",
+            "public class Test {",
+            "  public void testDeclaration() {",
+            "    String x = null;",
+            "    Varargs s = new Varargs(x);",
+            "    String y = \"hello\", z = null;",
+            "    Varargs s2 = new Varargs(y, z);",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  @Test
+  public void nullableTypeUseVarArgsFromBytecode() {
+    defaultCompilationHelper
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import com.uber.lib.Varargs;",
+            "public class Test {",
+            "  public void testTypeUse() {",
+            "    String[] x = null;",
+            "    Varargs.typeUse(x);",
+            "  }",
+            "}")
+        .doTest();
+  }
+
   @Test
   public void nullableTypeUseVarargs() {
     defaultCompilationHelper
@@ -515,4 +550,37 @@ public void testVarargsRestrictive() {
             "}")
         .doTest();
   }
+
+  /**
+   * Test for a restrictive @NonNull declaration annotation on a varargs parameter defined in
+   * bytecode
+   */
+  @Test
+  public void testVarargsRestrictiveBytecodes() {
+    makeTestHelperWithArgs(
+            Arrays.asList(
+                "-d",
+                temporaryFolder.getRoot().getAbsolutePath(),
+                "-XepOpt:NullAway:AnnotatedPackages=com.uber",
+                "-XepOpt:NullAway:UnannotatedSubPackages=com.uber.lib.unannotated",
+                "-XepOpt:NullAway:AcknowledgeRestrictiveAnnotations=true"))
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import com.uber.lib.unannotated.RestrictivelyAnnotatedVarargs;",
+            "public class Test {",
+            "  public void testDeclaration() {",
+            "    String x = null;",
+            "    String[] y = null;",
+            "    // BUG: Diagnostic contains: passing @Nullable parameter 'x'",
+            "    RestrictivelyAnnotatedVarargs.test(x);",
+            "    RestrictivelyAnnotatedVarargs.test(y);",
+            "    // BUG: Diagnostic contains: passing @Nullable parameter 'x'",
+            "    RestrictivelyAnnotatedVarargs.testTypeUse(x);",
+            "    // TODO should report an error; requires a fuller fix for #1027",
+            "    RestrictivelyAnnotatedVarargs.testTypeUse(y);",
+            "  }",
+            "}")
+        .doTest();
+  }
 }

test-java-lib/build.gradle

@@ -22,6 +22,7 @@ plugins {
 
 dependencies {
     annotationProcessor project(":nullaway")
+    implementation deps.test.jsr305Annotations
     implementation deps.build.jspecify
 
     compileOnly deps.build.javaxValidation

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

@@ -0,0 +1,10 @@
+package com.uber.lib;
+
+import javax.annotation.Nullable;
+
+public class Varargs {
+
+  public Varargs(@Nullable String... args) {}
+
+  public static void typeUse(String @org.jspecify.annotations.Nullable ... args) {}
+}

test-java-lib/src/main/java/com/uber/lib/unannotated/RestrictivelyAnnotatedVarargs.java

@@ -0,0 +1,11 @@
+package com.uber.lib.unannotated;
+
+import javax.annotation.Nonnull;
+
+public class RestrictivelyAnnotatedVarargs {
+
+  public static void test(@Nonnull String... args) {}
+
+  public static void testTypeUse(
+      @org.jspecify.annotations.NonNull String @org.jspecify.annotations.NonNull ... args) {}
+}