fix memory leak in builders: they should not hold on to the original collection (#193)

Author: zajac

core/commonMain/src/implementations/immutableList/PersistentVectorBuilder.kt

@@ -12,38 +12,52 @@ import kotlinx.collections.immutable.internal.MutabilityOwnership
 import kotlinx.collections.immutable.internal.assert
 import kotlinx.collections.immutable.internal.modCount
 
-internal class PersistentVectorBuilder<E>(private var vector: PersistentList<E>,
-                                          private var vectorRoot: Array<Any?>?,
-                                          private var vectorTail: Array<Any?>,
+internal class PersistentVectorBuilder<E>(vector: PersistentList<E>,
+                                          vectorRoot: Array<Any?>?,
+                                          vectorTail: Array<Any?>,
                                           internal var rootShift: Int) : AbstractMutableList<E>(), PersistentList.Builder<E> {
+
+    private var builtVector: PersistentList<E>? = vector
     private var ownership = MutabilityOwnership()
+
     internal var root = vectorRoot
-        private set
+        private set (value) {
+            if (value !== field) {
+                builtVector = null
+                field = value
+            }
+        }
+
     internal var tail = vectorTail
-        private set
+        private set (value) {
+            if (value !== field) {
+                builtVector = null
+                field = value
+            }
+        }
+
     override var size = vector.size
         private set
 
     internal fun getModCount() = modCount
 
     override fun build(): PersistentList<E> {
-        vector = if (root === vectorRoot && tail === vectorTail) {
-            vector
-        } else {
+        return builtVector ?: run {
+            val root = root
+            val tail = tail
             ownership = MutabilityOwnership()
-            vectorRoot = root
-            vectorTail = tail
-            if (root == null) {
+            val newlyBuiltVector: PersistentList<E> = if (root == null) {
                 if (tail.isEmpty()) {
                     persistentVectorOf()
                 } else {
                     SmallPersistentVector(tail.copyOf(size))
                 }
             } else {
-                PersistentVector(root!!, tail, size, rootShift)
+                PersistentVector(root, tail, size, rootShift)
             }
+            builtVector = newlyBuiltVector
+            newlyBuiltVector
         }
-        return vector
     }
 
     private fun rootSize(): Int {

core/commonMain/src/implementations/immutableMap/PersistentHashMapBuilder.kt

@@ -12,10 +12,18 @@ import kotlinx.collections.immutable.internal.DeltaCounter
 import kotlinx.collections.immutable.internal.MapImplementation
 import kotlinx.collections.immutable.internal.MutabilityOwnership
 
-internal class PersistentHashMapBuilder<K, V>(private var map: PersistentHashMap<K, V>) : PersistentMap.Builder<K, V>, AbstractMutableMap<K, V>() {
+internal class PersistentHashMapBuilder<K, V>(map: PersistentHashMap<K, V>) : PersistentMap.Builder<K, V>, AbstractMutableMap<K, V>() {
+    internal var builtMap: PersistentHashMap<K, V>? = map
+        private set 
     internal var ownership = MutabilityOwnership()
         private set
     internal var node = map.node
+        set(value) {
+            if (value !== field) {
+                field = value
+                builtMap = null
+            }
+        }
     internal var operationResult: V? = null
     internal var modCount = 0
 
@@ -27,13 +35,12 @@ internal class PersistentHashMapBuilder<K, V>(private var map: PersistentHashMap
         }
 
     override fun build(): PersistentHashMap<K, V> {
-        map = if (node === map.node) {
-            map
-        } else {
+        return builtMap ?: run {
+            val newlyBuiltMap = PersistentHashMap(node, size)
+            builtMap = newlyBuiltMap
             ownership = MutabilityOwnership()
-            PersistentHashMap(node, size)
+            newlyBuiltMap
         }
-        return map
     }
 
     override val entries: MutableSet<MutableMap.MutableEntry<K, V>>

core/commonMain/src/implementations/immutableSet/PersistentHashSetBuilder.kt

@@ -9,11 +9,17 @@ import kotlinx.collections.immutable.PersistentSet
 import kotlinx.collections.immutable.internal.DeltaCounter
 import kotlinx.collections.immutable.internal.MutabilityOwnership
 
-internal class PersistentHashSetBuilder<E>(private var set: PersistentHashSet<E>) : AbstractMutableSet<E>(), PersistentSet.Builder<E> {
+internal class PersistentHashSetBuilder<E>(set: PersistentHashSet<E>) : AbstractMutableSet<E>(), PersistentSet.Builder<E> {
+    private var builtSet: PersistentHashSet<E>? = set
     internal var ownership = MutabilityOwnership()
         private set
     internal var node = set.node
-        private set
+        private set (value) {
+            if (value !== field) {
+                builtSet = null
+                field = value
+            }
+        }
     internal var modCount = 0
         private set
 
@@ -25,13 +31,12 @@ internal class PersistentHashSetBuilder<E>(private var set: PersistentHashSet<E>
         }
 
     override fun build(): PersistentHashSet<E> {
-        set = if (node === set.node) {
-            set
-        } else {
+        return builtSet ?: run {
+            val newlyBuiltSet = PersistentHashSet(node, size)
             ownership = MutabilityOwnership()
-            PersistentHashSet(node, size)
+            builtSet = newlyBuiltSet
+            newlyBuiltSet
         }
-        return set
     }
 
     override fun contains(element: E): Boolean {

core/commonMain/src/implementations/persistentOrderedMap/PersistentOrderedMapBuilder.kt

@@ -12,7 +12,9 @@ import kotlinx.collections.immutable.internal.EndOfChain
 import kotlinx.collections.immutable.internal.MapImplementation
 import kotlinx.collections.immutable.internal.assert
 
-internal class PersistentOrderedMapBuilder<K, V>(private var map: PersistentOrderedMap<K, V>) : AbstractMutableMap<K, V>(), PersistentMap.Builder<K, V> {
+internal class PersistentOrderedMapBuilder<K, V>(map: PersistentOrderedMap<K, V>) : AbstractMutableMap<K, V>(), PersistentMap.Builder<K, V> {
+    private var builtMap: PersistentOrderedMap<K, V>? = map
+
     internal var firstKey = map.firstKey
         private set
 
@@ -23,15 +25,17 @@ internal class PersistentOrderedMapBuilder<K, V>(private var map: PersistentOrde
     override val size: Int get() = hashMapBuilder.size
 
     override fun build(): PersistentMap<K, V> {
-        val newHashMap = hashMapBuilder.build()
-        map = if (newHashMap === map.hashMap) {
+        return builtMap?.also { map ->
+            assert(hashMapBuilder.builtMap != null)
             assert(firstKey === map.firstKey)
             assert(lastKey === map.lastKey)
-            map
-        } else {
-            PersistentOrderedMap(firstKey, lastKey, newHashMap)
+        } ?: run {
+            assert(hashMapBuilder.builtMap == null)
+            val newHashMap = hashMapBuilder.build()
+            val newOrdered = PersistentOrderedMap(firstKey, lastKey, newHashMap)
+            builtMap = newOrdered
+            newOrdered
         }
-        return map
     }
 
     override val entries: MutableSet<MutableMap.MutableEntry<K, V>>
@@ -55,34 +59,36 @@ internal class PersistentOrderedMapBuilder<K, V>(private var map: PersistentOrde
 
     override fun put(key: K, value: @UnsafeVariance V): V? {
         val links = hashMapBuilder[key]
-        if (links != null) {
+        return if (links != null) {
             if (links.value === value) {
-                return value
+                value
+            } else {
+                builtMap = null
+                hashMapBuilder[key] = links.withValue(value)
+                links.value
             }
-            hashMapBuilder[key] = links.withValue(value)
-            return links.value
-        }
-
-        if (isEmpty()) {  //  isEmpty
-            firstKey = key
-            lastKey = key
-            hashMapBuilder[key] = LinkedValue(value)
-            return null
+        } else {
+            builtMap = null
+            if (isEmpty()) {
+                firstKey = key
+                lastKey = key
+                hashMapBuilder[key] = LinkedValue(value)
+            } else {
+                @Suppress("UNCHECKED_CAST")
+                val lastKey = lastKey as K
+                val lastLinks = hashMapBuilder[lastKey]!!
+                assert(!lastLinks.hasNext)
+                hashMapBuilder[lastKey] = lastLinks.withNext(key)
+                hashMapBuilder[key] = LinkedValue(value, previous = lastKey)
+                this.lastKey = key
+            }
+            null
         }
-        @Suppress("UNCHECKED_CAST")
-        val lastKey = lastKey as K
-        val lastLinks = hashMapBuilder[lastKey]!!
-        assert(!lastLinks.hasNext)
-
-        hashMapBuilder[lastKey] = lastLinks.withNext(key)
-        hashMapBuilder[key] = LinkedValue(value, previous = lastKey)
-        this.lastKey = key
-        return null
     }
 
     override fun remove(key: K): V? {
         val links = hashMapBuilder.remove(key) ?: return null
-
+        builtMap = null
         if (links.hasPrevious) {
             val previousLinks = hashMapBuilder[links.previous]!!
 //            assert(previousLinks.next == key)
@@ -115,6 +121,9 @@ internal class PersistentOrderedMapBuilder<K, V>(private var map: PersistentOrde
     }
 
     override fun clear() {
+        if (hashMapBuilder.isNotEmpty()) {
+            builtMap = null
+        }
         hashMapBuilder.clear()
         firstKey = EndOfChain
         lastKey = EndOfChain

core/commonMain/src/implementations/persistentOrderedSet/PersistentOrderedSetBuilder.kt

@@ -9,7 +9,8 @@ import kotlinx.collections.immutable.PersistentSet
 import kotlinx.collections.immutable.internal.EndOfChain
 import kotlinx.collections.immutable.internal.assert
 
-internal class PersistentOrderedSetBuilder<E>(private var set: PersistentOrderedSet<E>) : AbstractMutableSet<E>(), PersistentSet.Builder<E> {
+internal class PersistentOrderedSetBuilder<E>(set: PersistentOrderedSet<E>) : AbstractMutableSet<E>(), PersistentSet.Builder<E> {
+    private var builtSet: PersistentOrderedSet<E>? = set
     internal var firstElement = set.firstElement
     private var lastElement = set.lastElement
     internal val hashMapBuilder = set.hashMap.builder()
@@ -18,15 +19,17 @@ internal class PersistentOrderedSetBuilder<E>(private var set: PersistentOrdered
         get() = hashMapBuilder.size
 
     override fun build(): PersistentSet<E> {
-        val newMap = hashMapBuilder.build()
-        set = if (newMap === set.hashMap) {
+        return builtSet?.also { set ->
+            assert(hashMapBuilder.builtMap != null)
             assert(firstElement === set.firstElement)
             assert(lastElement === set.lastElement)
-            set
-        } else {
-            PersistentOrderedSet(firstElement, lastElement, newMap)
+        } ?: run {
+            assert(hashMapBuilder.builtMap == null)
+            val newMap = hashMapBuilder.build()
+            val newSet = PersistentOrderedSet(firstElement, lastElement, newMap)
+            builtSet = newSet
+            newSet
         }
-        return set
     }
 
     override fun contains(element: E): Boolean {
@@ -37,6 +40,7 @@ internal class PersistentOrderedSetBuilder<E>(private var set: PersistentOrdered
         if (hashMapBuilder.containsKey(element)) {
             return false
         }
+        builtSet = null
         if (isEmpty()) {
             firstElement = element
             lastElement = element
@@ -56,7 +60,7 @@ internal class PersistentOrderedSetBuilder<E>(private var set: PersistentOrdered
 
     override fun remove(element: E): Boolean {
         val links = hashMapBuilder.remove(element) ?: return false
-
+        builtSet = null
         if (links.hasPrevious) {
             val previousLinks = hashMapBuilder[links.previous]!!
 //            assert(previousLinks.next == element)
@@ -78,6 +82,9 @@ internal class PersistentOrderedSetBuilder<E>(private var set: PersistentOrdered
     }
 
     override fun clear() {
+        if (hashMapBuilder.isNotEmpty()) {
+            builtSet = null
+        }
         hashMapBuilder.clear()
         firstElement = EndOfChain
         lastElement = EndOfChain