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](jspecify/jdk@15eda89)
- 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. But I've taken the shortcut of not waiting for the JDK APIs.)
- 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 google/guava#6824.

PiperOrigin-RevId: 611445633

Author: cpovirk

j2kt/jre/java/common/java/util/Map.java

@@ -23,6 +23,7 @@
 import javaemul.internal.annotations.KtName;
 import javaemul.internal.annotations.KtNative;
 import javaemul.internal.annotations.KtProperty;
+import jsinterop.annotations.JsNonNull;
 import org.jspecify.nullness.NullMarked;
 import org.jspecify.nullness.Nullable;
 
@@ -71,7 +72,7 @@ Comparator<Map.Entry<K, V>> comparingByValue() {
   void clear();
 
   default @Nullable V compute(
-      K key, BiFunction<? super K, ? super @Nullable V, ? extends V> remappingFunction) {
+      K key, BiFunction<? super K, ? super @Nullable V, ? extends @Nullable V> remappingFunction) {
     return ktNative();
   }
 
@@ -81,7 +82,7 @@ default V computeIfAbsent(K key, Function<? super K, ? extends V> mappingFunctio
 
   @KtName("java_computeIfPresent")
   default @Nullable V computeIfPresent(
-      K key, BiFunction<? super K, ? super V, ? extends V> remappingFunction) {
+      K key, BiFunction<? super K, ? super @JsNonNull V, ? extends @Nullable V> remappingFunction) {
     return ktNative();
   }
 
@@ -115,7 +116,10 @@ default void forEach(BiConsumer<? super K, ? super V> action) {
 
   @KtName("java_merge")
   default @Nullable V merge(
-      K key, V value, BiFunction<? super V, ? super V, ? extends @Nullable V> remappingFunction) {
+      K key,
+      @JsNonNull V value,
+      BiFunction<? super @JsNonNull V, ? super @JsNonNull V, ? extends @Nullable V>
+          remappingFunction) {
     return ktNative();
   }
 

j2kt/jre/java/common/javaemul/lang/JavaMap.kt

@@ -43,7 +43,7 @@ interface JavaMap<K, V> : MutableMap<K, V> {
   fun computeIfAbsent(key: K, mappingFunction: Function<in K, out V>): V =
     default_computeIfAbsent(key, mappingFunction)
 
-  fun java_computeIfPresent(key: K, remappingFunction: BiFunction<in K, in V, out V?>): V? =
+  fun java_computeIfPresent(key: K, remappingFunction: BiFunction<in K, in V & Any, out V?>): V? =
     default_computeIfPresent(key, remappingFunction)
 
   fun forEach(action: BiConsumer<in K, in V>) = default_forEach(action)
@@ -66,7 +66,7 @@ interface JavaMap<K, V> : MutableMap<K, V> {
 
   // The Java `merge` function exists on Kotlin/JVM but is undocumented. So we rename our `merge` to
   // avoid a collision.
-  fun java_merge(key: K, value: V, remap: BiFunction<in V, in V, out V?>): V? =
+  fun java_merge(key: K, value: V & Any, remap: BiFunction<in V & Any, in V & Any, out V?>): V? =
     default_merge(key, value, remap)
 
   fun java_putAll(t: MutableMap<out K, out V>)
@@ -97,7 +97,7 @@ fun <K, V> MutableMap<K, V>.java_remove(key: Any?): V? = remove(key as K)
 
 fun <K, V> MutableMap<K, V>.java_computeIfPresent(
   key: K,
-  mappingFunction: BiFunction<in K, in V, out V?>
+  mappingFunction: BiFunction<in K, in V & Any, out V?>,
 ): V? =
   if (this is JavaMap) java_computeIfPresent(key, mappingFunction)
   else default_computeIfPresent(key, mappingFunction)
@@ -108,8 +108,8 @@ fun <K, V> MutableMap<K, V>.java_getOrDefault(key: Any?, defaultValue: V?): V? =
 
 fun <K, V> MutableMap<K, V>.java_merge(
   key: K,
-  value: V,
-  remap: BiFunction<in V, in V, out V?>
+  value: V & Any,
+  remap: BiFunction<in V & Any, in V & Any, out V?>,
 ): V? = if (this is JavaMap) java_merge(key, value, remap) else default_merge(key, value, remap)
 
 fun <K, V> MutableMap<K, V>.java_remove(key: Any?, value: Any?): Boolean =
@@ -121,7 +121,7 @@ internal inline fun <K, V> MutableMap<K, V>.default_forEach(action: BiConsumer<i
 
 internal fun <K, V> MutableMap<K, V>.default_compute(
   key: K,
-  remappingFunction: BiFunction<in K, in V?, out V?>
+  remappingFunction: BiFunction<in K, in V?, out V?>,
 ): V? {
   val oldValue = this[key]
   val newValue = remappingFunction.apply(key, oldValue)
@@ -142,7 +142,7 @@ internal fun <K, V> MutableMap<K, V>.default_compute(
 
 internal fun <K, V> MutableMap<K, V>.default_computeIfAbsent(
   key: K,
-  mappingFunction: Function<in K, out V>
+  mappingFunction: Function<in K, out V>,
 ): V {
   val oldValue = this[key]
   if (oldValue == null) {
@@ -155,7 +155,7 @@ internal fun <K, V> MutableMap<K, V>.default_computeIfAbsent(
 
 private fun <K, V> MutableMap<K, V>.default_computeIfPresent(
   key: K,
-  remappingFunction: BiFunction<in K, in V, out V?>
+  remappingFunction: BiFunction<in K, in V & Any, out V?>,
 ): V? {
   val oldValue = this[key]
   if (oldValue != null) {
@@ -176,8 +176,8 @@ internal fun <K, V> MutableMap<K, V>.default_getOrDefault(key: Any?, defaultValu
 
 private fun <K, V> MutableMap<K, V>.default_merge(
   key: K,
-  value: V,
-  remap: BiFunction<in V, in V, out V?>
+  value: V & Any,
+  remap: BiFunction<in V & Any, in V & Any, out V?>,
 ): V? {
   val oldValue = get(key)
   val newValue: V? = if (oldValue == null) value else remap.apply(oldValue, value)

j2kt/jre/java/native/java/util/stream/Collectors.java

@@ -351,7 +351,7 @@ public static Collector<CharSequence,?,String> joining(CharSequence delimiter) {
     return toCollection(ArrayList::new);
   }
 
-  public static <T extends @Nullable Object, K extends @Nullable Object, U extends @Nullable Object>
+  public static <T extends @Nullable Object, K extends @Nullable Object, U>
       Collector<T, ?, Map<K, U>> toMap(
           final Function<? super T, ? extends K> keyMapper,
           final Function<? super T, ? extends U> valueMapper) {
@@ -363,19 +363,15 @@ public static Collector<CharSequence,?,String> joining(CharSequence delimiter) {
         });
   }
 
-  public static <T extends @Nullable Object, K extends @Nullable Object, U extends @Nullable Object>
+  public static <T extends @Nullable Object, K extends @Nullable Object, U>
       Collector<T, ?, Map<K, U>> toMap(
           Function<? super T, ? extends K> keyMapper,
           Function<? super T, ? extends U> valueMapper,
           BinaryOperator<U> mergeFunction) {
     return toMap(keyMapper, valueMapper, mergeFunction, HashMap::new);
   }
 
-  public static <
-          T extends @Nullable Object,
-          K extends @Nullable Object,
-          U extends @Nullable Object,
-          M extends Map<K, U>>
+  public static <T extends @Nullable Object, K extends @Nullable Object, U, M extends Map<K, U>>
       Collector<T, ?, M> toMap(
           final Function<? super T, ? extends K> keyMapper,
           final Function<? super T, ? extends U> valueMapper,
@@ -385,12 +381,8 @@ public static Collector<CharSequence,?,String> joining(CharSequence delimiter) {
         mapSupplier,
         (map, item) -> {
           K key = keyMapper.apply(item);
-          U newValue = valueMapper.apply(item);
-          if (map.containsKey(key)) {
-            map.put(key, mergeFunction.apply(map.get(key), newValue));
-          } else {
-            map.put(key, newValue);
-          }
+          U value = valueMapper.apply(item);
+          map.merge(key, value, mergeFunction);
         },
         (m1, m2) -> mergeAll(m1, m2, mergeFunction),
         Collector.Characteristics.IDENTITY_FINISH);
@@ -418,8 +410,8 @@ D streamAndCollect(Collector<? super T, A, D> downstream, List<T> list) {
     return downstream.finisher().apply(a);
   }
 
-  private static <K extends @Nullable Object, V extends @Nullable Object, M extends Map<K, V>>
-      M mergeAll(M m1, M m2, BinaryOperator<V> mergeFunction) {
+  private static <K extends @Nullable Object, V, M extends Map<K, V>> M mergeAll(
+      M m1, M m2, BinaryOperator<V> mergeFunction) {
     for (Map.Entry<K, V> entry : m2.entrySet()) {
       m1.merge(entry.getKey(), entry.getValue(), mergeFunction);
     }

j2kt/jre/javatests/smoke/CollectionsTest.java

@@ -55,6 +55,7 @@
 import java.util.function.Predicate;
 import java.util.function.UnaryOperator;
 import java.util.stream.Stream;
+import jsinterop.annotations.JsNonNull;
 import org.jspecify.nullness.Nullable;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -334,7 +335,8 @@ public V computeIfAbsent(K key, Function<? super K, ? extends V> mappingFunction
     @Override
     @SuppressWarnings("nullness:override") // Checker expects @PolyNull (missing in JSpecify)
     public @Nullable V computeIfPresent(
-        K key, BiFunction<? super K, ? super V, ? extends @Nullable V> remappingFunction) {
+        K key,
+        BiFunction<? super K, ? super @JsNonNull V, ? extends @Nullable V> remappingFunction) {
       return super.computeIfPresent(key, remappingFunction);
     }
 
@@ -354,7 +356,9 @@ public void forEach(BiConsumer<? super K, ? super V> action) {
     // Checker framework uses @PolyNull V, JSpecify uses @Nullable V
     @SuppressWarnings({"nullness:override", "argument.type"})
     public @Nullable V merge(
-        K key, V value, BiFunction<? super V, ? super V, ? extends @Nullable V> remap) {
+        K key,
+        @JsNonNull V value,
+        BiFunction<? super @JsNonNull V, ? super @JsNonNull V, ? extends @Nullable V> remap) {
       return super.merge(key, value, remap);
     }