JSpecify: handle Nullability for return types of lambda expressions for Generic Types. (#854)

```java
class Test {
    interface A<T1 extends @Nullable Object> {
        T1 function(Object o);
    }
    static void testPositive() {
        // BUG: Diagnostic contains: returning @Nullable expression from method with @NonNull return type
        A<String> p = x -> null;
    }
    static void testNegative() {
        // We were reporting a false positive here since the indirect Nullability through the generic was not handled correctly.
        A<@Nullable String> p = x -> null;
    }
}
```
For the negative scenario mentioned above where we do not expect an
exception, we were getting a false positive. After computing the
Nullability correctly when it is indirectly applied through the generic
type we no longer get this false positive.

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

Similar changes for handling lambda parameters were made under
https://github.com/uber/NullAway/pull/852

---------

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

2 files changed, 69 insertions, 9 deletions

diff --git a/nullaway/src/main/java/com/uber/nullaway/NullAway.java b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
index d4ce4e8..ff4ee80 100644
--- a/nullaway/src/main/java/com/uber/nullaway/NullAway.java
+++ b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
@@ -380,16 +380,16 @@ public class NullAway extends BugChecker
     }
     Tree leaf = enclosingMethodOrLambda.getLeaf();
     Symbol.MethodSymbol methodSymbol;
+    LambdaExpressionTree lambdaTree = null;
     if (leaf instanceof MethodTree) {
       MethodTree enclosingMethod = (MethodTree) leaf;
       methodSymbol = ASTHelpers.getSymbol(enclosingMethod);
     } else {
       // we have a lambda
-      methodSymbol =
-          NullabilityUtil.getFunctionalInterfaceMethod(
-              (LambdaExpressionTree) leaf, state.getTypes());
+      lambdaTree = (LambdaExpressionTree) leaf;
+      methodSymbol = NullabilityUtil.getFunctionalInterfaceMethod(lambdaTree, state.getTypes());
     }
-    return checkReturnExpression(tree, retExpr, methodSymbol, state);
+    return checkReturnExpression(retExpr, methodSymbol, lambdaTree, tree, state);
   }
 
   @Override
@@ -853,8 +853,26 @@ public class NullAway extends BugChecker
         methodSymbol, state, isMethodAnnotated, methodReturnNullness);
   }
 
+  /**
+   * Checks that if a returned expression is {@code @Nullable}, the enclosing method does not have a
+   * {@code @NonNull} return type. Also performs an unboxing check on the returned expression.
+   * Finally, in JSpecify mode, also checks that the nullability of generic type arguments of the
+   * returned expression's type match the method return type.
+   *
+   * @param retExpr the expression being returned
+   * @param methodSymbol symbol for the enclosing method
+   * @param lambdaTree if return is inside a lambda, the tree for the lambda, otherwise {@code null}
+   * @param errorTree tree on which to report an error if needed
+   * @param state the visitor state
+   * @return {@link Description} of the returning {@code @Nullable} from {@code @NonNull} method
+   *     error if one is to be reported, otherwise {@link Description#NO_MATCH}
+   */
   private Description checkReturnExpression(
-      Tree tree, ExpressionTree retExpr, Symbol.MethodSymbol methodSymbol, VisitorState state) {
+      ExpressionTree retExpr,
+      Symbol.MethodSymbol methodSymbol,
+      @Nullable LambdaExpressionTree lambdaTree,
+      Tree errorTree,
+      VisitorState state) {
     Type returnType = methodSymbol.getReturnType();
     if (returnType.isPrimitive()) {
       // check for unboxing
@@ -867,20 +885,33 @@ public class NullAway extends BugChecker
       // support)
       return Description.NO_MATCH;
     }
-    // Do the generics checks here, since we need to check the type arguments regardless of the
-    // top-level nullability of the return type
+
+    // Check generic type arguments for returned expression here, since we need to check the type
+    // arguments regardless of the top-level nullability of the return type
     GenericsChecks.checkTypeParameterNullnessForFunctionReturnType(
         retExpr, methodSymbol, this, state);
+
+    // Now, perform the check for returning @Nullable from @NonNull.  First, we check if the return
+    // type is @Nullable, and if so, bail out.
     if (getMethodReturnNullness(methodSymbol, state, Nullness.NULLABLE).equals(Nullness.NULLABLE)) {
       return Description.NO_MATCH;
+    } else if (config.isJSpecifyMode()
+        && lambdaTree != null
+        && GenericsChecks.getGenericMethodReturnTypeNullness(
+                methodSymbol, ASTHelpers.getType(lambdaTree), state, config)
+            .equals(Nullness.NULLABLE)) {
+      // In JSpecify mode, the return type of a lambda may be @Nullable via a type argument
+      return Description.NO_MATCH;
     }
+
+    // Return type is @NonNull.  Check if the expression is @Nullable
     if (mayBeNullExpr(state, retExpr)) {
       return errorBuilder.createErrorDescriptionForNullAssignment(
           new ErrorMessage(
               MessageTypes.RETURN_NULLABLE,
               "returning @Nullable expression from method with @NonNull return type"),
           retExpr,
-          buildDescription(tree),
+          buildDescription(errorTree),
           state,
           methodSymbol);
     }
@@ -916,7 +947,7 @@ public class NullAway extends BugChecker
     if (tree.getBodyKind() == LambdaExpressionTree.BodyKind.EXPRESSION
         && funcInterfaceMethod.getReturnType().getKind() != TypeKind.VOID) {
       ExpressionTree resExpr = (ExpressionTree) tree.getBody();
-      return checkReturnExpression(tree, resExpr, funcInterfaceMethod, state);
+      return checkReturnExpression(resExpr, funcInterfaceMethod, tree, tree, state);
     }
     return Description.NO_MATCH;
   }
diff --git a/nullaway/src/test/java/com/uber/nullaway/NullAwayJSpecifyGenericsTests.java b/nullaway/src/test/java/com/uber/nullaway/NullAwayJSpecifyGenericsTests.java
index 73737b8..e722706 100644
--- a/nullaway/src/test/java/com/uber/nullaway/NullAwayJSpecifyGenericsTests.java
+++ b/nullaway/src/test/java/com/uber/nullaway/NullAwayJSpecifyGenericsTests.java
@@ -651,6 +651,35 @@ public class NullAwayJSpecifyGenericsTests extends NullAwayTestsBase {
   }
 
   @Test
+  public void testForLambdaReturnTypeInAnAssignment() {
+    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 void testPositive1() {",
+            "    // BUG: Diagnostic contains: returning @Nullable expression from method with @NonNull return type",
+            "    A<String> p = x -> null;",
+            "  }",
+            "  static void testPositive2() {",
+            "    // BUG: Diagnostic contains: returning @Nullable expression from method with @NonNull return type",
+            "    A<String> p = x -> { return null; };",
+            "  }",
+            "  static void testNegative1() {",
+            "    A<@Nullable String> p = x -> null;",
+            "  }",
+            "  static void testNegative2() {",
+            "    A<@Nullable String> p = x -> { return null; };",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  @Test
   public void testForDiamondInAnAssignment() {
     makeHelper()
         .addSourceLines(