Use CompositeDecoder::decodeCollectionSize to pre-allocate collections

benchmark/src/jmh/kotlin/kotlinx/benchmarks/cbor/CborCollectionsPreSizing.kt

@@ -0,0 +1,46 @@
+/*
+ * Copyright 2017-2026 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.benchmarks.cbor
+
+import kotlinx.serialization.ExperimentalSerializationApi
+import kotlinx.serialization.cbor.Cbor
+import kotlinx.serialization.decodeFromByteArray
+import kotlinx.serialization.encodeToByteArray
+import org.openjdk.jmh.annotations.*
+import java.util.concurrent.TimeUnit
+
+@Warmup(iterations = 7, time = 1)
+@Measurement(iterations = 7, time = 1)
+@BenchmarkMode(Mode.Throughput)
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+@State(Scope.Benchmark)
+@Fork(2)
+@OptIn(ExperimentalSerializationApi::class)
+open class CborCollectionsPreSizing {
+    @Param("1", "7", "17", "31")
+    var collectionSize: Int = 0
+    @Param("true", "false")
+    var definiteLengthEncoding: Boolean = false
+
+    private var cbor: Cbor = Cbor
+    private var encodedIntArray: ByteArray = ByteArray(0)
+    private var encodedIntMap: ByteArray = ByteArray(0)
+
+    @Setup
+    fun initializeCollections() {
+        cbor = Cbor { useDefiniteLengthEncoding = definiteLengthEncoding }
+        encodedIntArray = cbor.encodeToByteArray(IntArray(collectionSize) { it })
+        encodedIntMap = cbor.encodeToByteArray(IntArray(collectionSize) { it }.associateWith { it })
+    }
+
+    @Benchmark
+    fun array(): IntArray = cbor.decodeFromByteArray(encodedIntArray)
+
+    @Benchmark
+    fun map(): Map<Int, Int> = cbor.decodeFromByteArray(encodedIntMap)
+
+    @Benchmark
+    fun list(): List<Int> = cbor.decodeFromByteArray(encodedIntArray)
+}

core/api/kotlinx-serialization-core.klib.api

@@ -307,15 +307,16 @@ abstract class <#A: kotlin/Any?, #B: kotlin/Any?, #C: kotlinx.serialization.inte
     final val descriptor // kotlinx.serialization.internal/PrimitiveArraySerializer.descriptor|{}descriptor[0]
         final fun <get-descriptor>(): kotlinx.serialization.descriptors/SerialDescriptor // kotlinx.serialization.internal/PrimitiveArraySerializer.descriptor.<get-descriptor>|<get-descriptor>(){}[0]
 
-    abstract fun empty(): #B // kotlinx.serialization.internal/PrimitiveArraySerializer.empty|empty(){}[0]
+    abstract fun ofSize(kotlin/Int): #C // kotlinx.serialization.internal/PrimitiveArraySerializer.ofSize|ofSize(kotlin.Int){}[0]
+    abstract fun ofUnknownSize(): #C // kotlinx.serialization.internal/PrimitiveArraySerializer.ofUnknownSize|ofUnknownSize(){}[0]
     abstract fun readElement(kotlinx.serialization.encoding/CompositeDecoder, kotlin/Int, #C, kotlin/Boolean) // kotlinx.serialization.internal/PrimitiveArraySerializer.readElement|readElement(kotlinx.serialization.encoding.CompositeDecoder;kotlin.Int;1:2;kotlin.Boolean){}[0]
     abstract fun writeContent(kotlinx.serialization.encoding/CompositeEncoder, #B, kotlin/Int) // kotlinx.serialization.internal/PrimitiveArraySerializer.writeContent|writeContent(kotlinx.serialization.encoding.CompositeEncoder;1:1;kotlin.Int){}[0]
     final fun (#B).collectionIterator(): kotlin.collections/Iterator<#A> // kotlinx.serialization.internal/PrimitiveArraySerializer.collectionIterator|collectionIterator@1:1(){}[0]
     final fun (#C).builderSize(): kotlin/Int // kotlinx.serialization.internal/PrimitiveArraySerializer.builderSize|builderSize@1:2(){}[0]
     final fun (#C).checkCapacity(kotlin/Int) // kotlinx.serialization.internal/PrimitiveArraySerializer.checkCapacity|checkCapacity@1:2(kotlin.Int){}[0]
     final fun (#C).insert(kotlin/Int, #A) // kotlinx.serialization.internal/PrimitiveArraySerializer.insert|insert@1:2(kotlin.Int;1:0){}[0]
     final fun (#C).toResult(): #B // kotlinx.serialization.internal/PrimitiveArraySerializer.toResult|toResult@1:2(){}[0]
-    final fun builder(): #C // kotlinx.serialization.internal/PrimitiveArraySerializer.builder|builder(){}[0]
+    final fun builder(kotlin/Int): #C // kotlinx.serialization.internal/PrimitiveArraySerializer.builder|builder(kotlin.Int){}[0]
     final fun deserialize(kotlinx.serialization.encoding/Decoder): #B // kotlinx.serialization.internal/PrimitiveArraySerializer.deserialize|deserialize(kotlinx.serialization.encoding.Decoder){}[0]
     final fun serialize(kotlinx.serialization.encoding/Encoder, #B) // kotlinx.serialization.internal/PrimitiveArraySerializer.serialize|serialize(kotlinx.serialization.encoding.Encoder;1:1){}[0]
 }
@@ -804,11 +805,11 @@ sealed class <#A: kotlin/Any?, #B: kotlin/Any?, #C: kotlin/Any?, #D: kotlin.coll
 sealed class <#A: kotlin/Any?, #B: kotlin/Any?, #C: kotlin/Any?> kotlinx.serialization.internal/AbstractCollectionSerializer : kotlinx.serialization/KSerializer<#B> { // kotlinx.serialization.internal/AbstractCollectionSerializer|null[0]
     abstract fun (#B).collectionIterator(): kotlin.collections/Iterator<#A> // kotlinx.serialization.internal/AbstractCollectionSerializer.collectionIterator|collectionIterator@1:1(){}[0]
     abstract fun (#B).collectionSize(): kotlin/Int // kotlinx.serialization.internal/AbstractCollectionSerializer.collectionSize|collectionSize@1:1(){}[0]
-    abstract fun (#B).toBuilder(): #C // kotlinx.serialization.internal/AbstractCollectionSerializer.toBuilder|toBuilder@1:1(){}[0]
+    abstract fun (#B).toBuilder(kotlin/Int): #C // kotlinx.serialization.internal/AbstractCollectionSerializer.toBuilder|toBuilder@1:1(kotlin.Int){}[0]
     abstract fun (#C).builderSize(): kotlin/Int // kotlinx.serialization.internal/AbstractCollectionSerializer.builderSize|builderSize@1:2(){}[0]
     abstract fun (#C).checkCapacity(kotlin/Int) // kotlinx.serialization.internal/AbstractCollectionSerializer.checkCapacity|checkCapacity@1:2(kotlin.Int){}[0]
     abstract fun (#C).toResult(): #B // kotlinx.serialization.internal/AbstractCollectionSerializer.toResult|toResult@1:2(){}[0]
-    abstract fun builder(): #C // kotlinx.serialization.internal/AbstractCollectionSerializer.builder|builder(){}[0]
+    abstract fun builder(kotlin/Int): #C // kotlinx.serialization.internal/AbstractCollectionSerializer.builder|builder(kotlin.Int){}[0]
     abstract fun readAll(kotlinx.serialization.encoding/CompositeDecoder, #C, kotlin/Int, kotlin/Int) // kotlinx.serialization.internal/AbstractCollectionSerializer.readAll|readAll(kotlinx.serialization.encoding.CompositeDecoder;1:2;kotlin.Int;kotlin.Int){}[0]
     abstract fun readElement(kotlinx.serialization.encoding/CompositeDecoder, kotlin/Int, #C, kotlin/Boolean = ...) // kotlinx.serialization.internal/AbstractCollectionSerializer.readElement|readElement(kotlinx.serialization.encoding.CompositeDecoder;kotlin.Int;1:2;kotlin.Boolean){}[0]
     abstract fun serialize(kotlinx.serialization.encoding/Encoder, #B) // kotlinx.serialization.internal/AbstractCollectionSerializer.serialize|serialize(kotlinx.serialization.encoding.Encoder;1:1){}[0]

core/commonMain/src/kotlinx/serialization/internal/CollectionSerializers.kt

@@ -8,27 +8,33 @@ package kotlinx.serialization.internal
 import kotlinx.serialization.*
 import kotlinx.serialization.descriptors.*
 import kotlinx.serialization.encoding.*
+import kotlin.math.ceil
 import kotlin.reflect.*
 
 @InternalSerializationApi
 public sealed class AbstractCollectionSerializer<Element, Collection, Builder> : KSerializer<Collection> {
     protected abstract fun Collection.collectionSize(): Int
     protected abstract fun Collection.collectionIterator(): Iterator<Element>
-    protected abstract fun builder(): Builder
+    protected abstract fun builder(initialCapacityHint: Int): Builder
     protected abstract fun Builder.builderSize(): Int
     protected abstract fun Builder.toResult(): Collection
-    protected abstract fun Collection.toBuilder(): Builder
+    protected abstract fun Collection.toBuilder(expectedAdditionalSize: Int): Builder
     protected abstract fun Builder.checkCapacity(size: Int)
 
     abstract override fun serialize(encoder: Encoder, value: Collection)
 
     @InternalSerializationApi
     public fun merge(decoder: Decoder, previous: Collection?): Collection {
-        val builder = previous?.toBuilder() ?: builder()
-        val startIndex = builder.builderSize()
         val compositeDecoder = decoder.beginStructure(descriptor)
+        val expectedSize = compositeDecoder.decodeCollectionSize(descriptor)
+        val builder = if (previous != null) {
+            previous.toBuilder(expectedSize)
+        } else {
+            builder(expectedSize)
+        }
+        val startIndex = builder.builderSize()
         if (compositeDecoder.decodeSequentially()) {
-            readAll(compositeDecoder, builder, startIndex, readSize(compositeDecoder, builder))
+            readAll(compositeDecoder, builder, startIndex, expectedSize)
         } else {
             while (true) {
                 val index = compositeDecoder.decodeElementIndex(descriptor)
@@ -42,12 +48,6 @@ public sealed class AbstractCollectionSerializer<Element, Collection, Builder> :
 
     override fun deserialize(decoder: Decoder): Collection = merge(decoder, null)
 
-    private fun readSize(decoder: CompositeDecoder, builder: Builder): Int {
-        val size = decoder.decodeCollectionSize(descriptor)
-        builder.checkCapacity(size)
-        return size
-    }
-
     protected abstract fun readElement(decoder: CompositeDecoder, index: Int, builder: Builder, checkIndex: Boolean = true)
 
     protected abstract fun readAll(decoder: CompositeDecoder, builder: Builder, startIndex: Int, size: Int)
@@ -156,9 +156,14 @@ internal abstract class PrimitiveArraySerializer<Element, Array, Builder
     final override fun Builder.insert(index: Int, element: Element): Unit =
         error("This method lead to boxing and must not be used, use Builder.append instead")
 
-    final override fun builder(): Builder = empty().toBuilder()
+    final override fun builder(initialCapacityHint: Int): Builder {
+        if (initialCapacityHint == -1) return ofUnknownSize()
+        checkBuildersInitialCapacity(initialCapacityHint)
+        return ofSize(initialCapacityHint)
+    }
 
-    protected abstract fun empty(): Array
+    protected abstract fun ofSize(size: Int): Builder
+    protected abstract fun ofUnknownSize(): Builder
 
     abstract override fun readElement(
         decoder: CompositeDecoder,
@@ -189,13 +194,22 @@ internal class ReferenceArraySerializer<ElementKlass : Any, Element : ElementKla
 
     override fun Array<Element>.collectionSize(): Int = size
     override fun Array<Element>.collectionIterator(): Iterator<Element> = iterator()
-    override fun builder(): ArrayList<Element> = arrayListOf()
+    override fun builder(initialCapacityHint: Int): ArrayList<Element> {
+        if (initialCapacityHint == -1) return arrayListOf()
+        checkBuildersInitialCapacity(initialCapacityHint)
+        return ArrayList(initialCapacityHint)
+    }
     override fun ArrayList<Element>.builderSize(): Int = size
 
     @Suppress("UNCHECKED_CAST")
     override fun ArrayList<Element>.toResult(): Array<Element> = toNativeArrayImpl<ElementKlass, Element>(kClass)
 
-    override fun Array<Element>.toBuilder(): ArrayList<Element> = ArrayList(this.asList())
+    override fun Array<Element>.toBuilder(expectedAdditionalSize: Int): ArrayList<Element> {
+        checkBuildersInitialCapacity(expectedAdditionalSize)
+        val expectedCapacity = size + expectedAdditionalSize.coerceAtLeast(0)
+        return ArrayList<Element>(expectedAdditionalSize).also { it.addAll(this) }
+    }
+
     override fun ArrayList<Element>.checkCapacity(size: Int): Unit = ensureCapacity(size)
     override fun ArrayList<Element>.insert(index: Int, element: Element) {
         add(index, element)
@@ -213,10 +227,22 @@ internal abstract class CollectionSerializer<E, C: Collection<E>, B>(element: KS
 internal class ArrayListSerializer<E>(element: KSerializer<E>) : CollectionSerializer<E, List<E>, ArrayList<E>>(element) {
     override val descriptor: SerialDescriptor = ArrayListClassDesc(element.descriptor)
 
-    override fun builder(): ArrayList<E> = arrayListOf()
+    override fun builder(initialCapacityHint: Int): ArrayList<E> {
+        if (initialCapacityHint == -1) return arrayListOf()
+        checkBuildersInitialCapacity(initialCapacityHint)
+        return ArrayList(initialCapacityHint)
+    }
     override fun ArrayList<E>.builderSize(): Int = size
     override fun ArrayList<E>.toResult(): List<E> = this
-    override fun List<E>.toBuilder(): ArrayList<E> = this as? ArrayList<E> ?: ArrayList(this)
+    override fun List<E>.toBuilder(expectedAdditionalSize: Int): ArrayList<E> {
+        if (expectedAdditionalSize != -1) checkBuildersInitialCapacity(expectedAdditionalSize)
+        if (this is ArrayList) {
+            if (expectedAdditionalSize == -1) return this
+            return this.also { it.ensureCapacity(size + expectedAdditionalSize) }
+        }
+        if (expectedAdditionalSize == -1) return ArrayList(this)
+        return ArrayList<E>(size + expectedAdditionalSize).also { it.addAll(this) }
+    }
     override fun ArrayList<E>.checkCapacity(size: Int): Unit = ensureCapacity(size)
     override fun ArrayList<E>.insert(index: Int, element: E) { add(index, element) }
 }
@@ -227,10 +253,20 @@ internal class LinkedHashSetSerializer<E>(
 ) : CollectionSerializer<E, Set<E>, LinkedHashSet<E>>(eSerializer) {
     override val descriptor: SerialDescriptor = LinkedHashSetClassDesc(eSerializer.descriptor)
 
-    override fun builder(): LinkedHashSet<E> = linkedSetOf()
+    override fun builder(initialCapacityHint: Int): LinkedHashSet<E> {
+        if (initialCapacityHint == -1) return linkedSetOf()
+        checkBuildersInitialCapacity(initialCapacityHint)
+        return newLinkedHashSet(initialCapacityHint)
+    }
     override fun LinkedHashSet<E>.builderSize(): Int = size
     override fun LinkedHashSet<E>.toResult(): Set<E> = this
-    override fun Set<E>.toBuilder(): LinkedHashSet<E> = this as? LinkedHashSet<E> ?: LinkedHashSet(this)
+    override fun Set<E>.toBuilder(expectedAdditionalSize: Int): LinkedHashSet<E> {
+        if (this is LinkedHashSet<E>) return this
+        if (expectedAdditionalSize == -1) return LinkedHashSet(this)
+        checkBuildersInitialCapacity(expectedAdditionalSize)
+        return newLinkedHashSet<E>(size + expectedAdditionalSize).also { it.addAll(this) }
+    }
+
     override fun LinkedHashSet<E>.checkCapacity(size: Int) {}
     override fun LinkedHashSet<E>.insert(index: Int, element: E) { add(element) }
 }
@@ -241,10 +277,19 @@ internal class HashSetSerializer<E>(
 ) : CollectionSerializer<E, Set<E>, HashSet<E>>(eSerializer) {
     override val descriptor: SerialDescriptor = HashSetClassDesc(eSerializer.descriptor)
 
-    override fun builder(): HashSet<E> = HashSet()
+    override fun builder(initialCapacityHint: Int): HashSet<E> {
+        if (initialCapacityHint == -1) return HashSet()
+        checkBuildersInitialCapacity(initialCapacityHint)
+        return newHashSet(initialCapacityHint)
+    }
     override fun HashSet<E>.builderSize(): Int = size
     override fun HashSet<E>.toResult(): Set<E> = this
-    override fun Set<E>.toBuilder(): HashSet<E> = this as? HashSet<E> ?: HashSet(this)
+    override fun Set<E>.toBuilder(expectedAdditionalSize: Int): HashSet<E> {
+        if (this is HashSet<E>) return this
+        if (expectedAdditionalSize == -1) return HashSet(this)
+        checkBuildersInitialCapacity(expectedAdditionalSize)
+        return newHashSet<E>(size + expectedAdditionalSize).also { it.addAll(this) }
+    }
     override fun HashSet<E>.checkCapacity(size: Int) {}
     override fun HashSet<E>.insert(index: Int, element: E) { add(element) }
 }
@@ -257,10 +302,20 @@ internal class LinkedHashMapSerializer<K, V>(
     override val descriptor: SerialDescriptor = LinkedHashMapClassDesc(kSerializer.descriptor, vSerializer.descriptor)
     override fun Map<K, V>.collectionSize(): Int = size
     override fun Map<K, V>.collectionIterator(): Iterator<Map.Entry<K, V>> = iterator()
-    override fun builder(): LinkedHashMap<K, V> = LinkedHashMap()
+    override fun builder(initialCapacityHint: Int): LinkedHashMap<K, V> {
+        if (initialCapacityHint == -1) return LinkedHashMap()
+        checkBuildersInitialCapacity(initialCapacityHint)
+        return newLinkedHashMap(initialCapacityHint)
+    }
     override fun LinkedHashMap<K, V>.builderSize(): Int = size * 2
     override fun LinkedHashMap<K, V>.toResult(): Map<K, V> = this
-    override fun Map<K, V>.toBuilder(): LinkedHashMap<K, V> = this as? LinkedHashMap<K, V> ?: LinkedHashMap(this)
+    override fun Map<K, V>.toBuilder(expectedAdditionalSize: Int): LinkedHashMap<K, V> {
+        if (this is LinkedHashMap<K, V>) return this
+        if (expectedAdditionalSize == -1) return LinkedHashMap(this)
+        checkBuildersInitialCapacity(expectedAdditionalSize)
+        return newLinkedHashMap<K, V>(size + expectedAdditionalSize).also { it.putAll(this) }
+    }
+
     override fun LinkedHashMap<K, V>.checkCapacity(size: Int) {}
     override fun LinkedHashMap<K, V>.insertKeyValuePair(index: Int, key: K, value: V): Unit = set(key, value)
 }
@@ -273,10 +328,49 @@ internal class HashMapSerializer<K, V>(
     override val descriptor: SerialDescriptor = HashMapClassDesc(kSerializer.descriptor, vSerializer.descriptor)
     override fun Map<K, V>.collectionSize(): Int = size
     override fun Map<K, V>.collectionIterator(): Iterator<Map.Entry<K, V>> = iterator()
-    override fun builder(): HashMap<K, V> = HashMap()
+    override fun builder(initialCapacityHint: Int): HashMap<K, V> {
+        if (initialCapacityHint == -1) return HashMap()
+        checkBuildersInitialCapacity(initialCapacityHint)
+        return newHashMap(initialCapacityHint)
+    }
     override fun HashMap<K, V>.builderSize(): Int = size * 2
     override fun HashMap<K, V>.toResult(): Map<K, V> = this
-    override fun Map<K, V>.toBuilder(): HashMap<K, V> = this as? HashMap<K, V> ?: HashMap(this)
+    override fun Map<K, V>.toBuilder(expectedAdditionalSize: Int): HashMap<K, V> {
+        if (this is HashMap<K, V>) return this
+        if (expectedAdditionalSize == -1) return HashMap(this)
+        checkBuildersInitialCapacity(expectedAdditionalSize)
+        return newHashMap<K, V>(size + expectedAdditionalSize).also { it.putAll(this) }
+    }
+
     override fun HashMap<K, V>.checkCapacity(size: Int) {}
     override fun HashMap<K, V>.insertKeyValuePair(index: Int, key: K, value: V): Unit = set(key, value)
 }
+
+internal fun checkBuildersInitialCapacity(initialCapacityHint: Int) {
+    require(initialCapacityHint >= 0) { "builder's initial capacity must be non-negative, but was $initialCapacityHint" }
+}
+
+private fun estimateCapacityForHashMap(requiredCapacity: Int, loadFactor: Double): Int =
+    ceil(requiredCapacity / loadFactor).toInt()
+
+private const val DEFAULT_LOAD_FACTORY = 0.75f
+
+private fun <T> newHashSet(expectedCapacity: Int): HashSet<T> {
+    val capacity = estimateCapacityForHashMap(expectedCapacity, DEFAULT_LOAD_FACTORY.toDouble())
+    return HashSet(capacity, DEFAULT_LOAD_FACTORY)
+}
+
+private fun <K, V> newHashMap(expectedCapacity: Int): HashMap<K, V> {
+    val capacity = estimateCapacityForHashMap(expectedCapacity, DEFAULT_LOAD_FACTORY.toDouble())
+    return HashMap(capacity, DEFAULT_LOAD_FACTORY)
+}
+
+private fun <T> newLinkedHashSet(expectedCapacity: Int): LinkedHashSet<T> {
+    val capacity = estimateCapacityForHashMap(expectedCapacity, DEFAULT_LOAD_FACTORY.toDouble())
+    return LinkedHashSet(capacity, DEFAULT_LOAD_FACTORY)
+}
+
+private fun <K, V> newLinkedHashMap(expectedCapacity: Int): LinkedHashMap<K, V> {
+    val capacity = estimateCapacityForHashMap(expectedCapacity, DEFAULT_LOAD_FACTORY.toDouble())
+    return LinkedHashMap(capacity, DEFAULT_LOAD_FACTORY)
+}