grpc-pb: Support message in oneof

Signed-off-by: Johannes Zottele <[email protected]>

Author: Jozott00

grpc/grpc-core/src/commonTest/kotlin/kotlinx/rpc/grpc/pb/ProtosTest.kt

@@ -22,10 +22,7 @@ import test.recursive.RecursiveInternal
 import test.recursive.RecursiveReq
 import test.recursive.invoke
 import test.submsg.*
-import kotlin.test.Test
-import kotlin.test.assertEquals
-import kotlin.test.assertFailsWith
-import kotlin.test.assertNull
+import kotlin.test.*
 
 class ProtosTest {
 
@@ -155,6 +152,15 @@ class ProtosTest {
         }
         val decoded2 = encodeDecode(msg2, OneOfMsgInternal.CODEC)
         assertEquals(OneOfMsg.Field.Fixed(21u), decoded2.field)
+
+        val msg3 = OneOfMsg {
+            field = OneOfMsg.Field.Other(Other { arg2 = "test" })
+        }
+        val decoded3 = encodeDecode(msg3, OneOfMsgInternal.CODEC)
+        assertIs<OneOfMsg.Field.Other>(decoded3.field)
+        assertNull((decoded3.field as OneOfMsg.Field.Other).value.arg1)
+        assertEquals("test", (decoded3.field as OneOfMsg.Field.Other).value.arg2)
+        assertNull((decoded3.field as OneOfMsg.Field.Other).value.arg3)
     }
 
     @Test

grpc/grpc-core/src/commonTest/proto/oneof.proto

@@ -1,6 +1,9 @@
+import "sub_message.proto";
+
 message OneOfMsg {
   oneof field {
     int32 sint = 2;
     fixed64 fixed = 3;
+    test.submsg.Other other = 4;
   }
 }

protoc-gen/src/main/kotlin/kotlinx/rpc/protobuf/ModelToKotlinCommonGenerator.kt

@@ -282,44 +282,68 @@ class ModelToKotlinCommonGenerator(
     private fun CodeGenerator.readMatchCase(
         field: FieldDeclaration,
         lvalue: String = "msg.${field.name}",
-        wrapperCtor: (String) -> String = { it }
+        wrapperCtor: (String) -> String = { it },
+        beforeValueDecoding: CodeGenerator.() -> Unit = {},
     ) {
         when (val fieldType = field.type) {
             is FieldType.IntegralType -> whenCase("tag.fieldNr == ${field.number} && tag.wireType == $PB_PKG.WireType.${field.type.wireType.name}") {
+                beforeValueDecoding()
                 generateDecodeFieldValue(fieldType, lvalue, wrapperCtor = wrapperCtor)
             }
 
             is FieldType.List -> if (field.dec.isPacked) {
                 whenCase("tag.fieldNr == ${field.number} && tag.wireType == $PB_PKG.WireType.LENGTH_DELIMITED") {
+                    beforeValueDecoding()
                     generateDecodeFieldValue(fieldType, lvalue, isPacked = true, wrapperCtor = wrapperCtor)
                 }
             } else {
                 whenCase("tag.fieldNr == ${field.number} && tag.wireType == $PB_PKG.WireType.${fieldType.value.wireType.name}") {
+                    beforeValueDecoding()
                     generateDecodeFieldValue(fieldType, lvalue, isPacked = false, wrapperCtor = wrapperCtor)
                 }
             }
 
             is FieldType.Enum -> whenCase("tag.fieldNr == ${field.number} && tag.wireType == $PB_PKG.WireType.VARINT") {
+                beforeValueDecoding()
                 generateDecodeFieldValue(fieldType, lvalue, wrapperCtor = wrapperCtor)
             }
 
             is FieldType.OneOf -> {
                 fieldType.dec.variants.forEach { variant ->
                     val variantName = "${fieldType.dec.name.safeFullName()}.${variant.name}"
-                    readMatchCase(
-                        field = variant,
-                        lvalue = lvalue,
-                        wrapperCtor = { "$variantName($it)" }
-                    )
+                    if (variant.type is FieldType.Message) {
+                        // in case of a message, we must construct an empty message before reading the message
+                        readMatchCase(
+                            field = variant,
+                            lvalue = "field.value",
+                            beforeValueDecoding = {
+                                beforeValueDecoding()
+                                scope("val field = ($lvalue as? $variantName) ?: $variantName(${variant.type.internalCtor()}).also") {
+                                    // write the constructed oneof variant to the field
+                                    code("$lvalue = it")
+                                }
+                            })
+                    } else {
+                        readMatchCase(
+                            field = variant,
+                            lvalue = lvalue,
+                            wrapperCtor = { "$variantName($it)" },
+                            beforeValueDecoding = beforeValueDecoding
+                        )
+                    }
                 }
             }
 
             is FieldType.Message -> {
                 whenCase("tag.fieldNr == ${field.number} && tag.wireType == $PB_PKG.WireType.LENGTH_DELIMITED") {
-                    // check if the the current sub message object
-                    ifBranch(condition = "!msg.presenceMask[${field.presenceIdx}]", ifBlock = {
-                        code("$lvalue = ${fieldType.dec.value.internalClassFullName()}()")
-                    })
+                    if (field.presenceIdx != null) {
+                        // check if the current sub message object was already set, if not, set a new one
+                        // to set the field's presence tracker to true
+                        ifBranch(condition = "!msg.presenceMask[${field.presenceIdx}]", ifBlock = {
+                            code("$lvalue = ${fieldType.dec.value.internalClassFullName()}()")
+                        })
+                    }
+                    beforeValueDecoding()
                     generateDecodeFieldValue(fieldType, lvalue, wrapperCtor = wrapperCtor)
                 }
             }
@@ -820,6 +844,9 @@ class ModelToKotlinCommonGenerator(
         }
     }
 
+    private fun FieldType.Message.internalCtor() =
+        dec.value.internalClassFullName() + "()"
+
     private fun MessageDeclaration.internalClassFullName(): String {
         return name.safeFullName(MSG_INTERNAL_SUFFIX)
     }

protoc-gen/src/main/kotlin/kotlinx/rpc/protobuf/model/model.kt

@@ -75,10 +75,12 @@ data class FieldDeclaration(
 ) {
     val packedFixedSize = type.wireType == WireType.FIXED64 || type.wireType == WireType.FIXED32
 
+    val isPartOfOneof: Boolean = dec.realContainingOneof != null
+
     // aligns with edition settings and backward compatibility with proto2 and proto3
     val nullable: Boolean = (dec.hasPresence() && !dec.isRequired && !dec.hasDefaultValue()
             && !dec.isRepeated // repeated fields cannot be nullable (just empty)
-            && dec.realContainingOneof == null // upper conditions would match oneof inner fields
+            && !isPartOfOneof // upper conditions would match oneof inner fields
             && type !is FieldType.Message // messages must not be null (to conform protobuf standards)
             )
             || type is FieldType.OneOf // all OneOf fields are nullable