[K/N] HashMap/HashSet doesn't reclaim storage after removing elements #KT-53310

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

@@ -10,11 +10,17 @@ import java.io.InvalidObjectException
 import java.io.NotSerializableException
 
 internal class MapBuilder<K, V> private constructor(
+    // keys in insert order
     private var keysArray: Array<K>,
-    private var valuesArray: Array<V>?, // allocated only when actually used, always null in pure HashSet
+    // values in insert order, allocated only when actually used, always null in pure HashSet
+    private var valuesArray: Array<V>?,
+    // hash of a key by its index, -1 if a key at that index was removed
     private var presenceArray: IntArray,
+    // (index + 1) of a key by its hash, 0 if there is no key with that hash, -1 if collision chain continues to the hash-1
     private var hashArray: IntArray,
+    // max length of a collision chain
     private var maxProbeDistance: Int,
+    // index of the next key to be inserted
     private var length: Int
 ) : MutableMap<K, V>, Serializable {
     private var hashShift: Int = computeShift(hashSize)
@@ -165,7 +171,8 @@ internal class MapBuilder<K, V> private constructor(
 
     // ---------------------------- private ----------------------------
 
-    private val capacity: Int get() = keysArray.size
+    // Declared internal for testing
+    internal val capacity: Int get() = keysArray.size
     private val hashSize: Int get() = hashArray.size
 
     internal fun checkIsMutable() {
@@ -173,7 +180,19 @@ internal class MapBuilder<K, V> private constructor(
     }
 
     private fun ensureExtraCapacity(n: Int) {
-        ensureCapacity(length + n)
+        if (shouldCompact(extraCapacity = n)) {
+            rehash(hashSize)
+        } else {
+            ensureCapacity(length + n)
+        }
+    }
+
+    private fun shouldCompact(extraCapacity: Int): Boolean {
+        val spareCapacity = this.capacity - length
+        val gaps = length - size
+        return spareCapacity < extraCapacity                // there is no room for extraCapacity entries
+                && gaps + spareCapacity >= extraCapacity    // removing gaps prevents capacity expansion
+                && gaps >= this.capacity / 4                // at least 25% of current capacity is occupied by gaps
     }
 
     private fun ensureCapacity(capacity: Int) {
@@ -186,8 +205,6 @@ internal class MapBuilder<K, V> private constructor(
             presenceArray = presenceArray.copyOf(newSize)
             val newHashSize = computeHashSize(newSize)
             if (newHashSize > hashSize) rehash(newHashSize)
-        } else if (length + capacity - size > this.capacity) {
-            rehash(hashSize)
         }
     }
 

libraries/stdlib/jvm/test/collections/MapBuilderTest.kt

@@ -3,8 +3,11 @@
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
+@file:Suppress("INVISIBLE_REFERENCE", "INVISIBLE_MEMBER")
+
 package test.collections
 
+import kotlin.collections.builders.MapBuilder
 import kotlin.test.assertEquals
 import kotlin.test.assertFails
 import kotlin.test.Test
@@ -37,4 +40,16 @@ class MapBuilderTest {
             assertEquals(builderSize, size)
         }
     }
+
+    // KT-53310
+    @Test
+    fun reclaimStorage() {
+        val builder = MapBuilder<Int, Int>()
+        val initialCapacity = builder.capacity
+        repeat(20) {
+            builder[it] = it
+            builder.remove(it)
+        }
+        assertEquals(initialCapacity, builder.capacity)
+    }
 }

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

@@ -178,7 +178,19 @@ actual class HashMap<K, V> private constructor(
     }
 
     private fun ensureExtraCapacity(n: Int) {
-        ensureCapacity(length + n)
+        if (shouldCompact(extraCapacity = n)) {
+            rehash(hashSize)
+        } else {
+            ensureCapacity(length + n)
+        }
+    }
+
+    private fun shouldCompact(extraCapacity: Int): Boolean {
+        val spareCapacity = this.capacity - length
+        val gaps = length - size
+        return spareCapacity < extraCapacity                // there is no room for extraCapacity entries
+                && gaps + spareCapacity >= extraCapacity    // removing gaps prevents capacity expansion
+                && gaps >= this.capacity / 4                // at least 25% of current capacity is occupied by gaps
     }
 
     private fun ensureCapacity(capacity: Int) {
@@ -191,8 +203,6 @@ actual class HashMap<K, V> private constructor(
             presenceArray = presenceArray.copyOf(newSize)
             val newHashSize = computeHashSize(newSize)
             if (newHashSize > hashSize) rehash(newHashSize)
-        } else if (length + capacity - _size > this.capacity) {
-            rehash(hashSize)
         }
     }