From e9f45b0241e404fbc3122bfa74e40e2895123ee8 Mon Sep 17 00:00:00 2001
From: cpovirk <cpovirk@google.com>
Date: Wed, 8 Nov 2023 14:29:13 -0800
Subject: [PATCH] Improve nullness of types in some APIs related to `Map`
 merging, and fix `Collectors.toMap` null-handling.

- Restrict `Collections.toMap` value-type arguments to non-nullable types.
  - ...in J2KT, following what [we'd found in JSpecify research](https://github.com/jspecify/jdk/commit/15eda89bf8c3112d35ab7ffcc6ff6ce71502c5b5)
- Fix `Collections.toMap` to remove the key in question when `mergeFunction` returns `null`.
  - ...in J2KT
  - ...in J2CL
- Use `@NonNull` / `& Any` in a few places in `Map.merge` and `Map.computeIfPresent`.
  - ...in J2KT
  - ...in Guava `Map` implementations, even though we don't yet include `@NonNull` annotations in the JDK APIs that we build Guava against. (See post-submit discussion on cl/559605577.)
- Use `@Nullable` (to match the existing Kotlin `?` types) in the return types of `Map.computeIfPresent` and `Map.compute`.
  - ...in J2KT
- Test a bunch of this.

Note that the test for `mergeFunction` has to work around an overly restricted `toMap` signature that J2KT inherited from JSpecify. As discussed in a code comment there, this is fundamentally the same issue as we have in Guava with `ImmutableMap.toImmutableMap`, which is discussed as part of https://github.com/google/guava/issues/6824.

RELNOTES=n/a
PiperOrigin-RevId: 580659517
---
 guava/src/com/google/common/collect/Maps.java | 32 +++++++++++--------
 .../google/common/collect/Synchronized.java   | 14 ++++----
 2 files changed, 24 insertions(+), 22 deletions(-)

diff --git a/guava/src/com/google/common/collect/Maps.java b/guava/src/com/google/common/collect/Maps.java
index 34f09318a05c5..8fa916030306f 100644
--- a/guava/src/com/google/common/collect/Maps.java
+++ b/guava/src/com/google/common/collect/Maps.java
@@ -1735,15 +1735,15 @@ public V computeIfAbsent(
       throw new UnsupportedOperationException();
     }
 
-    /*
-     * TODO(cpovirk): Uncomment the @NonNull annotations below once our JDK stubs and J2KT
-     * emulations include them.
-     */
     @Override
     @CheckForNull
+    /*
+     * Our checker arguably should produce a nullness error here until we see @NonNull in JDK APIs.
+     * But it doesn't, which may be a sign that we still permit parameter contravariance in some
+     * cases?
+     */
     public V computeIfPresent(
-        K key,
-        BiFunction<? super K, ? super /*@NonNull*/ V, ? extends @Nullable V> remappingFunction) {
+        K key, BiFunction<? super K, ? super @NonNull V, ? extends @Nullable V> remappingFunction) {
       throw new UnsupportedOperationException();
     }
 
@@ -1757,11 +1757,11 @@ public V compute(
 
     @Override
     @CheckForNull
+    @SuppressWarnings("nullness") // TODO(b/262880368): Remove once we see @NonNull in JDK APIs
     public V merge(
         K key,
-        /*@NonNull*/ V value,
-        BiFunction<? super /*@NonNull*/ V, ? super /*@NonNull*/ V, ? extends @Nullable V>
-            function) {
+        @NonNull V value,
+        BiFunction<? super @NonNull V, ? super @NonNull V, ? extends @Nullable V> function) {
       throw new UnsupportedOperationException();
     }
 
@@ -3659,9 +3659,13 @@ public V computeIfAbsent(
      */
     @Override
     @CheckForNull
+    /*
+     * Our checker arguably should produce a nullness error here until we see @NonNull in JDK APIs.
+     * But it doesn't, which may be a sign that we still permit parameter contravariance in some
+     * cases?
+     */
     public V computeIfPresent(
-        K key,
-        BiFunction<? super K, ? super /*@NonNull*/ V, ? extends @Nullable V> remappingFunction) {
+        K key, BiFunction<? super K, ? super @NonNull V, ? extends @Nullable V> remappingFunction) {
       throw new UnsupportedOperationException();
     }
 
@@ -3675,11 +3679,11 @@ public V compute(
 
     @Override
     @CheckForNull
+    @SuppressWarnings("nullness") // TODO(b/262880368): Remove once we see @NonNull in JDK APIs
     public V merge(
         K key,
-        /*@NonNull*/ V value,
-        BiFunction<? super /*@NonNull*/ V, ? super /*@NonNull*/ V, ? extends @Nullable V>
-            function) {
+        @NonNull V value,
+        BiFunction<? super @NonNull V, ? super @NonNull V, ? extends @Nullable V> function) {
       throw new UnsupportedOperationException();
     }
 
diff --git a/guava/src/com/google/common/collect/Synchronized.java b/guava/src/com/google/common/collect/Synchronized.java
index 89916cd4d80c7..a8ad336b07432 100644
--- a/guava/src/com/google/common/collect/Synchronized.java
+++ b/guava/src/com/google/common/collect/Synchronized.java
@@ -49,6 +49,7 @@
 import java.util.function.UnaryOperator;
 import java.util.stream.Stream;
 import javax.annotation.CheckForNull;
+import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
@@ -1187,15 +1188,11 @@ public V computeIfAbsent(K key, Function<? super K, ? extends V> mappingFunction
       }
     }
 
-    /*
-     * TODO(cpovirk): Uncomment the @NonNull annotations below once our JDK stubs and J2KT
-     * emulations include them.
-     */
     @Override
     @CheckForNull
+    @SuppressWarnings("nullness") // TODO(b/262880368): Remove once we see @NonNull in JDK APIs
     public V computeIfPresent(
-        K key,
-        BiFunction<? super K, ? super /*@NonNull*/ V, ? extends @Nullable V> remappingFunction) {
+        K key, BiFunction<? super K, ? super @NonNull V, ? extends @Nullable V> remappingFunction) {
       synchronized (mutex) {
         return delegate().computeIfPresent(key, remappingFunction);
       }
@@ -1213,10 +1210,11 @@ public V compute(
 
     @Override
     @CheckForNull
+    @SuppressWarnings("nullness") // TODO(b/262880368): Remove once we see @NonNull in JDK APIs
     public V merge(
         K key,
-        /*@NonNull*/ V value,
-        BiFunction<? super /*@NonNull*/ V, ? super /*@NonNull*/ V, ? extends @Nullable V>
+        @NonNull V value,
+        BiFunction<? super @NonNull V, ? super @NonNull V, ? extends @Nullable V>
             remappingFunction) {
       synchronized (mutex) {
         return delegate().merge(key, value, remappingFunction);