Skip to content

Commit e3020cd

Browse files
Mr3zeeMr3zee
committed
KRPC-145 Map types in gRPC (#334)
1 parent 9bb284b commit e3020cd

File tree

7 files changed

+107
-22
lines changed

7 files changed

+107
-22
lines changed

‎protobuf-plugin/build.gradle.kts

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,6 @@ sourceSets {
3838
"**/empty_deprecated.proto",
3939
"**/example.proto",
4040
"**/funny_types.proto",
41-
"**/map.proto",
4241
"**/multiple_files.proto",
4342
"**/one_of.proto",
4443
"**/options.proto",

‎protobuf-plugin/src/main/kotlin/kotlinx/rpc/protobuf/ModelToKotlinGenerator.kt

Lines changed: 44 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -204,6 +204,7 @@ class ModelToKotlinGenerator(
204204
val uppercaseName = field.name.replaceFirstChar { ch -> ch.uppercase() }
205205
val setFieldCall = when (field.type) {
206206
is FieldType.List -> "addAll$uppercaseName"
207+
is FieldType.Map -> "putAll$uppercaseName"
207208
else -> "set$uppercaseName"
208209
}
209210

@@ -225,6 +226,7 @@ class ModelToKotlinGenerator(
225226
declaration.actualFields.forEach { field ->
226227
val javaName = when (field.type) {
227228
is FieldType.List -> "${field.name}List"
229+
is FieldType.Map -> "${field.name}Map"
228230
else -> field.name
229231
}
230232

@@ -270,6 +272,19 @@ class ModelToKotlinGenerator(
270272
}
271273
}
272274

275+
is FieldType.Map -> {
276+
val entry by entry
277+
278+
when (val value = entry.value) {
279+
is FieldType.Reference -> ".mapValues { it.value.toPlatform() }".also {
280+
val fq by value.value
281+
importRootDeclarationIfNeeded(fq, "toPlatform", true)
282+
}
283+
is FieldType.IntegralType -> ""
284+
else -> error("Unsupported type: $value")
285+
}
286+
}
287+
273288
else -> ""
274289
}
275290
}
@@ -296,6 +311,19 @@ class ModelToKotlinGenerator(
296311
}
297312
}
298313

314+
is FieldType.Map -> {
315+
val entry by entry
316+
317+
when (val value = entry.value) {
318+
is FieldType.Reference -> ".mapValues { it.value.toKotlin() }".also {
319+
val fq by value.value
320+
importRootDeclarationIfNeeded(fq, "toKotlin", true)
321+
}
322+
is FieldType.IntegralType -> ""
323+
else -> error("Unsupported type: $value")
324+
}
325+
}
326+
299327
else -> ""
300328
}
301329
}
@@ -306,7 +334,6 @@ class ModelToKotlinGenerator(
306334

307335
private fun FieldDeclaration.typeFqName(): String {
308336
return when (type) {
309-
// KRPC-156 Reference Types
310337
is FieldType.Reference -> {
311338
val value by type.value
312339
value.safeFullName()
@@ -326,12 +353,22 @@ class ModelToKotlinGenerator(
326353
"List<${fqValue.safeFullName()}>"
327354
}
328355

329-
// KRPC-145 Map Types
330-
// is FieldType.Map -> {
331-
// "Map<${type.keyName.simpleName}, ${type.valueName.simpleName}>"
332-
// }
333-
else -> {
334-
error("Unsupported type: $this")
356+
is FieldType.Map -> {
357+
val entry by type.entry
358+
359+
val fqKey = when (val key = entry.key) {
360+
is FieldType.Reference -> key.value.value
361+
is FieldType.IntegralType -> key.fqName
362+
else -> error("Unsupported type: $key")
363+
}
364+
365+
val fqValue = when (val value = entry.value) {
366+
is FieldType.Reference -> value.value.value
367+
is FieldType.IntegralType -> value.fqName
368+
else -> error("Unsupported type: $value")
369+
}
370+
371+
"Map<${fqKey.safeFullName()}, ${fqValue.safeFullName()}>"
335372
}
336373
}.withNullability(nullable)
337374
}

‎protobuf-plugin/src/main/kotlin/kotlinx/rpc/protobuf/ProtoToModelInterpreter.kt

Lines changed: 34 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ class ProtoToModelInterpreter(
2727
}
2828

2929
// package name of a currently parsed file
30-
private var packageName: FqName.Package by Delegates.notNull<FqName.Package>()
30+
private var packageName: FqName.Package by Delegates.notNull()
3131

3232
private fun DescriptorProtos.FileDescriptorProto.toModel(): FileDeclaration {
3333
val dependencies = dependencyList.map { depFilename ->
@@ -47,7 +47,7 @@ class ProtoToModelInterpreter(
4747
name = kotlinFileName(),
4848
packageName = packageName,
4949
dependencies = dependencies,
50-
messageDeclarations = messageTypeList.map { it.toModel(resolver, outerClass, parent = null) },
50+
messageDeclarations = messageTypeList.mapNotNull { it.toModel(resolver, outerClass, parent = null) },
5151
enumDeclarations = enumTypeList.map { it.toModel(resolver, outerClass, packageName) },
5252
serviceDeclarations = serviceList.map { it.toModel(resolver) },
5353
deprecated = options.deprecated,
@@ -88,15 +88,24 @@ class ProtoToModelInterpreter(
8888
return originalPackage
8989
}
9090

91+
private val mapEntries = mutableMapOf<FqName, Lazy<FieldType.Map.Entry>>()
92+
9193
private fun DescriptorProtos.DescriptorProto.toModel(
9294
parentResolver: NameResolver,
9395
outerClass: FqName,
9496
parent: FqName?,
95-
): MessageDeclaration {
97+
): MessageDeclaration? {
9698
val simpleName = name.fullProtoNameToKotlin(firstLetterUpper = true)
9799
val fqName = parentResolver.declarationFqName(simpleName, parent ?: packageName)
98100
val resolver = parentResolver.withScope(fqName)
99101

102+
if (options.mapEntry) {
103+
mapEntries[fqName] = resolveMapEntry(resolver)
104+
return null
105+
}
106+
107+
// resolve before fields
108+
val nestedDeclarations = nestedTypeList.mapNotNull { it.toModel(resolver, outerClass, fqName) }
100109
val fields = fieldList.mapNotNull {
101110
val oneOfName = if (it.hasOneofIndex()) {
102111
oneofDeclList[it.oneofIndex].name
@@ -113,14 +122,21 @@ class ProtoToModelInterpreter(
113122
actualFields = fields,
114123
oneOfDeclarations = oneofDeclList.mapIndexedNotNull { i, desc -> desc.toModel(i, resolver) },
115124
enumDeclarations = enumTypeList.map { it.toModel(resolver, outerClass, fqName) },
116-
nestedDeclarations = nestedTypeList.map { it.toModel(resolver, outerClass, fqName) },
125+
nestedDeclarations = nestedDeclarations,
117126
deprecated = options.deprecated,
118127
doc = null,
119128
).apply {
120129
messages[name] = this
121130
}
122131
}
123132

133+
private fun DescriptorProtos.DescriptorProto.resolveMapEntry(resolver: NameResolver): Lazy<FieldType.Map.Entry> = lazy {
134+
val keyType = fieldList[0].toModel(null, resolver) ?: error("Key type is null")
135+
val valueType = fieldList[1].toModel(null, resolver) ?: error("Value type is null")
136+
137+
FieldType.Map.Entry(keyType.type, valueType.type)
138+
}
139+
124140
private val oneOfFieldMembers = mutableMapOf<Int, MutableList<DescriptorProtos.FieldDescriptorProto>>()
125141

126142
private fun DescriptorProtos.FieldDescriptorProto.toModel(
@@ -176,19 +192,15 @@ class ProtoToModelInterpreter(
176192
typeName
177193
.substringAfter('.')
178194
.fullProtoNameToKotlin(firstLetterUpper = true)
179-
// TODO KRPC-146 Nested Types: parent full type resolution
180195
// KRPC-144 Import types
181-
.asReference { resolver.resolve(it) }
182-
.let { wrapWithLabel(it) }
196+
.let { resolvedType(it, resolver) }
183197
}
184198

185199
hasTypeName() -> {
186200
typeName
187201
.fullProtoNameToKotlin(firstLetterUpper = true)
188-
// TODO KRPC-146 Nested Types: parent full type resolution
189202
// KRPC-144 Import types: we assume local types now
190-
.asReference { resolver.resolve(it) }
191-
.let { wrapWithLabel(it) }
203+
.let { resolvedType(it, resolver) }
192204
}
193205

194206
else -> {
@@ -221,6 +233,18 @@ class ProtoToModelInterpreter(
221233
}
222234
}
223235

236+
private fun DescriptorProtos.FieldDescriptorProto.resolvedType(name: String, resolver: NameResolver): FieldType {
237+
val entry = mapEntries[resolver.resolveOrNull(name)]
238+
239+
if (entry != null) {
240+
return FieldType.Map(entry)
241+
}
242+
243+
val fieldType = FieldType.Reference(lazy { resolver.resolve(name) })
244+
245+
return wrapWithLabel(fieldType)
246+
}
247+
224248
private fun String.asReference(resolver: (String) -> FqName) =
225249
FieldType.Reference(lazy { resolver(this) })
226250

‎protobuf-plugin/src/main/kotlin/kotlinx/rpc/protobuf/model/FieldDeclaration.kt

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,10 @@ sealed interface FieldType {
1919
override val defaultValue: String = "emptyList()"
2020
}
2121

22-
data class Map(val keyName: FqName, val valueName: FqName) : FieldType {
22+
data class Map(val entry: Lazy<Entry>) : FieldType {
2323
override val defaultValue: String = "emptyMap()"
24+
25+
data class Entry(val key: FieldType, val value: FieldType)
2426
}
2527

2628
data class Reference(val value: Lazy<FqName>) : FieldType {

‎protobuf-plugin/src/test/kotlin/kotlinx/rpc/protobuf/test/TestReferenceService.kt

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,10 @@ class ReferenceTestServiceImpl : ReferenceTestService {
4141
override suspend fun Nested(message: Nested): Nested {
4242
return message
4343
}
44+
45+
override suspend fun Map(message: TestMap): TestMap {
46+
return message
47+
}
4448
}
4549

4650
class TestReferenceService : GrpcServerTest() {
@@ -197,4 +201,20 @@ class TestReferenceService : GrpcServerTest() {
197201
assertEquals("42", result.string)
198202
assertEquals(Nested.Inner2.NestedEnum.ZERO, result.enum)
199203
}
204+
205+
@Test
206+
fun testMap() = runGrpcTest { grpcClient ->
207+
val service = grpcClient.withService<ReferenceTestService>()
208+
val result = service.Map(TestMap {
209+
primitives = mapOf("1" to 2, "2" to 1)
210+
references = mapOf("ref" to kotlinx.rpc.protobuf.test.References {
211+
other = kotlinx.rpc.protobuf.test.Other {
212+
field = 42
213+
}
214+
})
215+
})
216+
217+
assertEquals(mapOf("1" to 2L, "2" to 1L), result.primitives)
218+
assertEquals(mapOf("ref" to 42), result.references.mapValues { it.value.other.field })
219+
}
200220
}

‎protobuf-plugin/src/test/proto/reference_service.proto

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import "enum.proto";
66
import "optional.proto";
77
import "repeated.proto";
88
import "nested.proto";
9+
import "test_map.proto";
910

1011
service ReferenceTestService {
1112
rpc Get(References) returns (kotlinx.rpc.protobuf.test.References);
@@ -17,4 +18,6 @@ service ReferenceTestService {
1718
rpc Repeated(kotlinx.rpc.protobuf.test.Repeated) returns (kotlinx.rpc.protobuf.test.Repeated);
1819

1920
rpc Nested(kotlinx.rpc.protobuf.test.Nested) returns (kotlinx.rpc.protobuf.test.Nested);
21+
22+
rpc Map(kotlinx.rpc.protobuf.test.TestMap) returns (kotlinx.rpc.protobuf.test.TestMap);
2023
}

‎protobuf-plugin/src/test/proto/map.proto renamed to ‎protobuf-plugin/src/test/proto/test_map.proto

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,9 @@ syntax = "proto3";
22

33
package kotlinx.rpc.protobuf.test;
44

5-
import "one_of.proto";
5+
import "reference_package.proto";
66

7-
message Map {
7+
message TestMap {
88
map<string, int64> primitives = 1;
9-
map<string, OneOf> references = 2;
9+
map<string, References> references = 2;
1010
}

0 commit comments

Comments
 (0)