Fix type parameter nullability matching for covariant types

nullaway/src/main/java/com/uber/nullaway/generics/CheckIdenticalNullabilityVisitor.java

@@ -13,17 +13,27 @@
  * Visitor that checks for identical nullability annotations at all nesting levels within two types.
  * Compares the Type it is called upon, i.e. the LHS type and the Type passed as an argument, i.e.
  * The RHS type.
+ *
+ * <p>In covariant mode, the check only fails if the RHS type argument is @Nullable but the LHS type
+ * argument is not. The reverse (LHS @Nullable, RHS non-null) is allowed.
  */
 public class CheckIdenticalNullabilityVisitor extends Types.DefaultTypeVisitor<Boolean, Type> {
   private final VisitorState state;
   private final GenericsChecks genericsChecks;
   private final Config config;
+  private final boolean covariant;
 
   CheckIdenticalNullabilityVisitor(
       VisitorState state, GenericsChecks genericsChecks, Config config) {
+    this(state, genericsChecks, config, false);
+  }
+
+  CheckIdenticalNullabilityVisitor(
+      VisitorState state, GenericsChecks genericsChecks, Config config, boolean covariant) {
     this.state = state;
     this.genericsChecks = genericsChecks;
     this.config = config;
+    this.covariant = covariant;
   }
 
   @Override
@@ -70,7 +80,9 @@ public Boolean visitClassType(Type.ClassType lhsType, Type rhsType) {
       }
       boolean isLHSNullableAnnotated = genericsChecks.isNullableAnnotated(lhsTypeArgument);
       boolean isRHSNullableAnnotated = genericsChecks.isNullableAnnotated(rhsTypeArgument);
-      if (isLHSNullableAnnotated != isRHSNullableAnnotated) {
+      if (covariant
+          ? (!isLHSNullableAnnotated && isRHSNullableAnnotated)
+          : (isLHSNullableAnnotated != isRHSNullableAnnotated)) {
         return false;
       }
       // nested generics
@@ -106,7 +118,9 @@ public Boolean visitArrayType(Type.ArrayType lhsType, Type rhsType) {
     Type rhsComponentType = rhsArrayType.getComponentType();
     boolean isLHSNullableAnnotated = genericsChecks.isNullableAnnotated(lhsComponentType);
     boolean isRHSNullableAnnotated = genericsChecks.isNullableAnnotated(rhsComponentType);
-    if (isRHSNullableAnnotated != isLHSNullableAnnotated) {
+    if (covariant
+        ? (!isLHSNullableAnnotated && isRHSNullableAnnotated)
+        : (isRHSNullableAnnotated != isLHSNullableAnnotated)) {
       return false;
     }
     return lhsComponentType.accept(this, rhsComponentType);

nullaway/src/main/java/com/uber/nullaway/generics/GenericsChecks.java

@@ -1442,6 +1442,16 @@ private boolean identicalTypeParameterNullability(
     return lhsType.accept(new CheckIdenticalNullabilityVisitor(state, this, config), rhsType);
   }
 
+  /**
+   * Like {@link #identicalTypeParameterNullability(Type, Type, VisitorState)}, but only fails when
+   * the RHS type argument is @Nullable and the LHS type argument is not. Used for method reference
+   * parameter checks where the referenced method may accept a wider type than required.
+   */
+  private boolean covariantTypeParameterNullability(
+      Type lhsType, Type rhsType, VisitorState state) {
+    return lhsType.accept(new CheckIdenticalNullabilityVisitor(state, this, config, true), rhsType);
+  }
+
   /**
    * Like {@link #identicalTypeParameterNullability(Type, Type, VisitorState)}, but allows for
    * covariant array subtyping at the top level.
@@ -1608,7 +1618,11 @@ public void compareGenericTypeParameterNullabilityForCall(
                     memberReferenceTree,
                     state,
                     (subtype, supertype, relationKind) -> {
-                      if (!subtypeParameterNullability(supertype, subtype, state)) {
+                      boolean valid =
+                          relationKind == MethodRefTypeRelationKind.PARAMETER
+                              ? covariantTypeParameterNullability(supertype, subtype, state)
+                              : subtypeParameterNullability(supertype, subtype, state);
+                      if (!valid) {
                         if (relationKind == MethodRefTypeRelationKind.RETURN) {
                           reportInvalidMethodReferenceReturnTypeError(
                               memberReferenceTree, supertype, subtype, state);

nullaway/src/test/java/com/uber/nullaway/jspecify/GenericMethodLambdaOrMethodRefArgTests.java

@@ -516,11 +516,13 @@ static void takeNullableNested(Foo<@Nullable String> f) {}
               void testNegative2(Foo<String> f1, Foo<@Nullable String> f2) {
                 takeMultiple(Test::takeNonNullNested, Test::takeNonNullNested, f1);
                 takeMultiple(Test::takeNullableNested, Test::takeNullableNested, f2);
-              }
-              void testPositive2(Foo<String> f1) {
-                // BUG: Diagnostic contains: parameter type of referenced method is Test.Foo<@Nullable String>, but parameter in functional interface method has type Test.Foo<String>, which has mismatched type parameter nullability
+                // takeNullableNested accepts more permissive type args than required by the FI
                 takeMultiple(Test::takeNonNullNested, Test::takeNullableNested, f1);
               }
+              void testPositive2(Foo<@Nullable String> f2) {
+                // BUG: Diagnostic contains: parameter type of referenced method is Test.Foo<String>, but parameter in functional interface method has type Test.Foo<@Nullable String>, which has mismatched type parameter nullability
+                takeMultiple(Test::takeNonNullNested, Test::takeNullableNested, f2);
+              }
             }
             """)
         .doTest();
@@ -1042,6 +1044,39 @@ static void test() {
         .doTest();
   }
 
+  @Test
+  public void issue1528MethodRefParamNullabilityCovariance() {
+    makeHelperWithInferenceFailureWarning()
+        .addSourceLines(
+            "Test.java",
+            """
+            import org.jspecify.annotations.NullMarked;
+            import org.jspecify.annotations.Nullable;
+            import java.util.Map;
+            import java.util.function.Function;
+            import java.util.stream.Stream;
+            @NullMarked
+            class Test {
+                static <T> void process(Stream<T> stream, Function<T, @Nullable String> mapper) {}
+                static @Nullable String fromEntry(Map.Entry<String, @Nullable String> entry) {
+                    return entry.getValue();
+                }
+                static @Nullable String fromEntryStrict(Map.Entry<String, String> entry) {
+                    return entry.getValue();
+                }
+                static void testOk(Map<String, String> map) {
+                    // referenced method param is more permissive (Entry<@Nullable>) than FI (Entry<String>)
+                    process(map.entrySet().stream(), Test::fromEntry);
+                }
+                static void testBad(Map<String, @Nullable String> map) {
+                    // BUG: Diagnostic contains: parameter type of referenced method
+                    process(map.entrySet().stream(), Test::fromEntryStrict);
+                }
+            }
+            """)
+        .doTest();
+  }
+
   private CompilationTestHelper makeHelperWithInferenceFailureWarning() {
     return makeTestHelperWithArgs(
         JSpecifyJavacConfig.withJSpecifyModeArgs(