Don't allocate temporary buffer in SmallPersistentVector.removeAll (#164)

---------

Co-authored-by: mcpiroman <[email protected]>

Author: Abduqodiri Qurbonzoda

benchmarks/commonMain/src/benchmarks/immutableList/RemoveAllPredicate.kt

@@ -0,0 +1,94 @@
+/*
+ * Copyright 2016-2023 JetBrains s.r.o.
+ * Use of this source code is governed by the Apache 2.0 License that can be found in the LICENSE.txt file.
+ */
+
+package benchmarks.immutableList
+
+import benchmarks.*
+import kotlinx.collections.immutable.PersistentList
+import kotlinx.benchmark.*
+import kotlinx.collections.immutable.persistentListOf
+import kotlin.random.Random
+
+@State(Scope.Benchmark)
+open class RemoveAllPredicate {
+    @Param(BM_1, BM_10, BM_100, BM_1000, BM_10000, BM_100000, BM_1000000, BM_10000000)
+    var size: Int = 0
+
+    private var persistentList = persistentListOf<String>()
+    private val truePredicate: (String) -> Boolean = { true }
+    private val falsePredicate: (String) -> Boolean = { false }
+    private var randomHalfElementsPredicate: (String) -> Boolean = truePredicate
+    private var randomTenElementsPredicate: (String) -> Boolean = truePredicate
+    private var randomOneElementPredicate: (String) -> Boolean = truePredicate
+    private var tailElementsPredicate: (String) -> Boolean = truePredicate
+
+    @Setup
+    fun prepare() {
+        val randomHalfElements = randomIndexes(size / 2).map { it.toString() }.toHashSet()
+        randomHalfElementsPredicate = { it in randomHalfElements }
+
+        val randomTenElements = randomIndexes(10).map { it.toString() }.toHashSet()
+        randomTenElementsPredicate = { it in randomTenElements }
+
+        val randomOneElement = Random.nextInt(size).toString()
+        randomOneElementPredicate = { it == randomOneElement }
+
+        val tailElements = List(tailSize()) { (size - 1 - it).toString() }.toHashSet()
+        tailElementsPredicate = { it in tailElements }
+
+        val allElements = List(size) { it.toString() }
+        persistentList = persistentListOf<String>().addAll(allElements)
+    }
+
+    // The benchmarks measure (time and memory spent in `removeAll` operation) / size
+    //
+    // Expected time: nearly constant
+    // Expected memory: nearly constant
+
+    /** Removes all elements. */
+    @Benchmark
+    fun removeAll_All(): PersistentList<String> {
+        return persistentList.removeAll(truePredicate)
+    }
+
+    /** Removes no elements. */
+    @Benchmark
+    fun removeAll_Non(): PersistentList<String> {
+        return persistentList.removeAll(falsePredicate)
+    }
+
+    /** Removes half of the elements randomly selected. */
+    @Benchmark
+    fun removeAll_RandomHalf(): PersistentList<String> {
+        return persistentList.removeAll(randomHalfElementsPredicate)
+    }
+
+    /** Removes 10 random elements. */
+    @Benchmark
+    fun removeAll_RandomTen(): PersistentList<String> {
+        return persistentList.removeAll(randomTenElementsPredicate)
+    }
+
+    /** Removes a random element. */
+    @Benchmark
+    fun removeAll_RandomOne(): PersistentList<String> {
+        return persistentList.removeAll(randomOneElementPredicate)
+    }
+
+    /** Removes last [tailSize] elements. */
+    @Benchmark
+    fun removeAll_Tail(): PersistentList<String> {
+        return persistentList.removeAll(tailElementsPredicate)
+    }
+
+    private fun randomIndexes(count: Int): List<Int> {
+        return List(count) { Random.nextInt(size) }
+    }
+
+    private fun tailSize(): Int {
+        val bufferSize = 32
+        return (size and (bufferSize - 1)).let { if (it == 0) bufferSize else it }
+    }
+}

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

@@ -49,30 +49,32 @@ internal class SmallPersistentVector<E>(private val buffer: Array<Any?>) : Immut
     }
 
     override fun removeAll(predicate: (E) -> Boolean): PersistentList<E> {
-        var newBuffer = buffer
         var newSize = size
-
-        var anyRemoved = false
+        var removeMask = 0
 
         for (index in 0 until size) {
             @Suppress("UNCHECKED_CAST")
             val element = buffer[index] as E
 
             if (predicate(element)) {
-                if (!anyRemoved) {
-                    newBuffer = buffer.copyOf()
-                    newSize = index
-
-                    anyRemoved = true
-                }
-            } else if (anyRemoved) {
-                newBuffer[newSize++] = element
+                newSize--
+                removeMask = removeMask or (1 shl index)
             }
         }
+
         return when (newSize) {
             size -> this
             0 -> EMPTY
-            else -> SmallPersistentVector(newBuffer.copyOfRange(0, newSize))
+            else -> {
+                val newBuffer = buffer.copyOf(newSize)
+                var newIndex = removeMask.countTrailingZeroBits()
+                for (index in newIndex + 1 until size) {
+                    if ((removeMask ushr index) and 1 == 0) {
+                        newBuffer[newIndex++] = buffer[index]
+                    }
+                }
+                SmallPersistentVector(newBuffer)
+            }
         }
     }