Fix bugs in reading varargs annotations from bytecodes

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

@@ -303,7 +303,18 @@
    */
   public static Stream<? extends AnnotationMirror> getTypeUseAnnotations(
       Symbol symbol, Config config) {
-    Stream<Attribute.TypeCompound> rawTypeAttributes = symbol.getRawTypeAttributes().stream();
+    // 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;
+    switch (symbol.getKind()) {
+      case PARAMETER:
+        typeAnnotationOwner = symbol.owner;
+        break;
+      default:
+        typeAnnotationOwner = 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 +324,47 @@
     } else {
       // filter for annotations directly on the type
       return rawTypeAttributes.filter(
-          t -> NullabilityUtil.isDirectTypeUseAnnotation(t, symbol, config));
+          t ->
+              targetTypeMatches(symbol, t.position)
+                  && 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:
+        switch (position.type) {
+          case 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;
+            }
+          default:
+            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 +539,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.
    *
@@ -514,12 +581,28 @@
     }
     // 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,24 @@ public void methodReferenceOnNullableVariable() {
             "}")
         .doTest();
   }
+
+  @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,39 @@ 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);",
+            "  }",
+            "}")
+        .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 +548,33 @@ 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);",
+            "  }",
+            "}")
+        .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,8 @@
+package com.uber.lib.unannotated;
+
+import javax.annotation.Nonnull;
+
+public class RestrictivelyAnnotatedVarargs {
+
+  public static void test(@Nonnull String... args) {}
+}