Fix bug with library-modeled return types, method references, and streams (#1706)

In `LibraryModelsHandler.onOverrideMethodType` we neglected to apply
top-level `@Nullable` annotations from library models. This impacted
cases like the new `streamMapNullableMapGetMethodReferenceTest`.

While investigating this issue, I realized that the test from #1462 now
passes, along with a similar test that actually uses the library types,
so enable those.

Fixes #1462 

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Bug Fixes**
* Improved handling of nullable return values and top-level annotations
in modeled library methods.
* Enhanced nullability analysis for stream transformations and method
references.

* **Tests**
* Added coverage for detecting unsafe dereferences of nullable stream
results.
* Added coverage for safely filtering nullable map lookups before
collection.
  * Enabled an existing stream nullability test.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: msridhar

nullaway/src/main/java/com/uber/nullaway/handlers/LibraryModelsHandler.java

@@ -498,10 +498,10 @@ public boolean isSingleArgNullImpliesFalseMethod(
   }
 
   /**
-   * Updates method types based on top-level parameter and nested annotation library models. For
-   * now, this method is only used in JSpecify mode, primarily for its handling of nested
-   * annotations. Outside of JSpecify mode, other handler methods are available for reasoning about
-   * top-level annotations only.
+   * Updates method types based on top-level parameter and return annotations and nested annotation
+   * library models. For now, this method is only used in JSpecify mode, primarily for its handling
+   * of nested annotations. Outside of JSpecify mode, other handler methods are available for
+   * reasoning about top-level annotations only.
    */
   @Override
   @SuppressWarnings({"ReferenceEquality", "TypeEquals"})
@@ -516,9 +516,17 @@ public Type.MethodType onOverrideMethodType(
         config.isJSpecifyMode()
             ? optimizedLibraryModels.explicitlyNullableParameters(methodSymbol)
             : ImmutableSet.of();
+    boolean isMethodUnannotated =
+        getCodeAnnotationInfo(state.context)
+            .isSymbolUnannotated(methodSymbol, this.config, mainHandler);
+    boolean modeledNullableReturn =
+        optimizedLibraryModels.hasNullableReturn(
+            methodSymbol, state.getTypes(), isMethodUnannotated);
     ImmutableSetMultimap<Integer, NestedAnnotationInfo> nestedAnnotations =
         optimizedLibraryModels.nestedAnnotationsForMethods(methodSymbol);
-    if (explicitlyNullableParameters.isEmpty() && nestedAnnotations.isEmpty()) {
+    if (explicitlyNullableParameters.isEmpty()
+        && !modeledNullableReturn
+        && nestedAnnotations.isEmpty()) {
       return methodType;
     }
     // update argument types, tracking if anything changed
@@ -546,10 +554,13 @@ public Type.MethodType onOverrideMethodType(
     // update return type
     Type returnType = methodType.restype;
     ImmutableSet<NestedAnnotationInfo> returnAnnotations = nestedAnnotations.get(-1);
-    Type updatedReturnType =
-        returnAnnotations.isEmpty()
-            ? returnType
-            : applyNestedAnnotations(returnType, returnAnnotations, state);
+    Type updatedReturnType = returnType;
+    if (modeledNullableReturn) {
+      updatedReturnType = applyTopLevelNullableAnnotation(updatedReturnType, state);
+    }
+    if (!returnAnnotations.isEmpty()) {
+      updatedReturnType = applyNestedAnnotations(updatedReturnType, returnAnnotations, state);
+    }
     if (updatedReturnType != returnType) {
       changed = true;
     }

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

@@ -160,6 +160,57 @@ void test() {
         .doTest();
   }
 
+  @Test
+  public void streamMapNullableTest() {
+    makeHelper()
+        .addSourceLines(
+            "Test.java",
+            """
+            import org.jspecify.annotations.*;
+            import java.util.*;
+            @NullMarked
+            class Test {
+                static @Nullable String mapToNull(String s) {
+                    return null;
+                }
+                static void test(List<String> list) {
+                    list.stream().map(Test::mapToNull).forEach(s -> {
+                        // BUG: Diagnostic contains: dereferenced expression 's' is @Nullable
+                        s.hashCode();
+                    });
+                }
+            }""")
+        .doTest();
+  }
+
+  @Test
+  public void streamMapNullableMapGetMethodReferenceTest() {
+    makeHelper()
+        .addSourceLines(
+            "Test.java",
+            """
+            import org.jspecify.annotations.NullMarked;
+            import java.util.List;
+            import java.util.Map;
+            import java.util.Objects;
+            @NullMarked
+            class Test {
+                static void unsafeGetValues(Map<String, String> values, List<String> keys) {
+                    keys.stream().map(values::get).forEach(value -> {
+                        // BUG: Diagnostic contains: dereferenced expression 'value' is @Nullable
+                        value.hashCode();
+                    });
+                }
+                static List<String> getValues(Map<String, String> values, List<String> keys) {
+                    return keys.stream()
+                        .map(values::get)
+                        .filter(Objects::nonNull)
+                        .toList();
+                }
+            }""")
+        .doTest();
+  }
+
   private CompilationTestHelper makeHelper() {
     return makeTestHelperWithArgs(
         JSpecifyJavacConfig.withJSpecifyModeArgs(List.of("-XepOpt:NullAway:OnlyNullMarked=true")));

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

@@ -1027,7 +1027,6 @@ static void test() {
         .doTest();
   }
 
-  @Ignore("https://github.com/uber/NullAway/issues/1462")
   @Test
   public void streamMapNullableTest() {
     makeHelper()