JSpecify: handle return types of method references in Java Generics (#847)

Adding new test case and the logic to handle the negative scenario of
the test case where we were getting a false positive earlier.

```Java
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() {
    // We were getting a false positive error on this line as the Nullability of the overriden method was not identified
        A<@Nullable String> p = Test::foo;
    }
}
```

All test cases in NullAwayJSpecifyGenericTests.java passed for these
changes.

---------

Co-authored-by: Manu Sridharan <msridhar@gmail.com>

3 files changed, 126 insertions, 13 deletions

diff --git a/nullaway/src/main/java/com/uber/nullaway/GenericsChecks.java b/nullaway/src/main/java/com/uber/nullaway/GenericsChecks.java
index cfba752..af79f42 100644
--- a/nullaway/src/main/java/com/uber/nullaway/GenericsChecks.java
+++ b/nullaway/src/main/java/com/uber/nullaway/GenericsChecks.java
@@ -702,11 +702,6 @@ public final class GenericsChecks {
   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 @@ public final class GenericsChecks {
     }
   }
 
-  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,
diff --git a/nullaway/src/main/java/com/uber/nullaway/NullAway.java b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
index 3a1f328..5d8dc83 100644
--- a/nullaway/src/main/java/com/uber/nullaway/NullAway.java
+++ b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
@@ -957,7 +957,8 @@ public class NullAway extends BugChecker
     // 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 @@ public class NullAway extends BugChecker
   }
 
   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;
   }
diff --git a/nullaway/src/test/java/com/uber/nullaway/NullAwayJSpecifyGenericsTests.java b/nullaway/src/test/java/com/uber/nullaway/NullAwayJSpecifyGenericsTests.java
index b4d9ff1..dd70683 100644
--- a/nullaway/src/test/java/com/uber/nullaway/NullAwayJSpecifyGenericsTests.java
+++ b/nullaway/src/test/java/com/uber/nullaway/NullAwayJSpecifyGenericsTests.java
@@ -485,6 +485,106 @@ public class NullAwayJSpecifyGenericsTests extends NullAwayTestsBase {
   }
 
   @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()
         .addSourceLines(