Fix memory leak in MapBuilder #KT-65518

libraries/stdlib/js/src/kotlin/collections/InternalHashMap.kt

@@ -134,11 +134,11 @@ internal class InternalHashMap<K, V> private constructor(
     }
 
     override fun remove(key: K): V? {
-        val index = removeKey(key)  // mutability gets checked here
+        checkIsMutable()
+        val index = findKey(key)
         if (index < 0) return null
-        val valuesArray = valuesArray!!
-        val oldValue = valuesArray[index]
-        valuesArray.resetAt(index)
+        val oldValue = valuesArray!![index]
+        removeEntryAt(index)
         return oldValue
     }
 
@@ -348,16 +348,9 @@ internal class InternalHashMap<K, V> private constructor(
         }
     }
 
-    private fun removeKey(key: K): Int {
-        checkIsMutable()
-        val index = findKey(key)
-        if (index < 0) return TOMBSTONE
-        removeKeyAt(index)
-        return index
-    }
-
-    private fun removeKeyAt(index: Int) {
+    private fun removeEntryAt(index: Int) {
         keysArray.resetAt(index)
+        valuesArray?.resetAt(index)
         removeHashAt(presenceArray[index])
         presenceArray[index] = TOMBSTONE
         _size--
@@ -459,7 +452,7 @@ internal class InternalHashMap<K, V> private constructor(
         val index = findKey(entry.key)
         if (index < 0) return false
         if (valuesArray!![index] != entry.value) return false
-        removeKeyAt(index)
+        removeEntryAt(index)
         return true
     }
 
@@ -467,7 +460,7 @@ internal class InternalHashMap<K, V> private constructor(
         checkIsMutable()
         val index = findValue(value)
         if (index < 0) return false
-        removeKeyAt(index)
+        removeEntryAt(index)
         return true
     }
 
@@ -508,7 +501,7 @@ internal class InternalHashMap<K, V> private constructor(
             checkForComodification()
             check(lastIndex != -1) { "Call next() before removing element from the iterator." }
             map.checkIsMutable()
-            map.removeKeyAt(lastIndex)
+            map.removeEntryAt(lastIndex)
             lastIndex = -1
             expectedModCount = map.modCount
         }

libraries/stdlib/jvm/src/kotlin/collections/builders/MapBuilder.kt

@@ -102,11 +102,11 @@ internal class MapBuilder<K, V> private constructor(
     }
 
     override fun remove(key: K): V? {
-        val index = removeKey(key)  // mutability gets checked here
+        checkIsMutable()
+        val index = findKey(key)
         if (index < 0) return null
-        val valuesArray = valuesArray!!
-        val oldValue = valuesArray[index]
-        valuesArray.resetAt(index)
+        val oldValue = valuesArray!![index]
+        removeEntryAt(index)
         return oldValue
     }
 
@@ -342,16 +342,17 @@ internal class MapBuilder<K, V> private constructor(
         }
     }
 
-    internal fun removeKey(key: K): Int {
+    internal fun removeKey(key: K): Boolean {
         checkIsMutable()
         val index = findKey(key)
-        if (index < 0) return TOMBSTONE
-        removeKeyAt(index)
-        return index
+        if (index < 0) return false
+        removeEntryAt(index)
+        return true
     }
 
-    private fun removeKeyAt(index: Int) {
+    private fun removeEntryAt(index: Int) {
         keysArray.resetAt(index)
+        valuesArray?.resetAt(index)
         removeHashAt(presenceArray[index])
         presenceArray[index] = TOMBSTONE
         size--
@@ -463,7 +464,7 @@ internal class MapBuilder<K, V> private constructor(
         val index = findKey(entry.key)
         if (index < 0) return false
         if (valuesArray!![index] != entry.value) return false
-        removeKeyAt(index)
+        removeEntryAt(index)
         return true
     }
 
@@ -471,7 +472,7 @@ internal class MapBuilder<K, V> private constructor(
         checkIsMutable()
         val index = findValue(element)
         if (index < 0) return false
-        removeKeyAt(index)
+        removeEntryAt(index)
         return true
     }
 
@@ -514,7 +515,7 @@ internal class MapBuilder<K, V> private constructor(
             checkForComodification()
             check(lastIndex != -1) { "Call next() before removing element from the iterator." }
             map.checkIsMutable()
-            map.removeKeyAt(lastIndex)
+            map.removeEntryAt(lastIndex)
             lastIndex = -1
             expectedModCount = map.modCount
         }
@@ -618,7 +619,7 @@ internal class MapBuilderKeys<E> internal constructor(
     override fun clear() = backing.clear()
     override fun add(element: E): Boolean = throw UnsupportedOperationException()
     override fun addAll(elements: Collection<E>): Boolean = throw UnsupportedOperationException()
-    override fun remove(element: E): Boolean = backing.removeKey(element) >= 0
+    override fun remove(element: E): Boolean = backing.removeKey(element)
     override fun iterator(): MutableIterator<E> = backing.keysIterator()
 
     override fun removeAll(elements: Collection<E>): Boolean {

libraries/stdlib/jvm/src/kotlin/collections/builders/SetBuilder.kt

@@ -34,7 +34,7 @@ internal class SetBuilder<E> internal constructor(
     override fun contains(element: E): Boolean = backing.containsKey(element)
     override fun clear() = backing.clear()
     override fun add(element: E): Boolean = backing.addKey(element) >= 0
-    override fun remove(element: E): Boolean = backing.removeKey(element) >= 0
+    override fun remove(element: E): Boolean = backing.removeKey(element)
     override fun iterator(): MutableIterator<E> = backing.keysIterator()
 
     override fun addAll(elements: Collection<E>): Boolean {

libraries/stdlib/native-wasm/src/kotlin/collections/HashMap.kt

@@ -136,11 +136,11 @@ public actual class HashMap<K, V> private constructor(
     }
 
     override actual fun remove(key: K): V? {
-        val index = removeKey(key)  // mutability gets checked here
+        checkIsMutable()
+        val index = findKey(key)
         if (index < 0) return null
-        val valuesArray = valuesArray!!
-        val oldValue = valuesArray[index]
-        valuesArray.resetAt(index)
+        val oldValue = valuesArray!![index]
+        removeEntryAt(index)
         return oldValue
     }
 
@@ -379,16 +379,17 @@ public actual class HashMap<K, V> private constructor(
         }
     }
 
-    internal fun removeKey(key: K): Int {
+    internal fun removeKey(key: K): Boolean {
         checkIsMutable()
         val index = findKey(key)
-        if (index < 0) return TOMBSTONE
-        removeKeyAt(index)
-        return index
+        if (index < 0) return false
+        removeEntryAt(index)
+        return true
     }
 
-    private fun removeKeyAt(index: Int) {
+    private fun removeEntryAt(index: Int) {
         keysArray.resetAt(index)
+        valuesArray?.resetAt(index)
         removeHashAt(presenceArray[index])
         presenceArray[index] = TOMBSTONE
         _size--
@@ -518,7 +519,7 @@ public actual class HashMap<K, V> private constructor(
         val index = findKey(entry.key)
         if (index < 0) return false
         if (valuesArray!![index] != entry.value) return false
-        removeKeyAt(index)
+        removeEntryAt(index)
         return true
     }
 
@@ -526,7 +527,7 @@ public actual class HashMap<K, V> private constructor(
         checkIsMutable()
         val index = findValue(element)
         if (index < 0) return false
-        removeKeyAt(index)
+        removeEntryAt(index)
         return true
     }
 
@@ -577,7 +578,7 @@ public actual class HashMap<K, V> private constructor(
             checkForComodification()
             check(lastIndex != -1) { "Call next() before removing element from the iterator." }
             map.checkIsMutable()
-            map.removeKeyAt(lastIndex)
+            map.removeEntryAt(lastIndex)
             lastIndex = -1
             expectedModCount = map.modCount
         }
@@ -682,7 +683,7 @@ internal class HashMapKeys<E> internal constructor(
     override fun clear() = backing.clear()
     override fun add(element: E): Boolean = throw UnsupportedOperationException()
     override fun addAll(elements: Collection<E>): Boolean = throw UnsupportedOperationException()
-    override fun remove(element: E): Boolean = backing.removeKey(element) >= 0
+    override fun remove(element: E): Boolean = backing.removeKey(element)
     override fun iterator(): MutableIterator<E> = backing.keysIterator()
 
     override fun removeAll(elements: Collection<E>): Boolean {

libraries/stdlib/native-wasm/src/kotlin/collections/HashSet.kt

@@ -70,7 +70,7 @@ public actual class HashSet<E> internal constructor(
     override fun getElement(element: E): E? = backing.getKey(element)
     override actual fun clear(): Unit = backing.clear()
     override actual fun add(element: E): Boolean = backing.addKey(element) >= 0
-    override actual fun remove(element: E): Boolean = backing.removeKey(element) >= 0
+    override actual fun remove(element: E): Boolean = backing.removeKey(element)
     override actual fun iterator(): MutableIterator<E> = backing.keysIterator()
 
     override actual fun addAll(elements: Collection<E>): Boolean {