JSpecify: handle return types of method references in Java Generics

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

@@ -702,11 +702,6 @@ public static void checkTypeParameterNullnessForMethodOverriding(
   public static Nullness getGenericMethodReturnTypeNullness(
       Symbol.MethodSymbol method, Symbol enclosingSymbol, VisitorState state, Config config) {
     Type enclosingType = getTypeForSymbol(enclosingSymbol, state);
-    if (enclosingType == null) {
-      // we have no additional information from generics, so return NONNULL (presence of a @Nullable
-      // annotation should have been handled by the caller)
-      return Nullness.NONNULL;
-    }
     return getGenericMethodReturnTypeNullness(method, enclosingType, state, config);
   }
 
@@ -738,8 +733,13 @@ private static Type getTypeForSymbol(Symbol symbol, VisitorState state) {
     }
   }
 
-  private static Nullness getGenericMethodReturnTypeNullness(
-      Symbol.MethodSymbol method, Type enclosingType, VisitorState state, Config config) {
+  static Nullness getGenericMethodReturnTypeNullness(
+      Symbol.MethodSymbol method, @Nullable Type enclosingType, VisitorState state, Config config) {
+    if (enclosingType == null) {
+      // we have no additional information from generics, so return NONNULL (presence of a @Nullable
+      // annotation should have been handled by the caller)
+      return Nullness.NONNULL;
+    }
     Type overriddenMethodType = state.getTypes().memberType(enclosingType, method);
     verify(
         overriddenMethodType instanceof ExecutableType,

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

@@ -957,7 +957,8 @@ private Description checkOverriding(
     // if the super method returns nonnull, overriding method better not return nullable
     // Note that, for the overriding method, the permissive default is non-null,
     // but it's nullable for the overridden one.
-    if (overriddenMethodReturnsNonNull(overriddenMethod, overridingMethod.owner, state)
+    if (overriddenMethodReturnsNonNull(
+            overriddenMethod, overridingMethod.owner, memberReferenceTree, state)
         && getMethodReturnNullness(overridingMethod, state, Nullness.NONNULL)
             .equals(Nullness.NULLABLE)
         && (memberReferenceTree == null
@@ -996,18 +997,30 @@ && getMethodReturnNullness(overridingMethod, state, Nullness.NONNULL)
   }
 
   private boolean overriddenMethodReturnsNonNull(
-      Symbol.MethodSymbol overriddenMethod, Symbol enclosingSymbol, VisitorState state) {
+      Symbol.MethodSymbol overriddenMethod,
+      Symbol enclosingSymbol,
+      @Nullable MemberReferenceTree memberReferenceTree,
+      VisitorState state) {
     Nullness methodReturnNullness =
         getMethodReturnNullness(overriddenMethod, state, Nullness.NULLABLE);
     if (!methodReturnNullness.equals(Nullness.NONNULL)) {
       return false;
     }
     // In JSpecify mode, for generic methods, we additionally need to check the return nullness
-    // using the type parameters from the type enclosing the overriding method
+    // using the type arguments from the type enclosing the overriding method
     if (config.isJSpecifyMode()) {
-      return GenericsChecks.getGenericMethodReturnTypeNullness(
-              overriddenMethod, enclosingSymbol, state, config)
-          .equals(Nullness.NONNULL);
+      if (memberReferenceTree != null) {
+        // For a method reference, we get generic type arguments from javac's inferred type for the
+        // tree, which properly preserves type-use annotations
+        return GenericsChecks.getGenericMethodReturnTypeNullness(
+                overriddenMethod, ASTHelpers.getType(memberReferenceTree), state, config)
+            .equals(Nullness.NONNULL);
+      } else {
+        // Use the enclosing class of the overriding method to find generic type arguments
+        return GenericsChecks.getGenericMethodReturnTypeNullness(
+                overriddenMethod, enclosingSymbol, state, config)
+            .equals(Nullness.NONNULL);
+      }
     }
     return true;
   }

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

@@ -484,6 +484,106 @@ public void testForMethodReferenceInAnAssignment() {
         .doTest();
   }
 
+  @Test
+  public void testForMethodReferenceForClassFieldAssignment() {
+    makeHelper()
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import org.jspecify.annotations.Nullable;",
+            "class Test {",
+            "  interface A<T1 extends @Nullable Object> {",
+            "    T1 function(Object o);",
+            "  }",
+            "  static @Nullable String foo(Object o) {",
+            "    return o.toString();",
+            "  }",
+            "  // BUG: Diagnostic contains: referenced method returns @Nullable",
+            "  A<String> positiveField = Test::foo;",
+            "  A<@Nullable String> negativeField = Test::foo;",
+            "}")
+        .doTest();
+  }
+
+  @Test
+  public void testForMethodReferenceReturnTypeInAnAssignment() {
+    makeHelper()
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import org.jspecify.annotations.Nullable;",
+            "class Test {",
+            "  interface A<T1 extends @Nullable Object> {",
+            "    T1 function(Object o);",
+            "  }",
+            "  static @Nullable String foo(Object o) {",
+            "    return o.toString();",
+            "  }",
+            "  static void testPositive() {",
+            "    // BUG: Diagnostic contains: referenced method returns @Nullable",
+            "    A<String> p = Test::foo;",
+            "  }",
+            "  static void testNegative() {",
+            "    A<@Nullable String> p = Test::foo;",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  @Test
+  public void testForMethodReferenceWhenReturned() {
+    makeHelper()
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import org.jspecify.annotations.Nullable;",
+            "class Test {",
+            "  interface A<T1 extends @Nullable Object> {",
+            "    T1 function(Object o);",
+            "  }",
+            "  static @Nullable String foo(Object o) {",
+            "    return o.toString();",
+            "  }",
+            "  static A<String> testPositive() {",
+            "    // BUG: Diagnostic contains: referenced method returns @Nullable",
+            "    return Test::foo;",
+            "  }",
+            "  static A<@Nullable String> testNegative() {",
+            "    return Test::foo;",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  @Test
+  public void testForMethodReferenceAsMethodParameter() {
+    makeHelper()
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import org.jspecify.annotations.Nullable;",
+            "class Test {",
+            "  interface A<T1 extends @Nullable Object> {",
+            "    T1 function(Object o);",
+            "  }",
+            "  static @Nullable String foo(Object o) {",
+            "    return o.toString();",
+            "  }",
+            "  static void fooPositive(A<String> a) {",
+            "  }",
+            "  static void fooNegative(A<@Nullable String> a) {",
+            "  }",
+            "  static void testPositive() {",
+            "    // BUG: Diagnostic contains: referenced method returns @Nullable",
+            "    fooPositive(Test::foo);",
+            "  }",
+            "  static void testNegative() {",
+            "    fooNegative(Test::foo);",
+            "  }",
+            "}")
+        .doTest();
+  }
+
   @Test
   public void testForLambdasInAnAssignment() {
     makeHelper()