From 026970c993b5b64f2e105bd01e75915fbadc87e2 Mon Sep 17 00:00:00 2001
From: Manu Sridharan <msridhar@gmail.com>
Date: Tue, 23 Jan 2024 10:17:18 -0800
Subject: [PATCH] more fixes, still needs cleanup

---
 .../src/main/java/com/uber/nullaway/NullAway.java  |  8 +++++---
 .../com/uber/nullaway/generics/GenericsChecks.java |  7 +++----
 .../nullaway/NullAwayJSpecifyGenericsTests.java    | 14 +++++++++-----
 3 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/nullaway/src/main/java/com/uber/nullaway/NullAway.java b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
index 552735ddcc..352f09810c 100644
--- a/nullaway/src/main/java/com/uber/nullaway/NullAway.java
+++ b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
@@ -905,7 +905,7 @@ private Description checkReturnExpression(
       if (GenericsChecks.getGenericMethodReturnTypeNullness(
                   methodSymbol, ASTHelpers.getType(lambdaTree), state, config)
               .equals(Nullness.NULLABLE)
-          || GenericsChecks.passingLambdaWithGenericReturnToUnmarkedCode(
+          || GenericsChecks.passingLambdaOrMethodRefWithGenericReturnToUnmarkedCode(
               methodSymbol, lambdaTree, state, config, codeAnnotationInfo)) {
         // In JSpecify mode, the return type of a lambda may be @Nullable via a type argument
         return Description.NO_MATCH;
@@ -1052,8 +1052,10 @@ private boolean overriddenMethodReturnsNonNull(
         // 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);
+                    overriddenMethod, ASTHelpers.getType(memberReferenceTree), state, config)
+                .equals(Nullness.NONNULL)
+            && !GenericsChecks.passingLambdaOrMethodRefWithGenericReturnToUnmarkedCode(
+                overriddenMethod, memberReferenceTree, state, config, codeAnnotationInfo);
       } else {
         // Use the enclosing class of the overriding method to find generic type arguments
         return GenericsChecks.getGenericMethodReturnTypeNullness(
diff --git a/nullaway/src/main/java/com/uber/nullaway/generics/GenericsChecks.java b/nullaway/src/main/java/com/uber/nullaway/generics/GenericsChecks.java
index 4427bc4259..4d34dab7e0 100644
--- a/nullaway/src/main/java/com/uber/nullaway/generics/GenericsChecks.java
+++ b/nullaway/src/main/java/com/uber/nullaway/generics/GenericsChecks.java
@@ -13,7 +13,6 @@
 import com.sun.source.tree.AssignmentTree;
 import com.sun.source.tree.ConditionalExpressionTree;
 import com.sun.source.tree.ExpressionTree;
-import com.sun.source.tree.LambdaExpressionTree;
 import com.sun.source.tree.MemberSelectTree;
 import com.sun.source.tree.MethodInvocationTree;
 import com.sun.source.tree.MethodTree;
@@ -846,9 +845,9 @@ public static String prettyTypeForError(Type type, VisitorState state) {
     return type.accept(new GenericTypePrettyPrintingVisitor(state), null);
   }
 
-  public static boolean passingLambdaWithGenericReturnToUnmarkedCode(
+  public static boolean passingLambdaOrMethodRefWithGenericReturnToUnmarkedCode(
       Symbol.MethodSymbol methodSymbol,
-      LambdaExpressionTree lambdaTree,
+      ExpressionTree expressionTree,
       VisitorState state,
       Config config,
       CodeAnnotationInfo codeAnnotationInfo) {
@@ -859,7 +858,7 @@ public static boolean passingLambdaWithGenericReturnToUnmarkedCode(
     }
     boolean callingUnannotated = false;
     TreePath path = state.getPath();
-    while (path != null && !path.getLeaf().equals(lambdaTree)) {
+    while (path != null && !path.getLeaf().equals(expressionTree)) {
       path = path.getParentPath();
     }
     verify(path != null, "did not find lambda tree in TreePath");
diff --git a/nullaway/src/test/java/com/uber/nullaway/NullAwayJSpecifyGenericsTests.java b/nullaway/src/test/java/com/uber/nullaway/NullAwayJSpecifyGenericsTests.java
index ed7873c64a..0958740d5e 100644
--- a/nullaway/src/test/java/com/uber/nullaway/NullAwayJSpecifyGenericsTests.java
+++ b/nullaway/src/test/java/com/uber/nullaway/NullAwayJSpecifyGenericsTests.java
@@ -1665,12 +1665,16 @@ public void testForNullReturnLambdaFromStreamMap() {
             "class Test {",
             "  @Nullable",
             "  static Integer foo(){",
-            "   return null;",
-            "   }",
+            "    return null;",
+            "  }",
+            "  @Nullable",
+            "  static Integer foo2(Integer i){",
+            "    return null;",
+            "  }",
             "  static void testNegative() {",
-            "   List<Integer> numbers = Arrays.asList(1, 2, 3, 4, 5);",
-            "   List<Integer> doubledNumbers = numbers.stream().map(number -> foo()).collect(Collectors.toList());",
-            // TODO need the same test but with a method reference
+            "    List<Integer> numbers = Arrays.asList(1, 2, 3, 4, 5);",
+            "    List<Integer> doubledNumbers = numbers.stream().map(number -> foo()).collect(Collectors.toList());",
+            "    List<Integer> other = numbers.stream().map(Test::foo2).collect(Collectors.toList());",
             "  }",
             "}")
         .doTest();