Proof of concept: Add implementation for computing the required bytes to encode a message

This commit introduces a new API which computes the required bytes to
encode a message without actually serializing a message. By extension,
this API is meant to be used as a cornerstone for implementing Kotlin#2075.

Signed-off-by: George Papadopoulos <[email protected]>

Author: GeorgePap-719

formats/protobuf/api/kotlinx-serialization-protobuf.api

@@ -21,6 +21,10 @@ public final class kotlinx/serialization/protobuf/ProtoBufKt {
 	public static synthetic fun ProtoBuf$default (Lkotlinx/serialization/protobuf/ProtoBuf;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lkotlinx/serialization/protobuf/ProtoBuf;
 }
 
+public final class kotlinx/serialization/protobuf/ProtoBufSerializedSizeCalculatorKt {
+	public static final fun getOrComputeSerializedSize (Lkotlinx/serialization/protobuf/ProtoBuf;Lkotlinx/serialization/SerializationStrategy;Ljava/lang/Object;)I
+}
+
 public final class kotlinx/serialization/protobuf/ProtoIntegerType : java/lang/Enum {
 	public static final field DEFAULT Lkotlinx/serialization/protobuf/ProtoIntegerType;
 	public static final field FIXED Lkotlinx/serialization/protobuf/ProtoIntegerType;

formats/protobuf/commonMain/src/kotlinx/serialization/protobuf/ProtoBufSerializedSize.kt

@@ -0,0 +1,82 @@
+package kotlinx.serialization.protobuf.internal
+
+internal fun computeEnumSizeNoTag(value: Int): Int = computeInt32SizeNoTag(value)
+
+internal fun computeByteArraySizeNoTag(value: ByteArray): Int = computeLengthDelimitedFieldSize(value.size)
+
+internal fun computeStringSizeNoTag(value: String): Int {
+    // java's implementation uses a custom method for some optimizations.
+    return computeLengthDelimitedFieldSize(value.length)
+}
+
+internal fun computeLengthDelimitedFieldSize(length: Int): Int = computeUInt32SizeNoTag(length) + length
+
+//TODO: should this also be named "compute" for consistency?
+internal fun getFixed64SizeNoTag(): Int = FIXED64_SIZE
+internal fun computeSInt64SizeNoTag(value: Long): Int = computeUInt64SizeNoTag(encodeZigZag64(value))
+internal fun computeInt64SizeNoTag(value: Long): Int = computeUInt64SizeNoTag(value)
+
+//TODO: should this also be named "compute" for consistency?
+internal fun getFixed32SizeNoTag() = FIXED32_SIZE
+internal fun computeSInt32SizeNoTag(value: Int) = computeUInt32SizeNoTag((encodeZigZag32(value)))
+internal fun computeInt32SizeNoTag(value: Int) =
+    if (value >= 0) computeUInt32SizeNoTag(value) else MAX_VARINT_SIZE
+
+/** Compute the number of bytes that would be needed to encode an uint32 field. */
+internal fun computeUInt32SizeNoTag(value: Int): Int = when {
+    value and (0.inv() shl 7) == 0 -> 1
+    value and (0.inv() shl 14) == 0 -> 2
+    value and (0.inv() shl 21) == 0 -> 3
+    value and (0.inv() shl 28) == 0 -> 4
+    else -> 5 // max varint32 size
+}
+
+private fun computeUInt64SizeNoTag(value: Long): Int {
+    var _value = value
+    // handle first two most common cases
+    if ((_value and (0L.inv() shl 7)) == 0L) {
+        return 1
+    }
+    if (_value < 0L) {
+        return 10
+    }
+    // rest cases
+    var size = 2
+    if ((_value and (0L.inv() shl 35)) != 0L) {
+        size += 4
+        _value = value ushr 28
+    }
+    if ((_value and (0L.inv() shl 21)) != 0L) {
+        size += 2
+        _value = value ushr 14
+    }
+    if ((_value and (0L.inv() shl 14)) != 0L) {
+        size += 1
+    }
+    return size
+}
+
+// helpers
+
+// per protobuf spec 1-10 bytes
+internal const val MAX_VARINT_SIZE = 10
+
+// after decoding the varint representing a field, the low 3 bits tell us the wire type,
+// and the rest of the integer tells us the field number.
+private const val TAG_TYPE_BITS = 3
+
+/**
+ * See [Scalar type values](https://developers.google.com/protocol-buffers/docs/proto#scalar).
+ */
+
+private const val FIXED32_SIZE = 4
+private const val FIXED64_SIZE = 8
+
+internal fun computeTagSize(protoId: Int): Int = computeUInt32SizeNoTag(makeTag(protoId, 0))
+private fun makeTag(protoId: Int, wireType: Int): Int = protoId shl TAG_TYPE_BITS or wireType
+
+// stream utils
+
+internal fun encodeZigZag64(value: Long): Long = (value shl 1) xor (value shr 63)
+
+internal fun encodeZigZag32(value: Int): Int = ((value shl 1) xor (value shr 31))

formats/protobuf/commonMain/src/kotlinx/serialization/protobuf/internal/ProtobufSizeComputers.kt

@@ -33,61 +33,61 @@ internal abstract class ProtobufTaggedEncoder : ProtobufTaggedBase(), Encoder, C
 
     protected open fun encodeTaggedInline(tag: ProtoDesc, inlineDescriptor: SerialDescriptor): Encoder = this.apply { pushTag(tag) }
 
-    public final override fun encodeNull() {
+    final override fun encodeNull() {
         if (nullableMode != NullableMode.ACCEPTABLE) {
             val message = when (nullableMode) {
                 NullableMode.OPTIONAL -> "'null' is not supported for optional properties in ProtoBuf"
                 NullableMode.COLLECTION -> "'null' is not supported for collection types in ProtoBuf"
                 NullableMode.NOT_NULL -> "'null' is not allowed for not-null properties"
-                else -> "'null' is not supported in ProtoBuf";
+                else -> "'null' is not supported in ProtoBuf"
             }
             throw SerializationException(message)
         }
     }
 
-    public final override fun encodeBoolean(value: Boolean) {
+    final override fun encodeBoolean(value: Boolean) {
         encodeTaggedBoolean(popTagOrDefault(), value)
     }
 
-    public final override fun encodeByte(value: Byte) {
+    final override fun encodeByte(value: Byte) {
         encodeTaggedByte(popTagOrDefault(), value)
     }
 
-    public final override fun encodeShort(value: Short) {
+    final override fun encodeShort(value: Short) {
         encodeTaggedShort(popTagOrDefault(), value)
     }
 
-    public final override fun encodeInt(value: Int) {
+    final override fun encodeInt(value: Int) {
         encodeTaggedInt(popTagOrDefault(), value)
     }
 
-    public final override fun encodeLong(value: Long) {
+    final override fun encodeLong(value: Long) {
         encodeTaggedLong(popTagOrDefault(), value)
     }
 
-    public final override fun encodeFloat(value: Float) {
+    final override fun encodeFloat(value: Float) {
         encodeTaggedFloat(popTagOrDefault(), value)
     }
 
-    public final override fun encodeDouble(value: Double) {
+    final override fun encodeDouble(value: Double) {
         encodeTaggedDouble(popTagOrDefault(), value)
     }
 
-    public final override fun encodeChar(value: Char) {
+    final override fun encodeChar(value: Char) {
         encodeTaggedChar(popTagOrDefault(), value)
     }
 
-    public final override fun encodeString(value: String) {
+    final override fun encodeString(value: String) {
         encodeTaggedString(popTagOrDefault(), value)
     }
 
-    public final override fun encodeEnum(
+    final override fun encodeEnum(
         enumDescriptor: SerialDescriptor,
         index: Int
     ): Unit = encodeTaggedEnum(popTagOrDefault(), enumDescriptor, index)
 
 
-    public final override fun endStructure(descriptor: SerialDescriptor) {
+    final override fun endStructure(descriptor: SerialDescriptor) {
         if (stackSize >= 0) {
             popTag()
         }
@@ -96,46 +96,48 @@ internal abstract class ProtobufTaggedEncoder : ProtobufTaggedBase(), Encoder, C
 
     protected open fun endEncode(descriptor: SerialDescriptor) {}
 
-    public final override fun encodeBooleanElement(descriptor: SerialDescriptor, index: Int, value: Boolean): Unit =
+    final override fun encodeBooleanElement(descriptor: SerialDescriptor, index: Int, value: Boolean): Unit =
         encodeTaggedBoolean(descriptor.getTag(index), value)
 
-    public final override fun encodeByteElement(descriptor: SerialDescriptor, index: Int, value: Byte): Unit =
+    final override fun encodeByteElement(descriptor: SerialDescriptor, index: Int, value: Byte): Unit =
         encodeTaggedByte(descriptor.getTag(index), value)
 
-    public final override fun encodeShortElement(descriptor: SerialDescriptor, index: Int, value: Short): Unit =
+    final override fun encodeShortElement(descriptor: SerialDescriptor, index: Int, value: Short): Unit =
         encodeTaggedShort(descriptor.getTag(index), value)
 
-    public final override fun encodeIntElement(descriptor: SerialDescriptor, index: Int, value: Int): Unit =
+    final override fun encodeIntElement(descriptor: SerialDescriptor, index: Int, value: Int): Unit =
         encodeTaggedInt(descriptor.getTag(index), value)
 
-    public final override fun encodeLongElement(descriptor: SerialDescriptor, index: Int, value: Long): Unit =
+    final override fun encodeLongElement(descriptor: SerialDescriptor, index: Int, value: Long): Unit =
         encodeTaggedLong(descriptor.getTag(index), value)
 
-    public final override fun encodeFloatElement(descriptor: SerialDescriptor, index: Int, value: Float): Unit =
+    final override fun encodeFloatElement(descriptor: SerialDescriptor, index: Int, value: Float): Unit =
         encodeTaggedFloat(descriptor.getTag(index), value)
 
-    public final override fun encodeDoubleElement(descriptor: SerialDescriptor, index: Int, value: Double): Unit =
+    final override fun encodeDoubleElement(descriptor: SerialDescriptor, index: Int, value: Double): Unit =
         encodeTaggedDouble(descriptor.getTag(index), value)
 
-    public final override fun encodeCharElement(descriptor: SerialDescriptor, index: Int, value: Char): Unit =
+    final override fun encodeCharElement(descriptor: SerialDescriptor, index: Int, value: Char): Unit =
         encodeTaggedChar(descriptor.getTag(index), value)
 
-    public final override fun encodeStringElement(descriptor: SerialDescriptor, index: Int, value: String): Unit =
+    final override fun encodeStringElement(descriptor: SerialDescriptor, index: Int, value: String): Unit =
         encodeTaggedString(descriptor.getTag(index), value)
 
-    public final override fun <T : Any?> encodeSerializableElement(
+    final override fun <T : Any?> encodeSerializableElement(
         descriptor: SerialDescriptor,
         index: Int,
         serializer: SerializationStrategy<T>,
         value: T
     ) {
         nullableMode = NullableMode.NOT_NULL
-
-        pushTag(descriptor.getTag(index))
+        val tag = descriptor.getTag(index)
+        println("will push tag:$tag in stack")
+        pushTag(tag)
+        println("total stackSize:${stackSize + 1}")
         encodeSerializableValue(serializer, value)
     }
 
-    public final override fun <T : Any> encodeNullableSerializableElement(
+    final override fun <T : Any> encodeNullableSerializableElement(
         descriptor: SerialDescriptor,
         index: Int,
         serializer: SerializationStrategy<T>,

formats/protobuf/commonMain/src/kotlinx/serialization/protobuf/internal/ProtobufTaggedEncoder.kt

@@ -0,0 +1,15 @@
+package kotlinx.serialization.protobuf
+
+import kotlinx.serialization.descriptors.SerialDescriptor
+
+internal actual fun createSerializedSizeCache(): SerializedSizeCache = JsHashMap()
+
+private class JsHashMap : SerializedSizeCache {
+    private val cache = mutableMapOf<SerialDescriptor, SerializedData>()
+
+    override fun get(descriptor: SerialDescriptor, key: SerializedSizeCacheKey): Int? = cache[descriptor]?.get(key)
+
+    override fun set(descriptor: SerialDescriptor, key: SerializedSizeCacheKey, serializedSize: Int) {
+        cache[descriptor] = mapOf(key to serializedSize)
+    }
+}

formats/protobuf/jsMain/src/kotlinx/serialization/protobuf/SerializedSizeCache.kt

@@ -0,0 +1,18 @@
+package kotlinx.serialization.protobuf
+
+import kotlinx.serialization.descriptors.SerialDescriptor
+import java.util.concurrent.ConcurrentHashMap
+
+internal actual fun createSerializedSizeCache(): SerializedSizeCache {
+    return ConcurrentHashMapSerializedCache()
+}
+
+private class ConcurrentHashMapSerializedCache : SerializedSizeCache {
+    private val cache = ConcurrentHashMap<SerialDescriptor, SerializedData>()
+
+    override fun get(descriptor: SerialDescriptor, key: SerializedSizeCacheKey): Int? = cache[descriptor]?.get(key)
+
+    override fun set(descriptor: SerialDescriptor, key: SerializedSizeCacheKey, serializedSize: Int) {
+        cache[descriptor] = mapOf(key to serializedSize)
+    }
+}