protoc-plugin: Add Map types

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

Author: Jozott00

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

@@ -7,12 +7,8 @@ package kotlinx.rpc.protobuf
 import org.slf4j.Logger
 import org.slf4j.helpers.NOPLogger
 
-data class CodeGenerationParameters(
-    val messageMode: RpcProtobufPlugin.MessageMode,
-)
 
 open class CodeGenerator(
-    val parameters: CodeGenerationParameters,
     private val indent: String,
     private val builder: StringBuilder = StringBuilder(),
     private val logger: Logger = NOPLogger.NOP_LOGGER,
@@ -68,7 +64,7 @@ open class CodeGenerator(
     }
 
     private fun withNextIndent(block: CodeGenerator.() -> Unit) {
-        CodeGenerator(parameters, "$indent$ONE_INDENT", builder, logger).block()
+        CodeGenerator("$indent$ONE_INDENT", builder, logger).block()
     }
 
     internal fun scope(
@@ -132,7 +128,7 @@ open class CodeGenerator(
             return
         }
 
-        val nested = CodeGenerator(parameters, "$indent$ONE_INDENT", logger = logger).apply(block)
+        val nested = CodeGenerator("$indent$ONE_INDENT", logger = logger).apply(block)
 
         if (nested.isEmpty) {
             newLine()
@@ -306,13 +302,12 @@ open class CodeGenerator(
 }
 
 class FileGenerator(
-    codeGenerationParameters: CodeGenerationParameters,
     var filename: String? = null,
     var packageName: String? = null,
     var packagePath: String? = packageName,
     var fileOptIns: List<String> = emptyList(),
     logger: Logger = NOPLogger.NOP_LOGGER,
-) : CodeGenerator(codeGenerationParameters, "", logger = logger) {
+) : CodeGenerator("", logger = logger) {
     private val imports = mutableListOf<String>()
 
     fun importPackage(name: String) {
@@ -354,10 +349,9 @@ class FileGenerator(
 }
 
 fun file(
-    codeGenerationParameters: CodeGenerationParameters,
     name: String? = null,
     packageName: String? = null,
     logger: Logger = NOPLogger.NOP_LOGGER,
     block: FileGenerator.() -> Unit,
-): FileGenerator = FileGenerator(codeGenerationParameters, name, packageName, packageName, emptyList(), logger)
+): FileGenerator = FileGenerator(name, packageName, packageName, emptyList(), logger)
     .apply(block)

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

@@ -15,7 +15,6 @@ private const val RPC_INTERNAL_PACKAGE_SUFFIX = "_rpc_internal"
 class ModelToKotlinCommonGenerator(
     private val model: Model,
     private val logger: Logger,
-    private val codeGenerationParameters: CodeGenerationParameters,
 ) {
     fun generateKotlinFiles(): List<FileGenerator> {
         return model.files.flatMap { it.generateKotlinFiles() }
@@ -36,7 +35,7 @@ class ModelToKotlinCommonGenerator(
     private fun FileDeclaration.generatePublicKotlinFile(): FileGenerator {
         currentPackage = packageName
 
-        return file(codeGenerationParameters, logger = logger) {
+        return file(logger = logger) {
             filename = this@generatePublicKotlinFile.name
             packageName = this@generatePublicKotlinFile.packageName.safeFullName()
             packagePath = this@generatePublicKotlinFile.packageName.safeFullName()
@@ -61,7 +60,7 @@ class ModelToKotlinCommonGenerator(
     private fun FileDeclaration.generateInternalKotlinFile(): FileGenerator {
         currentPackage = packageName
 
-        return file(codeGenerationParameters, logger = logger) {
+        return file(logger = logger) {
             filename = this@generateInternalKotlinFile.name
             packageName = this@generateInternalKotlinFile.packageName.safeFullName()
             packagePath =

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

@@ -20,27 +20,9 @@ import java.io.File
 class RpcProtobufPlugin {
     companion object {
         private const val DEBUG_OUTPUT_OPTION = "debugOutput"
-        private const val MESSAGE_MODE_OPTION = "messageMode"
-    }
-
-    enum class MessageMode {
-        Interface, Class;
-
-        companion object {
-            fun of(value: String?): MessageMode {
-                return when (value) {
-                    "interface" -> Interface
-                    "class" -> Class
-                    null -> error("Message mode is not specified, use --messageMode=interface or --messageMode=class")
-                    else -> error("Unknown message mode: $value")
-                }
-            }
-        }
     }
 
     private var debugOutput: String? = null
-    private lateinit var messageGenerationMode: MessageMode
-    private var targetCommon: Boolean = false
     private val logger: Logger by lazy {
         val debugOutput = debugOutput ?: return@lazy NOPLogger.NOP_LOGGER
 
@@ -71,7 +53,6 @@ class RpcProtobufPlugin {
         }
 
         debugOutput = parameters[DEBUG_OUTPUT_OPTION]
-        messageGenerationMode = MessageMode.of(parameters[MESSAGE_MODE_OPTION])
 
         val files = input.generateKotlinCommonFiles()
             .map { file ->
@@ -101,9 +82,8 @@ class RpcProtobufPlugin {
     }
 
     private fun CodeGeneratorRequest.generateKotlinCommonFiles(): List<FileGenerator> {
-        val model = this.toCommonModel()
-        val fileGenerator =
-            ModelToKotlinCommonGenerator(model, logger, CodeGenerationParameters(messageGenerationMode))
+        val model = this.toModel()
+        val fileGenerator = ModelToKotlinCommonGenerator(model, logger)
         return fileGenerator.generateKotlinFiles()
     }
 }

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

@@ -11,18 +11,33 @@ import kotlinx.rpc.protobuf.model.*
 
 private val modelCache = mutableMapOf<Descriptors.GenericDescriptor, Any>()
 
-fun CodeGeneratorRequest.toCommonModel(): Model {
+/**
+ * Converts a [CodeGeneratorRequest] into the protoc plugin [Model] of the protobuf.
+ *
+ * @return a [Model] instance containing a list of [FileDeclaration]s that represent
+ *         the converted protobuf files.
+ */
+fun CodeGeneratorRequest.toModel(): Model {
     val protoFileMap = protoFileList.associateBy { it.name }
     val fileDescriptors = mutableMapOf<String, Descriptors.FileDescriptor>()
 
     val files = protoFileList.map { protoFile -> protoFile.toDescriptor(protoFileMap, fileDescriptors) }
 
     return Model(
-        files = files.map { it.toCommonModel() }
+        files = files.map { it.toModel() }
     )
 }
 
 
+/**
+ * Converts a [DescriptorProtos.FileDescriptorProto] instance to a [Descriptors.FileDescriptor],
+ * resolving its dependencies.
+ *
+ * @param protoFileMap a map of file names to `FileDescriptorProto` objects which are available for resolution.
+ * @param cache a mutable map that stores already resolved `FileDescriptor` objects by file name to prevent
+ *  redundant computations.
+ * @return the resolved `FileDescriptor` instance corresponding to this `FileDescriptorProto`.
+ */
 private fun DescriptorProtos.FileDescriptorProto.toDescriptor(
     protoFileMap: Map<String, DescriptorProtos.FileDescriptorProto>,
     cache: MutableMap<String, Descriptors.FileDescriptor>
@@ -39,6 +54,38 @@ private fun DescriptorProtos.FileDescriptorProto.toDescriptor(
     return fileDescriptor
 }
 
+/**
+ * Returns the fully qualified name [FqName] of this descriptor, resolving it to the most specific
+ * declaration or package name based on its type.
+ *
+ * Depending on the type of the descriptor, the fully qualified name is computed recursively,
+ * using the containing type or file, and appropriately converting names.
+ *
+ * @return The fully qualified name represented as an instance of FqName, specific to the descriptor's context.
+ */
+private fun Descriptors.GenericDescriptor.fqName(): FqName {
+    val nameCapital = name.simpleProtoNameToKotlin(firstLetterUpper = true)
+    val nameLower = name.simpleProtoNameToKotlin()
+    return when (this) {
+        is Descriptors.FileDescriptor -> FqName.Package.fromString(`package`)
+        is Descriptors.Descriptor -> FqName.Declaration(nameCapital, containingType?.fqName() ?: file.fqName())
+        is Descriptors.FieldDescriptor -> {
+            val usedName = if (realContainingOneof != null) nameCapital else nameLower
+            FqName.Declaration(usedName, containingType?.fqName() ?: file.fqName())
+        }
+
+        is Descriptors.EnumValueDescriptor -> FqName.Declaration(name, type.fqName())
+        is Descriptors.OneofDescriptor -> FqName.Declaration(nameCapital, containingType?.fqName() ?: file.fqName())
+        is Descriptors.ServiceDescriptor -> FqName.Declaration(nameCapital, file?.fqName() ?: file.fqName())
+        is Descriptors.MethodDescriptor -> FqName.Declaration(nameLower, service?.fqName() ?: file.fqName())
+        else -> error("Unknown generic descriptor: $this")
+    }
+}
+
+/**
+ * Caches the `descriptor.toModel()` result in the [modelCache] to ensure that only a single object
+ * per descriptor exists.
+ */
 private inline fun <D, reified T> D.cached(block: (D) -> T): T
         where D : Descriptors.GenericDescriptor, T : Any {
     if (modelCache.containsKey(this)) {
@@ -49,35 +96,35 @@ private inline fun <D, reified T> D.cached(block: (D) -> T): T
     return declaration
 }
 
-private fun Descriptors.FileDescriptor.toCommonModel(): FileDeclaration = cached {
+private fun Descriptors.FileDescriptor.toModel(): FileDeclaration = cached {
     return FileDeclaration(
         name = kotlinFileName(),
         packageName = FqName.Package.fromString(`package`),
-        dependencies = dependencies.map { it.toCommonModel() },
-        messageDeclarations = messageTypes.map { it.toCommonModel() },
-        enumDeclarations = enumTypes.map { it.toCommonModel() },
-        serviceDeclarations = services.map { it.toCommonModel() },
+        dependencies = dependencies.map { it.toModel() },
+        messageDeclarations = messageTypes.map { it.toModel() },
+        enumDeclarations = enumTypes.map { it.toModel() },
+        serviceDeclarations = services.map { it.toModel() },
         doc = null,
         dec = this,
     )
 }
 
-private fun Descriptors.Descriptor.toCommonModel(): MessageDeclaration = cached {
-    val regularFields = fields.filter { field -> field.realContainingOneof == null }.map { it.toCommonModel() }
+private fun Descriptors.Descriptor.toModel(): MessageDeclaration = cached {
+    val regularFields = fields.filter { field -> field.realContainingOneof == null }.map { it.toModel() }
 
     return MessageDeclaration(
         name = fqName(),
         actualFields = regularFields,
         // get all oneof declarations that are not created from an optional in proto3 https://github.com/googleapis/api-linter/issues/1323
-        oneOfDeclarations = oneofs.filter { it.fields[0].realContainingOneof != null }.map { it.toCommonModel() },
-        enumDeclarations = enumTypes.map { it.toCommonModel() },
-        nestedDeclarations = nestedTypes.map { it.toCommonModel() },
+        oneOfDeclarations = oneofs.filter { it.fields[0].realContainingOneof != null }.map { it.toModel() },
+        enumDeclarations = enumTypes.map { it.toModel() },
+        nestedDeclarations = nestedTypes.map { it.toModel() },
         doc = null,
         dec = this,
     )
 }
 
-private fun Descriptors.FieldDescriptor.toCommonModel(): FieldDeclaration = cached {
+private fun Descriptors.FieldDescriptor.toModel(): FieldDeclaration = cached {
     toProto().hasProto3Optional()
     return FieldDeclaration(
         name = fqName().simpleName,
@@ -88,15 +135,15 @@ private fun Descriptors.FieldDescriptor.toCommonModel(): FieldDeclaration = cach
 }
 
 
-private fun Descriptors.OneofDescriptor.toCommonModel(): OneOfDeclaration = cached {
+private fun Descriptors.OneofDescriptor.toModel(): OneOfDeclaration = cached {
     return OneOfDeclaration(
         name = fqName(),
-        variants = fields.map { it.toCommonModel() },
+        variants = fields.map { it.toModel() },
         dec = this,
     )
 }
 
-private fun Descriptors.EnumDescriptor.toCommonModel(): EnumDeclaration = cached {
+private fun Descriptors.EnumDescriptor.toModel(): EnumDeclaration = cached {
     val entriesMap = mutableMapOf<Int, EnumDeclaration.Entry>()
     val aliases = mutableListOf<EnumDeclaration.Alias>()
 
@@ -105,7 +152,7 @@ private fun Descriptors.EnumDescriptor.toCommonModel(): EnumDeclaration = cached
             val original = entriesMap.getValue(value.number)
             aliases.add(value.toAliasModel(original))
         } else {
-            entriesMap[value.number] = value.toCommonModel()
+            entriesMap[value.number] = value.toModel()
         }
     }
 
@@ -122,7 +169,7 @@ private fun Descriptors.EnumDescriptor.toCommonModel(): EnumDeclaration = cached
     )
 }
 
-private fun Descriptors.EnumValueDescriptor.toCommonModel(): EnumDeclaration.Entry = cached {
+private fun Descriptors.EnumValueDescriptor.toModel(): EnumDeclaration.Entry = cached {
     return EnumDeclaration.Entry(
         name = fqName(),
         doc = null,
@@ -140,19 +187,19 @@ private fun Descriptors.EnumValueDescriptor.toAliasModel(original: EnumDeclarati
     )
 }
 
-private fun Descriptors.ServiceDescriptor.toCommonModel(): ServiceDeclaration = cached {
+private fun Descriptors.ServiceDescriptor.toModel(): ServiceDeclaration = cached {
     return ServiceDeclaration(
         name = fqName(),
-        methods = methods.map { it.toCommonModel() },
+        methods = methods.map { it.toModel() },
         dec = this,
     )
 }
 
-private fun Descriptors.MethodDescriptor.toCommonModel(): MethodDeclaration = cached {
+private fun Descriptors.MethodDescriptor.toModel(): MethodDeclaration = cached {
     return MethodDeclaration(
         name = name,
-        inputType = inputType.toCommonModel(),
-        outputType = outputType.toCommonModel(),
+        inputType = inputType.toModel(),
+        outputType = outputType.toModel(),
         dec = this,
     )
 }
@@ -176,41 +223,25 @@ private fun Descriptors.FieldDescriptor.modelType(): FieldType {
         Descriptors.FieldDescriptor.Type.SFIXED64 -> FieldType.IntegralType.SFIXED64
         Descriptors.FieldDescriptor.Type.SINT32 -> FieldType.IntegralType.SINT32
         Descriptors.FieldDescriptor.Type.SINT64 -> FieldType.IntegralType.SINT64
-        Descriptors.FieldDescriptor.Type.ENUM -> FieldType.Reference(lazy { enumType!!.toCommonModel().name })
-        Descriptors.FieldDescriptor.Type.MESSAGE -> FieldType.Reference(lazy { messageType!!.toCommonModel().name })
+        Descriptors.FieldDescriptor.Type.ENUM -> FieldType.Reference(lazy { enumType!!.toModel().name })
+        Descriptors.FieldDescriptor.Type.MESSAGE -> FieldType.Reference(lazy { messageType!!.toModel().name })
         Descriptors.FieldDescriptor.Type.GROUP -> error("GROUP type is unsupported")
     }
 
+    if (isMapField) {
+        val keyType = messageType.findFieldByName("key").modelType()
+        val valType = messageType.findFieldByName("value").modelType()
+        val mapEntry = FieldType.Map.Entry(keyType, valType)
+        return FieldType.Map(lazy { mapEntry })
+    }
+
     if (isRepeated) {
         return FieldType.List(baseType)
     }
 
-    // TODO: Handle map type
-
     return baseType
 }
 
-//// GenericDescriptor Extensions ////
-
-private fun Descriptors.GenericDescriptor.fqName(): FqName {
-    val nameCapital = name.simpleProtoNameToKotlin(firstLetterUpper = true)
-    val nameLower = name.simpleProtoNameToKotlin()
-    return when (this) {
-        is Descriptors.FileDescriptor -> FqName.Package.fromString(`package`)
-        is Descriptors.Descriptor -> FqName.Declaration(nameCapital, containingType?.fqName() ?: file.fqName())
-        is Descriptors.FieldDescriptor -> {
-            val usedName = if (realContainingOneof != null) nameCapital else nameLower
-            FqName.Declaration(usedName, containingType?.fqName() ?: file.fqName())
-        }
-
-        is Descriptors.EnumValueDescriptor -> FqName.Declaration(name, type.fqName())
-        is Descriptors.OneofDescriptor -> FqName.Declaration(nameCapital, containingType?.fqName() ?: file.fqName())
-        is Descriptors.ServiceDescriptor -> FqName.Declaration(nameCapital, file?.fqName() ?: file.fqName())
-        is Descriptors.MethodDescriptor -> FqName.Declaration(nameLower, service?.fqName() ?: file.fqName())
-        else -> error("Unknown generic descriptor: $this")
-    }
-}
-
 //// Utility Extensions ////
 
 private fun Descriptors.FileDescriptor.kotlinFileName(): String {