grpc-native: Implement all packed fixed sized field decoding

Jozott00 · Jozott00 · commit 6fac0b22abca · 2025-07-25T12:36:05.000+02:00
Signed-off-by: Johannes Zottele &lt;official@johannes-zottele.at&gt;
diff --git a/grpc/grpc-core/src/commonMain/kotlin/kotlinx/rpc/grpc/internal/WireDecoder.kt b/grpc/grpc-core/src/commonMain/kotlin/kotlinx/rpc/grpc/internal/WireDecoder.kt
@@ -28,7 +28,7 @@ import kotlinx.io.Buffer
  * }
  * ```
  */
-internal interface WireDecoder: AutoCloseable {
+internal interface WireDecoder : AutoCloseable {
     fun readTag(): KTag?
     fun readBool(): Boolean?
     fun readInt32(): Int?
@@ -47,6 +47,11 @@ internal interface WireDecoder: AutoCloseable {
     fun readString(): String?
     fun readBytes(): ByteArray?
     fun readPackedFixed32(): List<UInt>?
+    fun readPackedFixed64(): List<ULong>?
+    fun readPackedSFixed32(): List<Int>?
+    fun readPackedSFixed64(): List<Long>?
+    fun readPackedFloat(): List<Float>?
+    fun readPackedDouble(): List<Double>?
 }
 
 /**
diff --git a/grpc/grpc-core/src/commonMain/kotlin/kotlinx/rpc/grpc/internal/WireEncoder.kt b/grpc/grpc-core/src/commonMain/kotlin/kotlinx/rpc/grpc/internal/WireEncoder.kt
@@ -14,7 +14,9 @@ import kotlinx.io.Sink
  *
  * [flush] must be called to ensure that all data is written to the [Sink].
  */
+@OptIn(ExperimentalUnsignedTypes::class)
 internal interface WireEncoder {
+    fun flush()
     fun writeBool(field: Int, value: Boolean): Boolean
     fun writeInt32(fieldNr: Int, value: Int): Boolean
     fun writeInt64(fieldNr: Int, value: Long): Boolean
@@ -29,10 +31,14 @@ internal interface WireEncoder {
     fun writeFloat(fieldNr: Int, value: Float): Boolean
     fun writeDouble(fieldNr: Int, value: Double): Boolean
     fun writeEnum(fieldNr: Int, value: Int): Boolean
-    fun writeString(fieldNr: Int, value: String): Boolean
-    fun flush()
     fun writeBytes(fieldNr: Int, value: ByteArray): Boolean
+    fun writeString(fieldNr: Int, value: String): Boolean
     fun writePackedFixed32(fieldNr: Int, value: UIntArray): Boolean
+    fun writePackedFixed64(fieldNr: Int, value: ULongArray): Boolean
+    fun writePackedSFixed32(fieldNr: Int, value: IntArray): Boolean
+    fun writePackedSFixed64(fieldNr: Int, value: LongArray): Boolean
+    fun writePackedFloat(fieldNr: Int, value: FloatArray): Boolean
+    fun writePackedDouble(fieldNr: Int, value: DoubleArray): Boolean
 }
 
 
diff --git a/grpc/grpc-core/src/nativeMain/kotlin/kotlinx/rpc/grpc/internal/WireDecoder.native.kt b/grpc/grpc-core/src/nativeMain/kotlin/kotlinx/rpc/grpc/internal/WireDecoder.native.kt
@@ -11,10 +11,11 @@ import libprotowire.*
 import kotlin.experimental.ExperimentalNativeApi
 import kotlin.math.min
 
+// maximum buffer size to allocate as contiguous memory in bytes
 private const val MAX_PACKED_BULK_SIZE: Int = 1_000_000
 
 @OptIn(ExperimentalForeignApi::class, ExperimentalNativeApi::class)
-internal class WireDecoderNative(private val source: Buffer): WireDecoder {
+internal class WireDecoderNative(private val source: Buffer) : WireDecoder {
 
     // wraps the source in a class that allows to pass data from the source buffer to the C++ encoder
     // without copying it to an intermediate byte array.
@@ -201,6 +202,48 @@ internal class WireDecoderNative(private val source: Buffer): WireDecoder {
         return bytes
     }
 
+    override fun readPackedFixed32() = readPackedFixedInternal(
+        UInt.SIZE_BYTES,
+        ::UIntArray,
+        Pinned<UIntArray>::addressOf,
+        UIntArray::asList,
+    )
+
+    override fun readPackedFixed64() = readPackedFixedInternal(
+        ULong.SIZE_BYTES,
+        ::ULongArray,
+        Pinned<ULongArray>::addressOf,
+        ULongArray::asList,
+    )
+
+    override fun readPackedSFixed32() = readPackedFixedInternal(
+        Int.SIZE_BYTES,
+        ::IntArray,
+        Pinned<IntArray>::addressOf,
+        IntArray::asList,
+    )
+
+    override fun readPackedSFixed64() = readPackedFixedInternal(
+        Long.SIZE_BYTES,
+        ::LongArray,
+        Pinned<LongArray>::addressOf,
+        LongArray::asList,
+    )
+
+    override fun readPackedFloat() = readPackedFixedInternal(
+        Float.SIZE_BYTES,
+        ::FloatArray,
+        Pinned<FloatArray>::addressOf,
+        FloatArray::asList,
+    )
+
+    override fun readPackedDouble() = readPackedFixedInternal(
+        Double.SIZE_BYTES,
+        ::DoubleArray,
+        Pinned<DoubleArray>::addressOf,
+        DoubleArray::asList,
+    )
+    
     /*
      * Based on the length of the packed repeated field, one of two list strategies is chosen.
      * If the length is less or equal a specific threshold (MAX_PACKED_BULK_SIZE),
@@ -210,39 +253,52 @@ internal class WireDecoderNative(private val source: Buffer): WireDecoder {
      *
      * Note that this implementation assumes a little endian memory order.
      */
-    override fun readPackedFixed32(): List<UInt>? {
+    private inline fun <T : Any, R : Any> readPackedFixedInternal(
+        sizeBytes: Int,
+        crossinline createArray: (Int) -> R,
+        crossinline getAddress: Pinned<R>.(Int) -> COpaquePointer,
+        crossinline asList: (R) -> List<T>
+    ): List<T>? {
+        // fetch the size of the packed repeated field
         var byteLen = readInt32() ?: return null
         if (byteLen < 0) return null
         if (source.size < byteLen) return null
-        if (byteLen % UInt.SIZE_BYTES != 0 ) return null
-        val count = byteLen / UInt.SIZE_BYTES
+        if (byteLen % sizeBytes != 0) return null
         if (byteLen == 0) return emptyList()
 
-        if (count <= MAX_PACKED_BULK_SIZE) {
-            // this implementation assumes that the program is running on little endian machines.
-            val arr = UIntArray(count)
-            arr.usePinned {
-                pw_decoder_read_raw_bytes(raw, it.addressOf(0), byteLen)
-            }
-            return ArrayList(arr)
-        } else {
-            val bufByteLen = MAX_PACKED_BULK_SIZE
-            val bufLen = bufByteLen / UInt.SIZE_BYTES
-            val buffer = UIntArray(bufLen)
-            var list = persistentListOf<UInt>()
-            buffer.usePinned {
+        // allocate the buffer array (has at most MAX_PACKED_BULK_SIZE bytes)
+        val bufByteLen = minOf(byteLen, MAX_PACKED_BULK_SIZE)
+        val bufElemCount = bufByteLen / sizeBytes
+        val buffer = createArray(bufElemCount)
+
+        buffer.usePinned {
+            val bufAddr = it.getAddress(0)
+
+            if (byteLen == bufByteLen) {
+                // the whole packed field fits into the buffer -> copy into buffer and returns it as a list.
+                pw_decoder_read_raw_bytes(raw, bufAddr, byteLen)
+                return asList(buffer)
+            } else {
+                // the packed field is too large for the buffer, so we load it into a persistent list
+                var chunkedList = persistentListOf<T>()
+
                 while (byteLen > 0) {
-                    val written = min(bufByteLen, byteLen)
-                    pw_decoder_read_raw_bytes(raw, it.addressOf(0), written)
-                    list = if (written == bufByteLen) {
-                        list.addAll(buffer)
+                    // copy data into the buffer.
+                    val copySize = min(bufByteLen, byteLen)
+                    pw_decoder_read_raw_bytes(raw, bufAddr, copySize)
+
+                    // add buffer to the chunked list
+                    chunkedList = if (copySize == bufByteLen) {
+                        chunkedList.addAll(asList(buffer))
                     } else {
-                        list.addAll(buffer.copyOfRange(0, written / UInt.SIZE_BYTES))
+                        chunkedList.addAll(asList(buffer).subList(0, copySize / sizeBytes))
                     }
-                    byteLen -= written
+
+                    byteLen -= copySize
                 }
+
+                return chunkedList
             }
-            return list
         }
     }
 }
diff --git a/grpc/grpc-core/src/nativeMain/kotlin/kotlinx/rpc/grpc/internal/WireEncoder.native.kt b/grpc/grpc-core/src/nativeMain/kotlin/kotlinx/rpc/grpc/internal/WireEncoder.native.kt
@@ -7,16 +7,14 @@ package kotlinx.rpc.grpc.internal
 import kotlinx.cinterop.*
 import kotlinx.io.Sink
 import libprotowire.*
-import kotlin.collections.isEmpty
 import kotlin.experimental.ExperimentalNativeApi
 import kotlin.native.ref.createCleaner
-import kotlin.text.isEmpty
 
 
 // TODO: Evaluate if we should implement a ZeroCopyOutputSink (similar to the ZeroCopyInputSource)
 //      to reduce the number of copies during encoding.
 @OptIn(ExperimentalForeignApi::class, ExperimentalNativeApi::class)
-internal class WireEncoderNative(private val sink: Sink): WireEncoder {
+internal class WireEncoderNative(private val sink: Sink) : WireEncoder {
     /**
      * The context object provides a stable reference to the kotlin context.
      * This is required, as functions must be static and cannot capture environment references.
@@ -32,6 +30,7 @@ internal class WireEncoderNative(private val sink: Sink): WireEncoder {
 
     // create context as a stable reference that can be passed to static function callback
     private val context = StableRef.create(this.Ctx())
+
     // construct encoder with a callback that calls write() on this.context
     internal val raw = run {
         pw_encoder_new(context.asCPointer(), staticCFunction { ctx, buf, size ->
@@ -126,16 +125,46 @@ internal class WireEncoderNative(private val sink: Sink): WireEncoder {
         }
     }
 
-    override fun writePackedFixed32(fieldNr: Int, value: UIntArray): Boolean {
-        if (value.isEmpty()) {
-            return pw_encoder_write_bytes(raw, fieldNr, null, 0)
-        }
-        val bytes = value.size * UInt.SIZE_BYTES
-        return value.usePinned {
-            pw_encoder_write_bytes(raw, fieldNr, it.addressOf(0), bytes)
-        }
-    }
+    override fun writePackedFixed32(fieldNr: Int, value: UIntArray) =
+        writePackedInternal(fieldNr, value, UIntArray::size, UInt.SIZE_BYTES)
+        { it.addressOf(0) }
 
+    override fun writePackedFixed64(fieldNr: Int, value: ULongArray) =
+        writePackedInternal(fieldNr, value, ULongArray::size, ULong.SIZE_BYTES)
+        { it.addressOf(0) }
+
+    override fun writePackedSFixed32(fieldNr: Int, value: IntArray) =
+        writePackedInternal(fieldNr, value, IntArray::size, Int.SIZE_BYTES)
+        { it.addressOf(0) }
+
+    override fun writePackedSFixed64(fieldNr: Int, value: LongArray) =
+        writePackedInternal(fieldNr, value, LongArray::size, Long.SIZE_BYTES)
+        { it.addressOf(0) }
+
+    override fun writePackedFloat(fieldNr: Int, value: FloatArray) =
+        writePackedInternal(fieldNr, value, FloatArray::size, Float.SIZE_BYTES)
+        { it.addressOf(0) }
+
+    override fun writePackedDouble(fieldNr: Int, value: DoubleArray) =
+        writePackedInternal(fieldNr, value, DoubleArray::size, Double.SIZE_BYTES)
+        { it.addressOf(0) }
 }
 
-internal actual fun WireEncoder(sink: Sink): WireEncoder = WireEncoderNative(sink)
+internal actual fun WireEncoder(sink: Sink): WireEncoder = WireEncoderNative(sink)
+
+
+@OptIn(ExperimentalForeignApi::class)
+private inline fun <A : Any> WireEncoderNative.writePackedInternal(
+    fieldNr: Int,
+    value: A,
+    crossinline sizeOf: A.() -> Int,
+    sizeBytes: Int,
+    crossinline ptr: (Pinned<A>) -> COpaquePointer
+): Boolean {
+    val len = sizeOf(value)
+    if (len == 0) return pw_encoder_write_bytes(raw, fieldNr, null, 0)
+    val bytes = len * sizeBytes
+    return value.usePinned { pinned ->
+        pw_encoder_write_bytes(raw, fieldNr, ptr(pinned), bytes)
+    }
+}
diff --git a/grpc/grpc-core/src/nativeTest/kotlin/kotlinx/rpc/grpc/internal/WireCodecTest.kt b/grpc/grpc-core/src/nativeTest/kotlin/kotlinx/rpc/grpc/internal/WireCodecTest.kt
@@ -5,19 +5,9 @@
 package kotlinx.rpc.grpc.internal
 
 import kotlinx.cinterop.ExperimentalForeignApi
-import kotlinx.cinterop.addressOf
-import kotlinx.cinterop.convert
-import kotlinx.cinterop.readBytes
-import kotlinx.cinterop.usePinned
 import kotlinx.io.Buffer
-import platform.posix.memcpy
 import kotlin.experimental.ExperimentalNativeApi
-import kotlin.test.Test
-import kotlin.test.assertEquals
-import kotlin.test.assertFalse
-import kotlin.test.assertNotNull
-import kotlin.test.assertNull
-import kotlin.test.assertTrue
+import kotlin.test.*
 
 // TODO: Move this to the commonTest
 @OptIn(ExperimentalForeignApi::class, ExperimentalNativeApi::class)
@@ -624,31 +614,6 @@ class WireCodecTest {
         assertTrue(buffer.exhausted())
     }
 
-    @Test
-    fun testPackedFixed32EncodeDecode() {
-        val buffer = Buffer()
-        val encoder = WireEncoder(buffer)
-
-        val ints = UIntArray(100000) { it.toUInt() }
-
-        assertTrue(encoder.writePackedFixed32(1, ints))
-        encoder.flush()
-
-        val decoder = WireDecoder(buffer)
-        val tag = decoder.readTag()
-        assertNotNull(tag)
-        assertEquals(1, tag.fieldNr)
-        assertEquals(WireType.LENGTH_DELIMITED, tag.wireType)
-
-        val actualInts = decoder.readPackedFixed32()
-        assertNotNull(actualInts)
-        assertEquals(ints.size, actualInts.size)
-        assertEquals(ints.toList(), actualInts)
-
-        decoder.close()
-        assertTrue(buffer.exhausted())
-    }
-
     @Test
     fun testDoubleEncodeDecode() {
         val fieldNr = 21
@@ -698,4 +663,75 @@ class WireCodecTest {
         decoder.close()
         assertTrue(buffer.exhausted())
     }
+
+
+    private inline fun <A, T> runPackedTest(
+        arr: A,
+        crossinline write: WireEncoder.(Int, A) -> Boolean,
+        crossinline read: WireDecoder.() -> List<T>?,
+        crossinline asList: (A) -> List<T>
+    ) {
+        val buf = Buffer()
+        with(WireEncoder(buf)) {
+            assertTrue(write(1, arr))
+            flush()
+        }
+        WireDecoder(buf).use { dec ->
+            dec.readTag()!!.apply {
+                assertEquals(1, fieldNr)
+                assertEquals(WireType.LENGTH_DELIMITED, wireType)
+            }
+            val test = dec.read()
+            assertEquals(asList(arr), test)
+        }
+        assertTrue(buf.exhausted())
+    }
+
+    @Test
+    fun testPackedFixed32() = runPackedTest(
+        UIntArray(1_000_000) { UInt.MAX_VALUE + it.toUInt() },
+        WireEncoder::writePackedFixed32,
+        WireDecoder::readPackedFixed32,
+        UIntArray::asList
+    )
+
+    @Test
+    fun testPackedFixed64() = runPackedTest(
+        ULongArray(1_000_000) { UInt.MAX_VALUE + it.toULong() },
+        WireEncoder::writePackedFixed64,
+        WireDecoder::readPackedFixed64,
+        ULongArray::asList
+    )
+
+    @Test
+    fun testPackedSFixed32() = runPackedTest(
+        IntArray(1_000_000) { Int.MAX_VALUE + it },
+        WireEncoder::writePackedSFixed32,
+        WireDecoder::readPackedSFixed32,
+        IntArray::asList
+    )
+
+    @Test
+    fun testPackedSFixed64() = runPackedTest(
+        LongArray(1_000_000) { Long.MAX_VALUE + it },
+        WireEncoder::writePackedSFixed64,
+        WireDecoder::readPackedSFixed64,
+        LongArray::asList
+    )
+
+    @Test
+    fun testPackedFloat() = runPackedTest(
+        FloatArray(1_000_000) { it.toFloat() / 3.3f * ((it and 1) * 2 - 1) },
+        WireEncoder::writePackedFloat,
+        WireDecoder::readPackedFloat,
+        FloatArray::asList
+    )
+
+    @Test
+    fun testPackedDouble() = runPackedTest(
+        DoubleArray(1_000_000) { it.toDouble() / 3.3 * ((it and 1) * 2 - 1) },
+        WireEncoder::writePackedDouble,
+        WireDecoder::readPackedDouble,
+        DoubleArray::asList
+    )
 }
