fixup! Align ConcurrentMap.getOrPut behavior with MutableMap.getOrPut #KT-67339

# Conflicts:
#	libraries/stdlib/jdk8/src/kotlin/internal/jdk8/JDK8PlatformImplementations.kt
#	libraries/stdlib/jvm/src/kotlin/internal/PlatformImplementations.kt
diff --git a/libraries/stdlib/jdk8/src/kotlin/internal/jdk8/JDK8PlatformImplementations.kt b/libraries/stdlib/jdk8/src/kotlin/internal/jdk8/JDK8PlatformImplementations.kt
index e93a06e..a340742 100644
--- a/libraries/stdlib/jdk8/src/kotlin/internal/jdk8/JDK8PlatformImplementations.kt
+++ b/libraries/stdlib/jdk8/src/kotlin/internal/jdk8/JDK8PlatformImplementations.kt
@@ -17,6 +17,7 @@
 @file:Suppress("INVISIBLE_REFERENCE", "INVISIBLE_MEMBER", "CANNOT_OVERRIDE_INVISIBLE_MEMBER")
 package kotlin.internal.jdk8
 
+import java.util.concurrent.ConcurrentMap
 import java.util.regex.MatchResult
 import java.util.regex.Matcher
 import kotlin.internal.jdk7.JDK7PlatformImplementations
@@ -71,4 +72,15 @@
                 Instant.fromEpochMilliseconds(System.currentTimeMillis())
         }
     }
+
+    // computeIfAbsent doesn't put the newValue if it's null, hence require non-null newValue
+    override fun <K, V, NewV : V & Any> computeIfAbsent(map: ConcurrentMap<K, V>, key: K, newValue: NewV): V =
+        if (sdkIsNullOrAtLeast(24)) {
+            // computeIfAbsent is available since Android SDK 24.
+            map.computeIfAbsent(key) { newValue }
+        } else {
+            // Use putIfAbsent for older SDKs as a fallback.
+            // Note: If the key was mapped to a null value, putIfAbsent won't replace it, but the new value will still be returned.
+            map.putIfAbsent(key, newValue) ?: newValue
+        }
 }
diff --git a/libraries/stdlib/jvm/src/kotlin/collections/MapsJVM.kt b/libraries/stdlib/jvm/src/kotlin/collections/MapsJVM.kt
index 9f46d6a..ffc9401 100644
--- a/libraries/stdlib/jvm/src/kotlin/collections/MapsJVM.kt
+++ b/libraries/stdlib/jvm/src/kotlin/collections/MapsJVM.kt
@@ -14,6 +14,7 @@
 import java.util.TreeMap
 import java.util.concurrent.ConcurrentMap
 import kotlin.collections.builders.MapBuilder
+import kotlin.internal.IMPLEMENTATIONS
 
 /**
  * Returns a new read-only map, mapping only the specified key to the
@@ -63,37 +64,42 @@
  *
  * Returns the value for the given [key] if the value is present and not `null`.
  * Otherwise, calls the [defaultValue] function,
- * puts its result into the map under the given key and returns the call result.
+ * puts its result into the map under the given key, and returns the call result.
+ *
+ * The result of [defaultValue] is put into the map even if it is `null`.
+ * If [defaultValue] throws an exception, the exception is rethrown.
  *
  * This function guarantees not to put the new value into the map if the key is already
- * mapped to a non-null value. However, the [defaultValue] function may still be invoked.
+ * associated with a non-null value. However, the [defaultValue] function may still be invoked.
  *
- * This function relies on [ConcurrentMap.computeIfAbsent]. Hence, the `ConcurrentMap` implementations
- * that support `null` values must override the default `computeIfAbsent` implementation.
+ * This function relies on [ConcurrentMap.computeIfAbsent]. Therefore, `ConcurrentMap` implementations
+ * that support `null` values must override the default `computeIfAbsent` implementation, so that
+ * the result of the `mappingFunction` is put into the map both when there is no existing value for the key
+ * and when the key is associated with a `null` value.
+ *
+ * @throws NullPointerException if the specified [key] or the result of [defaultValue] is `null`,
+ *   and this concurrent map does not support `null` keys or values.
  *
  * @sample samples.collections.Maps.Usage.getOrPut
  */
 public inline fun <K, V> ConcurrentMap<K, V>.getOrPut(key: K, defaultValue: () -> V): V {
-    // Do not use computeIfAbsent on JVM8 as it would change locking behavior
-    this.get(key)?.let { return it }
-
-    val newValue = defaultValue()
-    return try {
-        // computeIfAbsent doesn't put the newValue if it's null
-        if (newValue == null) {
-            this.putIfAbsent(key, newValue) ?: newValue
-        } else {
-            this.computeIfAbsent(key) { newValue }
-        }
-    } catch (e: LinkageError) {
-        // In Android projects without the desugared library, computeIfAbsent is available since SDK 24.
-        // See https://developer.android.com/studio/write/java8-support-table for more info.
-        // Use putIfAbsent as a fallback in this case.
-        // Note: If the key was mapped to a null value, putIfAbsent won't replace it, but the new value will still be returned.
-        this.putIfAbsent(key, newValue) ?: newValue
-    }
+    // Do not call defaultValue() inside computeIfAbsent mappingFunction as it would change locking behavior
+    return this.get(key) ?: this.internalGetOrPutIfNull(key, newValue = defaultValue())
 }
 
+@PublishedApi
+internal fun <K, V> ConcurrentMap<K, V>.internalGetOrPutIfNull(key: K, newValue: V): V =
+    if (newValue != null) {
+        // Returns the current (existing or computed) value associated with the specified key
+        IMPLEMENTATIONS.computeIfAbsent(this, key, newValue)
+    } else {
+        // Returns the previous value associated with the specified key.
+        //   If the key is already mapped, returns the mapped value;
+        //   otherwise, puts the newValue and returns null, which is the value that was put.
+        @Suppress("UNCHECKED_CAST")
+        this.putIfAbsent(key, newValue) as V
+    }
+
 
 /**
  * Converts this [Map] to a [SortedMap]. The resulting [SortedMap] determines the equality and order of keys according to their natural sorting order.
diff --git a/libraries/stdlib/jvm/src/kotlin/internal/PlatformImplementations.kt b/libraries/stdlib/jvm/src/kotlin/internal/PlatformImplementations.kt
index 0615fa9..9cf1b33 100644
--- a/libraries/stdlib/jvm/src/kotlin/internal/PlatformImplementations.kt
+++ b/libraries/stdlib/jvm/src/kotlin/internal/PlatformImplementations.kt
@@ -6,6 +6,7 @@
 package kotlin.internal
 
 import java.lang.reflect.Method
+import java.util.concurrent.ConcurrentMap
 import java.util.regex.MatchResult
 import kotlin.random.FallbackThreadLocalRandom
 import kotlin.random.Random
@@ -49,6 +50,10 @@
     public open fun getSystemClock(): Clock {
         throw UnsupportedOperationException("getSystemClock should not be called on the base PlatformImplementations.")
     }
+
+    public open fun <K, V, NewV : V & Any> computeIfAbsent(map: ConcurrentMap<K, V>, key: K, newValue: NewV): V {
+        throw UnsupportedOperationException("computeIfAbsent should not be called on the base PlatformImplementations.")
+    }
 }
 
 
diff --git a/libraries/stdlib/jvm/test/collections/MapJVMTest.kt b/libraries/stdlib/jvm/test/collections/MapJVMTest.kt
index c827469..9836499 100644
--- a/libraries/stdlib/jvm/test/collections/MapJVMTest.kt
+++ b/libraries/stdlib/jvm/test/collections/MapJVMTest.kt
@@ -7,6 +7,7 @@
 
 import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.ConcurrentMap
+import java.util.function.Function
 import kotlin.test.*
 
 class MapJVMTest {
@@ -63,16 +64,123 @@
         assertEquals(listOf(1, 3, 5), map.keys.toList())
         assertEquals(listOf('b', 'd', 'f'), map.values.toList())
     }
-    
-    @Test fun getOrPutFailsOnConcurrentMap() {
-        val map = ConcurrentHashMap<String, Int>()
 
-        // not an error anymore
-        expect(1) {
-            map.getOrPut("x") { 1 }
+    @Test
+    fun getOrPutOnConcurrentHashMap() {
+        val map = ConcurrentHashMap<String?, String?>()
+
+        assertEquals("v1", map.getOrPut("k1") { "v1" })
+        assertEquals("v1", map.getOrPut("k1") { "newV1" })
+        // Doesn't throw because defaultValue() wasn't called or its result wasn't tried to be put
+        assertEquals("v1", map.getOrPut("k1") { null })
+
+        // Doesn't support null values
+        assertFailsWith<NullPointerException> {
+            map.getOrPut("k2") { null }
         }
-        expect(1) {
-            (map as MutableMap<String, Int>).getOrPut("x") { 1 }
+        assertFalse(map.containsKey("k2"))
+        assertEquals("v2", map.getOrPut("k2") { "v2" })
+        // Doesn't throw because defaultValue() wasn't called or its result wasn't tried to be put
+        assertEquals("v2", map.getOrPut("k2") { null })
+
+        // Doesn't support null keys
+        assertFailsWith<NullPointerException> {
+            map.getOrPut(null) { "v3" }
         }
+        // Doesn't support null keys and values
+        assertFailsWith<NullPointerException> {
+            map.getOrPut(null) { null }
+        }
+
+        val expected = setOf(
+            "k1" to "v1",
+            "k2" to "v2"
+        )
+        assertEquals(expected, map.entries.map { it.toPair() }.toSet())
+    }
+
+    @Test
+    fun getOrPutOnConcurrentMap() {
+        val map = SimpleConcurrentMap<String?, String?>()
+
+        assertEquals("v1", map.getOrPut("k1") { "v1" })
+        assertEquals("v1", map.getOrPut("k1") { "newV1" })
+        assertEquals("v1", map.getOrPut("k1") { null })
+
+        assertEquals(null, map.getOrPut("k2") { null })
+        assertTrue(map.containsKey("k2"))
+        assertEquals("v2", map.getOrPut("k2") { "v2" }) // replace null value
+        assertEquals("v2", map.getOrPut("k2") { null })
+
+        assertEquals("v3", map.getOrPut(null) { "v3" })
+        assertEquals("v3", map.getOrPut(null) { "newV3" })
+        assertEquals("v3", map.getOrPut(null) { null })
+
+        val expected = listOf(
+            "k1" to "v1",
+            "k2" to "v2",
+            null to "v3"
+        )
+        assertContentEquals(expected, map.entries.map { it.toPair() })
     }
 }
+
+
+// Allows null values. Not actually multi-thread safe.
+private class SimpleConcurrentMap<K, V> : AbstractMutableMap<K, V>(), ConcurrentMap<K, V> {
+    private val backing = mutableMapOf<K, V>()
+
+    override val size: Int get() = backing.size
+
+    override fun get(key: K): V? =
+        backing[key]
+
+    override fun getOrDefault(key: K, defaultValue: V): V =
+        if (backing.containsKey(key)) {
+            @Suppress("UNCHECKED_CAST")
+            backing[key] as V
+        } else {
+            defaultValue
+        }
+
+    override fun put(key: K, value: V): V? =
+        backing.put(key, value)
+
+    override fun putIfAbsent(key: K, value: V): V? =
+        if (!backing.containsKey(key)) {
+            backing.put(key, value)
+        } else {
+            backing.get(key)
+        }
+
+    // Implementations which support null values must override this default implementation.
+    override fun computeIfAbsent(key: K, mappingFunction: Function<in K, out V>): V =
+        backing.get(key) ?: run {
+            val newValue = mappingFunction.apply(key)
+            if (newValue != null) {
+                backing.put(key, newValue)
+            }
+            newValue
+        }
+
+    override fun remove(key: K, value: V): Boolean =
+        backing.remove(key, value)
+
+    override fun replace(key: K, oldValue: V, newValue: V): Boolean =
+        if (backing.containsKey(key) && backing[key] == oldValue) {
+            backing.put(key, newValue)
+            true
+        } else {
+            false
+        }
+
+    override fun replace(key: K, value: V): V? =
+        if (backing.containsKey(key)) {
+            backing.put(key, value)
+        } else {
+            null
+        }
+
+    override val entries: MutableSet<MutableMap.MutableEntry<K, V>>
+        get() = backing.entries
+}
\ No newline at end of file
diff --git a/libraries/tools/binary-compatibility-validator/reference-public-api/kotlin-stdlib-runtime-merged.txt b/libraries/tools/binary-compatibility-validator/reference-public-api/kotlin-stdlib-runtime-merged.txt
index f84b68a..e3fc312 100644
--- a/libraries/tools/binary-compatibility-validator/reference-public-api/kotlin-stdlib-runtime-merged.txt
+++ b/libraries/tools/binary-compatibility-validator/reference-public-api/kotlin-stdlib-runtime-merged.txt
@@ -2566,6 +2566,7 @@
 	public static final fun getOrPut (Ljava/util/concurrent/ConcurrentMap;Ljava/lang/Object;Lkotlin/jvm/functions/Function0;)Ljava/lang/Object;
 	public static final fun getValue (Ljava/util/Map;Ljava/lang/Object;)Ljava/lang/Object;
 	public static final fun hashMapOf ([Lkotlin/Pair;)Ljava/util/HashMap;
+	public static final fun internalGetOrPutIfNull (Ljava/util/concurrent/ConcurrentMap;Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;
 	public static final fun linkedMapOf ([Lkotlin/Pair;)Ljava/util/LinkedHashMap;
 	public static final fun map (Ljava/util/Map;Lkotlin/jvm/functions/Function1;)Ljava/util/List;
 	public static final fun mapCapacity (I)I