Kotlin · Mr3zee · Apr 30, 2025 · Apr 29, 2025
diff --git a/protobuf-plugin/build.gradle.kts b/protobuf-plugin/build.gradle.kts
@@ -38,7 +38,6 @@ sourceSets {
                 "**/empty_deprecated.proto",
                 "**/example.proto",
                 "**/funny_types.proto",
-                "**/map.proto",
                 "**/multiple_files.proto",
                 "**/one_of.proto",
                 "**/options.proto",

diff --git a/protobuf-plugin/src/main/kotlin/kotlinx/rpc/protobuf/ModelToKotlinGenerator.kt b/protobuf-plugin/src/main/kotlin/kotlinx/rpc/protobuf/ModelToKotlinGenerator.kt
@@ -204,6 +204,7 @@ class ModelToKotlinGenerator(
                     val uppercaseName = field.name.replaceFirstChar { ch -> ch.uppercase() }
                     val setFieldCall = when (field.type) {
                         is FieldType.List -> "addAll$uppercaseName"
+                        is FieldType.Map -> "putAll$uppercaseName"
                         else -> "set$uppercaseName"
                     }
 
@@ -225,6 +226,7 @@ class ModelToKotlinGenerator(
                 declaration.actualFields.forEach { field ->
                     val javaName = when (field.type) {
                         is FieldType.List -> "${field.name}List"
+                        is FieldType.Map -> "${field.name}Map"
                         else -> field.name
                     }
 
@@ -270,6 +272,19 @@ class ModelToKotlinGenerator(
                 }
             }
 
+            is FieldType.Map -> {
+                val entry by entry
+
+                when (val value = entry.value) {
+                    is FieldType.Reference -> ".mapValues { it.value.toPlatform() }".also {
+                        val fq by value.value
+                        importRootDeclarationIfNeeded(fq, "toPlatform", true)
+                    }
+                    is FieldType.IntegralType -> ""
+                    else -> error("Unsupported type: $value")
+                }
+            }
+
             else -> ""
         }
     }
@@ -296,6 +311,19 @@ class ModelToKotlinGenerator(
                 }
             }
 
+            is FieldType.Map -> {
+                val entry by entry
+
+                when (val value = entry.value) {
+                    is FieldType.Reference -> ".mapValues { it.value.toKotlin() }".also {
+                        val fq by value.value
+                        importRootDeclarationIfNeeded(fq, "toKotlin", true)
+                    }
+                    is FieldType.IntegralType -> ""
+                    else -> error("Unsupported type: $value")
+                }
+            }
+
             else -> ""
         }
     }
@@ -306,7 +334,6 @@ class ModelToKotlinGenerator(
 
     private fun FieldDeclaration.typeFqName(): String {
         return when (type) {
-            // KRPC-156 Reference Types
             is FieldType.Reference -> {
                 val value by type.value
                 value.safeFullName()
@@ -326,12 +353,22 @@ class ModelToKotlinGenerator(
                 "List<${fqValue.safeFullName()}>"
             }
 
-            // KRPC-145 Map Types
-//            is FieldType.Map -> {
-//                "Map<${type.keyName.simpleName}, ${type.valueName.simpleName}>"
-//            }
-            else -> {
-                error("Unsupported type: $this")
+            is FieldType.Map -> {
+                val entry by type.entry
+
+                val fqKey = when (val key = entry.key) {
+                    is FieldType.Reference -> key.value.value
+                    is FieldType.IntegralType -> key.fqName
+                    else -> error("Unsupported type: $key")
+                }
+
+                val fqValue = when (val value = entry.value) {
+                    is FieldType.Reference -> value.value.value
+                    is FieldType.IntegralType -> value.fqName
+                    else -> error("Unsupported type: $value")
+                }
+
+                "Map<${fqKey.safeFullName()}, ${fqValue.safeFullName()}>"
             }
         }.withNullability(nullable)
     }

diff --git a/protobuf-plugin/src/main/kotlin/kotlinx/rpc/protobuf/ProtoToModelInterpreter.kt b/protobuf-plugin/src/main/kotlin/kotlinx/rpc/protobuf/ProtoToModelInterpreter.kt
@@ -27,7 +27,7 @@ class ProtoToModelInterpreter(
     }
 
     // package name of a currently parsed file
-    private var packageName: FqName.Package by Delegates.notNull<FqName.Package>()
+    private var packageName: FqName.Package by Delegates.notNull()
 
     private fun DescriptorProtos.FileDescriptorProto.toModel(): FileDeclaration {
         val dependencies = dependencyList.map { depFilename ->
@@ -47,7 +47,7 @@ class ProtoToModelInterpreter(
             name = kotlinFileName(),
             packageName = packageName,
             dependencies = dependencies,
-            messageDeclarations = messageTypeList.map { it.toModel(resolver, outerClass, parent = null) },
+            messageDeclarations = messageTypeList.mapNotNull { it.toModel(resolver, outerClass, parent = null) },
             enumDeclarations = enumTypeList.map { it.toModel(resolver, outerClass, packageName) },
             serviceDeclarations = serviceList.map { it.toModel(resolver) },
             deprecated = options.deprecated,
@@ -88,15 +88,24 @@ class ProtoToModelInterpreter(
         return originalPackage
     }
 
+    private val mapEntries = mutableMapOf<FqName, Lazy<FieldType.Map.Entry>>()
+
     private fun DescriptorProtos.DescriptorProto.toModel(
         parentResolver: NameResolver,
         outerClass: FqName,
         parent: FqName?,
-    ): MessageDeclaration {
+    ): MessageDeclaration? {
         val simpleName = name.fullProtoNameToKotlin(firstLetterUpper = true)
         val fqName = parentResolver.declarationFqName(simpleName, parent ?: packageName)
         val resolver = parentResolver.withScope(fqName)
 
+        if (options.mapEntry) {
+            mapEntries[fqName] = resolveMapEntry(resolver)
+            return null
+        }
+
+        // resolve before fields
+        val nestedDeclarations = nestedTypeList.mapNotNull { it.toModel(resolver, outerClass, fqName) }
         val fields = fieldList.mapNotNull {
             val oneOfName = if (it.hasOneofIndex()) {
                 oneofDeclList[it.oneofIndex].name
@@ -113,14 +122,21 @@ class ProtoToModelInterpreter(
             actualFields = fields,
             oneOfDeclarations = oneofDeclList.mapIndexedNotNull { i, desc -> desc.toModel(i, resolver) },
             enumDeclarations = enumTypeList.map { it.toModel(resolver, outerClass, fqName) },
-            nestedDeclarations = nestedTypeList.map { it.toModel(resolver, outerClass, fqName) },
+            nestedDeclarations = nestedDeclarations,
             deprecated = options.deprecated,
             doc = null,
         ).apply {
             messages[name] = this
         }
     }
 
+    private fun DescriptorProtos.DescriptorProto.resolveMapEntry(resolver: NameResolver): Lazy<FieldType.Map.Entry> = lazy {
+        val keyType = fieldList[0].toModel(null, resolver) ?: error("Key type is null")
+        val valueType = fieldList[1].toModel(null, resolver) ?: error("Value type is null")
+
+        FieldType.Map.Entry(keyType.type, valueType.type)
+    }
+
     private val oneOfFieldMembers = mutableMapOf<Int, MutableList<DescriptorProtos.FieldDescriptorProto>>()
 
     private fun DescriptorProtos.FieldDescriptorProto.toModel(
@@ -176,19 +192,15 @@ class ProtoToModelInterpreter(
                 typeName
                     .substringAfter('.')
                     .fullProtoNameToKotlin(firstLetterUpper = true)
-                    // TODO KRPC-146 Nested Types: parent full type resolution
                     //  KRPC-144 Import types
-                    .asReference { resolver.resolve(it) }
-                    .let { wrapWithLabel(it) }
+                    .let { resolvedType(it, resolver) }
             }
 
             hasTypeName() -> {
                 typeName
                     .fullProtoNameToKotlin(firstLetterUpper = true)
-                    // TODO KRPC-146 Nested Types: parent full type resolution
                     //  KRPC-144 Import types: we assume local types now
-                    .asReference { resolver.resolve(it) }
-                    .let { wrapWithLabel(it) }
+                    .let { resolvedType(it, resolver) }
             }
 
             else -> {
@@ -221,6 +233,18 @@ class ProtoToModelInterpreter(
         }
     }
 
+    private fun DescriptorProtos.FieldDescriptorProto.resolvedType(name: String, resolver: NameResolver): FieldType {
+        val entry = mapEntries[resolver.resolveOrNull(name)]
+
+        if (entry != null) {
+            return FieldType.Map(entry)
+        }
+
+        val fieldType = FieldType.Reference(lazy { resolver.resolve(name) })
+
+        return wrapWithLabel(fieldType)
+    }
+
     private fun String.asReference(resolver: (String) -> FqName) =
         FieldType.Reference(lazy { resolver(this) })
 

diff --git a/protobuf-plugin/src/main/kotlin/kotlinx/rpc/protobuf/model/FieldDeclaration.kt b/protobuf-plugin/src/main/kotlin/kotlinx/rpc/protobuf/model/FieldDeclaration.kt
@@ -19,8 +19,10 @@ sealed interface FieldType {
         override val defaultValue: String = "emptyList()"
     }
 
-    data class Map(val keyName: FqName, val valueName: FqName) : FieldType {
+    data class Map(val entry: Lazy<Entry>) : FieldType {
         override val defaultValue: String = "emptyMap()"
+
+        data class Entry(val key: FieldType, val value: FieldType)
     }
 
     data class Reference(val value: Lazy<FqName>) : FieldType {

diff --git a/protobuf-plugin/src/test/kotlin/kotlinx/rpc/protobuf/test/TestReferenceService.kt b/protobuf-plugin/src/test/kotlin/kotlinx/rpc/protobuf/test/TestReferenceService.kt
@@ -41,6 +41,10 @@ class ReferenceTestServiceImpl : ReferenceTestService {
     override suspend fun Nested(message: Nested): Nested {
         return message
     }
+
+    override suspend fun Map(message: TestMap): TestMap {
+        return message
+    }
 }
 
 class TestReferenceService : GrpcServerTest() {
@@ -197,4 +201,20 @@ class TestReferenceService : GrpcServerTest() {
         assertEquals("42", result.string)
         assertEquals(Nested.Inner2.NestedEnum.ZERO, result.enum)
     }
+
+    @Test
+    fun testMap() = runGrpcTest { grpcClient ->
+        val service = grpcClient.withService<ReferenceTestService>()
+        val result = service.Map(TestMap {
+            primitives = mapOf("1" to 2, "2" to 1)
+            references = mapOf("ref" to kotlinx.rpc.protobuf.test.References {
+                other = kotlinx.rpc.protobuf.test.Other {
+                    field = 42
+                }
+            })
+        })
+
+        assertEquals(mapOf("1" to 2L, "2" to 1L), result.primitives)
+        assertEquals(mapOf("ref" to 42), result.references.mapValues { it.value.other.field })
+    }
 }
diff --git a/protobuf-plugin/src/test/proto/reference_service.proto b/protobuf-plugin/src/test/proto/reference_service.proto
@@ -6,6 +6,7 @@ import "enum.proto";
 import "optional.proto";
 import "repeated.proto";
 import "nested.proto";
+import "test_map.proto";
 
 service ReferenceTestService {
   rpc Get(References) returns (kotlinx.rpc.protobuf.test.References);
@@ -17,4 +18,6 @@ service ReferenceTestService {
   rpc Repeated(kotlinx.rpc.protobuf.test.Repeated) returns (kotlinx.rpc.protobuf.test.Repeated);
 
   rpc Nested(kotlinx.rpc.protobuf.test.Nested) returns (kotlinx.rpc.protobuf.test.Nested);
+
+  rpc Map(kotlinx.rpc.protobuf.test.TestMap) returns (kotlinx.rpc.protobuf.test.TestMap);
 }
diff --git a/protobuf-plugin/src/test/proto/map.proto → ...obuf-plugin/src/test/proto/test_map.proto b/protobuf-plugin/src/test/proto/map.proto → ...obuf-plugin/src/test/proto/test_map.proto
@@ -2,9 +2,9 @@ syntax = "proto3";
 
 package kotlinx.rpc.protobuf.test;
 
-import "one_of.proto";
+import "reference_package.proto";
 
-message Map {
+message TestMap {
   map<string, int64> primitives = 1;
-  map<string, OneOf> references = 2;
+  map<string, References> references = 2;
 }