Skip to content

Commit 5fc6fd0

Browse files
Mr3zeeMr3zee
committed
Compiler plugin tinkering for gRPC
1 parent 32bbabb commit 5fc6fd0

File tree

13 files changed

+152
-37
lines changed

13 files changed

+152
-37
lines changed

‎compiler-plugin/compiler-plugin-backend/src/main/core/kotlinx/rpc/codegen/extension/RpcIrContext.kt‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023-2024 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
2+
* Copyright 2023-2025 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
33
*/
44

55
package kotlinx.rpc.codegen.extension
@@ -105,6 +105,14 @@ internal class RpcIrContext(
105105
getRpcIrClassSymbol("RpcServiceDescriptor", "descriptor")
106106
}
107107

108+
val grpcServiceDescriptor by lazy {
109+
getIrClassSymbol("kotlinx.rpc.grpc.descriptor", "GrpcServiceDescriptor")
110+
}
111+
112+
val grpcDelegate by lazy {
113+
getIrClassSymbol("kotlinx.rpc.grpc.descriptor", "GrpcDelegate")
114+
}
115+
108116
val rpcType by lazy {
109117
getRpcIrClassSymbol("RpcType", "descriptor")
110118
}
@@ -262,6 +270,10 @@ internal class RpcIrContext(
262270
rpcServiceDescriptor.namedProperty("fqName")
263271
}
264272

273+
val grpcServiceDescriptorDelegate by lazy {
274+
grpcServiceDescriptor.namedProperty("delegate")
275+
}
276+
265277
private fun IrClassSymbol.namedProperty(name: String): IrPropertySymbol {
266278
return owner.properties.single { it.name.asString() == name }.symbol
267279
}
@@ -276,7 +288,7 @@ internal class RpcIrContext(
276288
return getIrClassSymbol("kotlinx.rpc$suffix", name)
277289
}
278290

279-
private fun getIrClassSymbol(packageName: String, name: String): IrClassSymbol {
291+
fun getIrClassSymbol(packageName: String, name: String): IrClassSymbol {
280292
return versionSpecificApi.referenceClass(pluginContext, packageName, name)
281293
?: error("Unable to find symbol. Package: $packageName, name: $name")
282294
}

‎compiler-plugin/compiler-plugin-backend/src/main/core/kotlinx/rpc/codegen/extension/RpcIrServiceProcessor.kt‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023-2024 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
2+
* Copyright 2023-2025 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
33
*/
44

55
package kotlinx.rpc.codegen.extension
@@ -9,14 +9,17 @@ import org.jetbrains.kotlin.cli.common.messages.MessageCollector
99
import org.jetbrains.kotlin.ir.IrStatement
1010
import org.jetbrains.kotlin.ir.declarations.IrClass
1111
import org.jetbrains.kotlin.ir.util.hasAnnotation
12+
import org.jetbrains.kotlin.ir.util.isInterface
1213
import org.jetbrains.kotlin.ir.visitors.IrElementTransformer
1314

1415
internal class RpcIrServiceProcessor(
1516
@Suppress("unused")
1617
private val logger: MessageCollector,
1718
) : IrElementTransformer<RpcIrContext> {
1819
override fun visitClass(declaration: IrClass, data: RpcIrContext): IrStatement {
19-
if (declaration.hasAnnotation(RpcClassId.rpcAnnotation)) {
20+
if ((declaration.hasAnnotation(RpcClassId.rpcAnnotation)
21+
|| declaration.hasAnnotation(RpcClassId.grpcAnnotation)) && declaration.isInterface
22+
) {
2023
processService(declaration, data)
2124
}
2225

‎compiler-plugin/compiler-plugin-backend/src/main/core/kotlinx/rpc/codegen/extension/RpcStubGenerator.kt‎

Lines changed: 66 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023-2024 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
2+
* Copyright 2023-2025 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
33
*/
44

55
package kotlinx.rpc.codegen.extension
@@ -44,6 +44,10 @@ private object Descriptor {
4444
const val CREATE_INSTANCE = "createInstance"
4545
}
4646

47+
private object GrpcDescriptor {
48+
const val DELEGATE = "delegate"
49+
}
50+
4751
@Suppress("detekt.LargeClass", "detekt.TooManyFunctions")
4852
internal class RpcStubGenerator(
4953
private val declaration: ServiceDeclaration,
@@ -123,7 +127,10 @@ internal class RpcStubGenerator(
123127

124128
clientProperty()
125129

126-
coroutineContextProperty()
130+
// not for gRPC
131+
if (!declaration.isGrpc) {
132+
coroutineContextProperty()
133+
}
127134

128135
declaration.fields.forEach {
129136
rpcFlowField(it)
@@ -550,7 +557,15 @@ internal class RpcStubGenerator(
550557
overriddenSymbols = listOf(method.function.symbol)
551558

552559
body = irBuilder(symbol).irBlockBody {
553-
+irReturn(
560+
val call = if (declaration.isGrpc) {
561+
irRpcMethodClientCall(
562+
method = method,
563+
functionThisReceiver = functionThisReceiver,
564+
isMethodObject = isMethodObject,
565+
methodClass = methodClass,
566+
arguments = arguments,
567+
)
568+
} else {
554569
irCall(
555570
callee = ctx.functions.scopedClientCall,
556571
type = method.function.returnType,
@@ -600,7 +615,9 @@ internal class RpcStubGenerator(
600615

601616
putValueArgument(1, lambda)
602617
}
603-
)
618+
}
619+
620+
+irReturn(call)
604621
}
605622
}
606623
}
@@ -868,7 +885,10 @@ internal class RpcStubGenerator(
868885
stubCompanionObjectThisReceiver = thisReceiver
869886
?: error("Stub companion object expected to have thisReceiver: ${name.asString()}")
870887

871-
superTypes = listOf(ctx.rpcServiceDescriptor.typeWith(declaration.serviceType))
888+
superTypes = listOfNotNull(
889+
ctx.rpcServiceDescriptor.typeWith(declaration.serviceType),
890+
if (declaration.isGrpc) ctx.grpcServiceDescriptor.typeWith(declaration.serviceType) else null,
891+
)
872892

873893
generateCompanionObjectConstructor()
874894

@@ -901,6 +921,10 @@ internal class RpcStubGenerator(
901921
generateCreateInstanceFunction()
902922

903923
generateGetFieldsFunction()
924+
925+
if (declaration.isGrpc) {
926+
generateGrpcDelegateProperty()
927+
}
904928
}
905929

906930
/**
@@ -1488,6 +1512,43 @@ internal class RpcStubGenerator(
14881512
}
14891513
}
14901514

1515+
/**
1516+
* override val delegate: GrpcDelegate = MyServiceDelegate
1517+
*/
1518+
private fun IrClass.generateGrpcDelegateProperty() {
1519+
addProperty {
1520+
name = Name.identifier(GrpcDescriptor.DELEGATE)
1521+
visibility = DescriptorVisibilities.PUBLIC
1522+
}.apply {
1523+
overriddenSymbols = listOf(ctx.properties.grpcServiceDescriptorDelegate)
1524+
1525+
addBackingFieldUtil {
1526+
visibility = DescriptorVisibilities.PRIVATE
1527+
type = ctx.grpcDelegate.defaultType
1528+
vsApi { isFinalVS = true }
1529+
}.apply {
1530+
initializer = factory.createExpressionBody(
1531+
IrGetObjectValueImpl(
1532+
startOffset = UNDEFINED_OFFSET,
1533+
endOffset = UNDEFINED_OFFSET,
1534+
type = ctx.grpcDelegate.defaultType,
1535+
symbol = ctx.getIrClassSymbol(
1536+
declaration.service.packageFqName?.asString()
1537+
?: error("Expected package name fro service ${declaration.service.name}"),
1538+
"${declaration.service.name.asString()}Delegate",
1539+
),
1540+
)
1541+
)
1542+
}
1543+
1544+
addDefaultGetter(this@generateGrpcDelegateProperty, ctx.irBuiltIns) {
1545+
visibility = DescriptorVisibilities.PUBLIC
1546+
overriddenSymbols = listOf(ctx.properties.grpcServiceDescriptorDelegate.owner.getterOrFail.symbol)
1547+
}
1548+
}
1549+
}
1550+
1551+
14911552
// Associated object annotation works on JS, WASM, and Native platforms.
14921553
// See https://kotlinlang.org/api/latest/jvm/stdlib/kotlin.reflect/find-associated-object.html
14931554
private fun addAssociatedObjectAnnotationIfPossible() {

‎compiler-plugin/compiler-plugin-backend/src/main/core/kotlinx/rpc/codegen/extension/ServiceDeclaration.kt‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,17 @@
11
/*
2-
* Copyright 2023-2024 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
2+
* Copyright 2023-2025 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
33
*/
44

55
package kotlinx.rpc.codegen.extension
66

7+
import kotlinx.rpc.codegen.common.RpcClassId
78
import org.jetbrains.kotlin.ir.declarations.IrClass
89
import org.jetbrains.kotlin.ir.declarations.IrProperty
910
import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
1011
import org.jetbrains.kotlin.ir.declarations.IrValueParameter
1112
import org.jetbrains.kotlin.ir.types.IrType
1213
import org.jetbrains.kotlin.ir.util.defaultType
14+
import org.jetbrains.kotlin.ir.util.hasAnnotation
1315
import org.jetbrains.kotlin.ir.util.kotlinFqName
1416

1517
class ServiceDeclaration(
@@ -18,6 +20,7 @@ class ServiceDeclaration(
1820
val methods: List<Method>,
1921
val fields: List<FlowField>,
2022
) {
23+
val isGrpc = service.hasAnnotation(RpcClassId.grpcAnnotation)
2124
val fqName = service.kotlinFqName.asString()
2225

2326
val serviceType = service.defaultType

‎compiler-plugin/compiler-plugin-common/src/main/core/kotlinx/rpc/codegen/common/Names.kt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023-2024 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
2+
* Copyright 2023-2025 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
33
*/
44

55
package kotlinx.rpc.codegen.common
@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.name.Name
1212
object RpcClassId {
1313
val remoteServiceInterface = ClassId(FqName("kotlinx.rpc"), Name.identifier("RemoteService"))
1414
val rpcAnnotation = ClassId(FqName("kotlinx.rpc.annotations"), Name.identifier("Rpc"))
15+
val grpcAnnotation = ClassId(FqName("kotlinx.rpc.grpc.annotations"), Name.identifier("Grpc"))
1516
val checkedTypeAnnotation = ClassId(FqName("kotlinx.rpc.annotations"), Name.identifier("CheckedTypeAnnotation"))
1617

1718
val serializableAnnotation = ClassId(FqName("kotlinx.serialization"), Name.identifier("Serializable"))

‎compiler-plugin/compiler-plugin-k2/src/main/core/kotlinx/rpc/codegen/FirGenerationKeys.kt‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023-2024 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
2+
* Copyright 2023-2025 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
33
*/
44

55
package kotlinx.rpc.codegen
@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.name.Name
1313
import org.jetbrains.kotlinx.serialization.compiler.fir.SerializationPluginKey
1414

1515
internal class RpcGeneratedStubKey(
16+
val isGrpc: Boolean,
1617
private val serviceName: Name,
1718
val functions: List<FirFunctionSymbol<*>>,
1819
) : GeneratedDeclarationKey() {
@@ -25,6 +26,7 @@ internal val FirBasedSymbol<*>.generatedRpcServiceStubKey: RpcGeneratedStubKey?
2526
(origin as? FirDeclarationOrigin.Plugin)?.key as? RpcGeneratedStubKey
2627

2728
internal class RpcGeneratedRpcMethodClassKey(
29+
val isGrpc: Boolean,
2830
val rpcMethod: FirFunctionSymbol<*>,
2931
) : GeneratedDeclarationKey() {
3032
val isObject = rpcMethod.valueParameterSymbols.isEmpty()

‎compiler-plugin/compiler-plugin-k2/src/main/core/kotlinx/rpc/codegen/FirRpcAdditionalCheckers.kt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023-2024 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
2+
* Copyright 2023-2025 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
33
*/
44

55
package kotlinx.rpc.codegen
@@ -25,6 +25,7 @@ class FirRpcAdditionalCheckers(
2525
) : FirAdditionalCheckersExtension(session) {
2626
override fun FirDeclarationPredicateRegistrar.registerPredicates() {
2727
register(FirRpcPredicates.rpc)
28+
register(FirRpcPredicates.grpc)
2829
register(FirRpcPredicates.checkedAnnotationMeta)
2930
}
3031

‎compiler-plugin/compiler-plugin-k2/src/main/core/kotlinx/rpc/codegen/FirRpcPredicates.kt‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,10 @@ object FirRpcPredicates {
1616
metaAnnotated(RpcClassId.rpcAnnotation.asSingleFqName(), includeItself = true)
1717
}
1818

19+
internal val grpc = DeclarationPredicate.create {
20+
annotated(RpcClassId.grpcAnnotation.asSingleFqName()) // @Grpc
21+
}
22+
1923
internal val checkedAnnotationMeta = DeclarationPredicate.create {
2024
metaAnnotated(RpcClassId.checkedTypeAnnotation.asSingleFqName(), includeItself = false)
2125
}

‎compiler-plugin/compiler-plugin-k2/src/main/core/kotlinx/rpc/codegen/FirRpcServiceGenerator.kt‎

Lines changed: 21 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023-2024 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
2+
* Copyright 2023-2025 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
33
*/
44

55
package kotlinx.rpc.codegen
@@ -75,6 +75,7 @@ class FirRpcServiceGenerator(
7575

7676
override fun FirDeclarationPredicateRegistrar.registerPredicates() {
7777
register(FirRpcPredicates.rpc)
78+
register(FirRpcPredicates.grpc)
7879
}
7980

8081
/**
@@ -112,7 +113,7 @@ class FirRpcServiceGenerator(
112113
val rpcMethodClassKey = classSymbol.generatedRpcMethodClassKey
113114

114115
return when {
115-
rpcMethodClassKey != null -> {
116+
rpcMethodClassKey != null && !rpcMethodClassKey.isGrpc -> {
116117
when {
117118
!rpcMethodClassKey.isObject -> setOf(
118119
SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT,
@@ -130,7 +131,10 @@ class FirRpcServiceGenerator(
130131
SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT
131132
}
132133

133-
classSymbol.isInterface && session.predicateBasedProvider.matches(FirRpcPredicates.rpc, classSymbol) -> {
134+
classSymbol.isInterface && (
135+
session.predicateBasedProvider.matches(FirRpcPredicates.rpc, classSymbol) ||
136+
session.predicateBasedProvider.matches(FirRpcPredicates.grpc, classSymbol)
137+
) -> {
134138
setOf(RpcNames.SERVICE_STUB_NAME)
135139
}
136140

@@ -159,7 +163,7 @@ class FirRpcServiceGenerator(
159163
generateRpcMethodClass(owner, name, rpcServiceStubKey)
160164
}
161165

162-
owner.generatedRpcMethodClassKey != null -> {
166+
owner.generatedRpcMethodClassKey != null && owner.generatedRpcMethodClassKey?.isGrpc == false -> {
163167
generateNestedClassLikeDeclarationWithSerialization(owner, name)
164168
}
165169

@@ -191,7 +195,7 @@ class FirRpcServiceGenerator(
191195
val methodName = name.rpcMethodName
192196
val rpcMethod = rpcServiceStubKey.functions.singleOrNull { it.name == methodName }
193197
?: return null
194-
val rpcMethodClassKey = RpcGeneratedRpcMethodClassKey(rpcMethod)
198+
val rpcMethodClassKey = RpcGeneratedRpcMethodClassKey(rpcServiceStubKey.isGrpc, rpcMethod)
195199
val classKind = if (rpcMethodClassKey.isObject) ClassKind.OBJECT else ClassKind.CLASS
196200

197201
val rpcMethodClass = createNestedClass(
@@ -204,13 +208,15 @@ class FirRpcServiceGenerator(
204208
modality = Modality.FINAL
205209
}
206210

207-
rpcMethodClass.addAnnotation(RpcClassId.serializableAnnotation, session)
211+
if (!session.predicateBasedProvider.matches(FirRpcPredicates.grpc, owner)) {
212+
rpcMethodClass.addAnnotation(RpcClassId.serializableAnnotation, session)
213+
}
208214

209215
/**
210216
* Required to pass isSerializableObjectAndNeedsFactory check
211217
* from [SerializationFirSupertypesExtension].
212218
*/
213-
if (!isJvmOrMetadata && rpcMethodClassKey.isObject) {
219+
if (!isJvmOrMetadata && rpcMethodClassKey.isObject && !rpcMethodClassKey.isGrpc) {
214220
rpcMethodClass.replaceSuperTypeRefs(createSerializationFactorySupertype())
215221
}
216222

@@ -265,7 +271,13 @@ class FirRpcServiceGenerator(
265271
.filterIsInstance<FirFunction>()
266272
.map { it.symbol }
267273

268-
return createNestedClass(owner, RpcNames.SERVICE_STUB_NAME, RpcGeneratedStubKey(owner.name, functions)) {
274+
val key = RpcGeneratedStubKey(
275+
isGrpc = session.predicateBasedProvider.matches(FirRpcPredicates.grpc, owner),
276+
serviceName = owner.name,
277+
functions = functions,
278+
)
279+
280+
return createNestedClass(owner, RpcNames.SERVICE_STUB_NAME, key) {
269281
visibility = Visibilities.Public
270282
modality = Modality.FINAL
271283
}.symbol
@@ -299,7 +311,7 @@ class FirRpcServiceGenerator(
299311
context: MemberGenerationContext,
300312
rpcMethodClassKey: RpcGeneratedRpcMethodClassKey,
301313
): Set<Name> {
302-
return if (rpcMethodClassKey.isObject) {
314+
return if (rpcMethodClassKey.isObject && !rpcMethodClassKey.isGrpc) {
303315
// add .serializer() method for a serializable object
304316
serializationExtension.getCallableNamesForClass(classSymbol, context)
305317
} else {

0 commit comments

Comments
 (0)