Speedup ProtoWriteType.from (#2879)

* Added map-related protobuf benchmark * Reimplemented ProtoWireType lookup
Kotlin · Dec 13, 2024 · d62f3b8 · d62f3b8
1 parent db217e4
commit d62f3b8
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 3 deletions.
diff --git a/benchmark/src/jmh/kotlin/kotlinx/benchmarks/protobuf/ProtoMapBenchmark.kt b/benchmark/src/jmh/kotlin/kotlinx/benchmarks/protobuf/ProtoMapBenchmark.kt
@@ -0,0 +1,32 @@
+/*
+ * Copyright 2017-2024 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.benchmarks.protobuf
+
+import kotlinx.serialization.Serializable
+import kotlinx.serialization.encodeToByteArray
+import kotlinx.serialization.protobuf.ProtoBuf
+import org.openjdk.jmh.annotations.*
+import java.util.concurrent.TimeUnit
+
+@Warmup(iterations = 5, time = 1)
+@Measurement(iterations = 5, time = 1)
+@BenchmarkMode(Mode.Throughput)
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+@State(Scope.Benchmark)
+@Fork(1)
+open class ProtoMapBenchmark {
+
+    @Serializable
+    class Holder(val map: Map<String, Int>)
+
+    private val value = Holder((0..128).associateBy { it.toString() })
+    private val bytes = ProtoBuf.encodeToByteArray(value)
+
+    @Benchmark
+    fun toBytes() = ProtoBuf.encodeToByteArray(Holder.serializer(), value)
+
+    @Benchmark
+    fun fromBytes() = ProtoBuf.decodeFromByteArray(Holder.serializer(), bytes)
+}
diff --git a/formats/protobuf/commonMain/src/kotlinx/serialization/protobuf/internal/Helpers.kt b/formats/protobuf/commonMain/src/kotlinx/serialization/protobuf/internal/Helpers.kt
@@ -22,8 +22,17 @@ internal enum class ProtoWireType(val typeId: Int) {
     ;
 
     companion object {
-        fun from(typeId: Int): ProtoWireType {
-            return ProtoWireType.entries.find { it.typeId == typeId } ?: INVALID
+        private val entryArray = Array(8) { typeId ->
+            ProtoWireType.entries.find { it.typeId == typeId } ?: INVALID
+        }
+
+        /**
+         * Extracts three least significant bits from the [value] and
+         * returns [ProtoWireType] with corresponding type id, or [ProtoWireType.INVALID]
+         * if there are no such a type.
+         */
+        fun fromLeastSignificantBits(value: Int): ProtoWireType {
+            return entryArray[value and 0b111]
         }
     }
 

diff --git a/formats/protobuf/commonMain/src/kotlinx/serialization/protobuf/internal/ProtobufReader.kt b/formats/protobuf/commonMain/src/kotlinx/serialization/protobuf/internal/ProtobufReader.kt
@@ -42,7 +42,7 @@ internal class ProtobufReader(private val input: ByteArrayInput) {
             -1
         } else {
             currentId = header ushr 3
-            currentType = ProtoWireType.from(header and 0b111)
+            currentType = ProtoWireType.fromLeastSignificantBits(header)
             currentId
         }
     }