Grpc services followup

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

@@ -255,8 +255,8 @@ internal class RpcIrContext(
             grpcServiceDescriptor.namedFunction("delegate")
         }
 
-        val grpcMessageCodecResolverResolve by lazy {
-            grpcMessageCodecResolver.namedFunction("resolve")
+        val grpcMessageCodecResolverResolveOrNull by lazy {
+            grpcMessageCodecResolver.namedFunction("resolveOrNull")
         }
 
         private fun IrClassSymbol.namedFunction(name: String): IrSimpleFunction {

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

@@ -1155,11 +1155,11 @@ internal class RpcStubGenerator(
             "Only methods are allowed here"
         }
 
-        check(callable.arguments.size == 1) {
-            "Only single argument methods are allowed here"
+        check(callable.arguments.size <= 1) {
+            "Only single or none argument methods are allowed here"
         }
 
-        val requestParameterType = callable.arguments[0].type
+        val requestParameterType = callable.arguments.getOrNull(0)?.type ?: ctx.irBuiltIns.unitType
         val responseParameterType = callable.function.returnType
 
         val requestType: IrType = requestParameterType.unwrapFlow()
@@ -1262,7 +1262,7 @@ internal class RpcStubGenerator(
                     startOffset = UNDEFINED_OFFSET,
                     endOffset = UNDEFINED_OFFSET,
                     type = ctx.grpcMessageCodec.typeWith(type),
-                    symbol = ctx.functions.grpcMessageCodecResolverResolve.symbol,
+                    symbol = ctx.functions.grpcMessageCodecResolverResolveOrNull.symbol,
                     typeArgumentsCount = 0,
                     valueArgumentsCount = 1,
                 )

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.name.Name
 object RpcClassId {
     val rpcAnnotation = ClassId(FqName("kotlinx.rpc.annotations"), Name.identifier("Rpc"))
     val grpcAnnotation = ClassId(FqName("kotlinx.rpc.grpc.annotations"), Name.identifier("Grpc"))
+    val withCodecAnnotation = ClassId(FqName("kotlinx.rpc.grpc.codec"), Name.identifier("WithCodec"))
     val checkedTypeAnnotation = ClassId(FqName("kotlinx.rpc.annotations"), Name.identifier("CheckedTypeAnnotation"))
 
     val serializableAnnotation = ClassId(FqName("kotlinx.serialization"), Name.identifier("Serializable"))
@@ -19,6 +20,8 @@ object RpcClassId {
     val flow = ClassId(FqName("kotlinx.coroutines.flow"), Name.identifier("Flow"))
     val sharedFlow = ClassId(FqName("kotlinx.coroutines.flow"), Name.identifier("SharedFlow"))
     val stateFlow = ClassId(FqName("kotlinx.coroutines.flow"), Name.identifier("StateFlow"))
+
+    val messageCodec = ClassId(FqName("kotlinx.rpc.grpc.codec"), Name.identifier("MessageCodec"))
 }
 
 object RpcNames {

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

@@ -23,6 +23,8 @@ class FirRpcAdditionalCheckers(
 ) : FirAdditionalCheckersExtension(session) {
     override fun FirDeclarationPredicateRegistrar.registerPredicates() {
         register(FirRpcPredicates.rpc)
+        register(FirRpcPredicates.grpc)
+        register(FirRpcPredicates.withCodec)
         register(FirRpcPredicates.checkedAnnotationMeta)
     }
 

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

@@ -12,6 +12,14 @@ object FirRpcPredicates {
         metaAnnotated(RpcClassId.rpcAnnotation.asSingleFqName(), includeItself = true) // @Rpc
     }
 
+    internal val grpc = DeclarationPredicate.create {
+        metaAnnotated(RpcClassId.grpcAnnotation.asSingleFqName(), includeItself = true)
+    }
+
+    internal val withCodec = DeclarationPredicate.create {
+        annotated(RpcClassId.withCodecAnnotation.asSingleFqName())
+    }
+
     internal val checkedAnnotationMeta = DeclarationPredicate.create {
         metaAnnotated(RpcClassId.checkedTypeAnnotation.asSingleFqName(), includeItself = false)
     }

compiler-plugin/compiler-plugin-k2/src/main/kotlin/kotlinx/rpc/codegen/checkers/FirGrpcServiceDeclarationChecker.kt

@@ -0,0 +1,76 @@
+/*
+ * Copyright 2023-2025 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.rpc.codegen.checkers
+
+import kotlinx.rpc.codegen.FirRpcPredicates
+import kotlinx.rpc.codegen.checkers.diagnostics.FirGrpcDiagnostics
+import kotlinx.rpc.codegen.checkers.util.functionParametersRecursionCheck
+import kotlinx.rpc.codegen.common.RpcClassId
+import kotlinx.rpc.codegen.vsApi
+import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
+import org.jetbrains.kotlin.diagnostics.reportOn
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
+import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.fir.types.isMarkedNullable
+
+object FirGrpcServiceDeclarationChecker {
+    fun check(
+        declaration: FirRegularClass,
+        context: CheckerContext,
+        reporter: DiagnosticReporter,
+    ) {
+        if (!context.session.predicateBasedProvider.matches(FirRpcPredicates.grpc, declaration)) {
+            return
+        }
+
+        vsApi {
+            declaration
+                .declarationsVS(context.session)
+                .filterIsInstance<FirNamedFunctionSymbol>()
+        }.onEach { function ->
+            if (function.valueParameterSymbols.size > 1) {
+                reporter.reportOn(
+                    source = function.source,
+                    factory = FirGrpcDiagnostics.MULTIPLE_PARAMETERS_IN_GRPC_SERVICE,
+                    context = context,
+                )
+            }
+
+            if (function.valueParameterSymbols.size == 1) {
+                val parameterSymbol = function.valueParameterSymbols[0]
+                if (parameterSymbol.resolvedReturnType.isMarkedNullable) {
+                    reporter.reportOn(
+                        source = parameterSymbol.source,
+                        factory = FirGrpcDiagnostics.NULLABLE_PARAMETER_IN_GRPC_SERVICE,
+                        context = context,
+                    )
+                }
+
+                functionParametersRecursionCheck(
+                    function = function,
+                    context = context,
+                ) { source, symbol, parents ->
+                    if (symbol.classId == RpcClassId.flow && parents.isNotEmpty()) {
+                        reporter.reportOn(
+                            source = source,
+                            factory = FirGrpcDiagnostics.NON_TOP_LEVEL_CLIENT_STREAMING_IN_RPC_SERVICE,
+                            context = context,
+                        )
+                    }
+                }
+            }
+
+            if (function.resolvedReturnType.isMarkedNullable) {
+                reporter.reportOn(
+                    source = function.resolvedReturnTypeRef.source,
+                    factory = FirGrpcDiagnostics.NULLABLE_RETURN_TYPE_IN_GRPC_SERVICE,
+                    context = context,
+                )
+            }
+        }
+    }
+}

compiler-plugin/compiler-plugin-k2/src/main/kotlin/kotlinx/rpc/codegen/checkers/FirRpcCheckers.kt

@@ -18,6 +18,8 @@ class FirRpcDeclarationCheckers(ctx: FirCheckersContext) : DeclarationCheckers()
         FirRpcAnnotationCheckerVS(),
         FirRpcStrictModeClassCheckerVS(),
         FirRpcServiceDeclarationCheckerVS(ctx),
+        FirGrpcServiceDeclarationCheckerVS(),
+        FirWithCodecDeclarationCheckerVS(),
     )
 
     override val classCheckers: Set<FirClassChecker> = setOf(

compiler-plugin/compiler-plugin-k2/src/main/kotlin/kotlinx/rpc/codegen/checkers/FirRpcStrictModeClassChecker.kt

@@ -6,26 +6,19 @@ package kotlinx.rpc.codegen.checkers
 
 import kotlinx.rpc.codegen.FirRpcPredicates
 import kotlinx.rpc.codegen.checkers.diagnostics.FirRpcStrictModeDiagnostics
+import kotlinx.rpc.codegen.checkers.util.functionParametersRecursionCheck
 import kotlinx.rpc.codegen.common.RpcClassId
 import kotlinx.rpc.codegen.vsApi
 import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.diagnostics.KtDiagnosticFactory0
 import org.jetbrains.kotlin.diagnostics.reportOn
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
-import org.jetbrains.kotlin.fir.analysis.checkers.extractArgumentsTypeRefAndSource
-import org.jetbrains.kotlin.fir.analysis.checkers.toClassLikeSymbol
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.utils.isSuspend
 import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
-import org.jetbrains.kotlin.fir.scopes.impl.toConeType
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
-import org.jetbrains.kotlin.fir.types.FirTypeRef
-import org.jetbrains.kotlin.utils.memoryOptimizedMap
-import org.jetbrains.kotlin.utils.memoryOptimizedPlus
 
 object FirRpcStrictModeClassChecker {
     fun check(
@@ -37,124 +30,51 @@ object FirRpcStrictModeClassChecker {
             return
         }
 
-        val serializablePropertiesProvider = context.session.serializablePropertiesProvider
         vsApi { declaration.declarationsVS(context.session) }.forEach { declaration ->
             when (declaration) {
                 is FirPropertySymbol -> {
                     reporter.reportOn(declaration.source, FirRpcStrictModeDiagnostics.FIELD_IN_RPC_SERVICE, context)
                 }
 
                 is FirNamedFunctionSymbol -> {
-                    checkFunction(declaration, context, reporter, serializablePropertiesProvider)
-                }
-
-                else -> {}
-            }
-        }
-    }
-
-    private fun checkFunction(
-        function: FirNamedFunctionSymbol,
-        context: CheckerContext,
-        reporter: DiagnosticReporter,
-        serializablePropertiesProvider: FirSerializablePropertiesProvider,
-    ) {
-        fun reportOn(element: KtSourceElement?, checker: FirRpcStrictModeDiagnostics.() -> KtDiagnosticFactory0?) {
-            reporter.reportOn(element, FirRpcStrictModeDiagnostics.checker() ?: return, context)
-        }
-
-        val returnClassSymbol = vsApi {
-            function.resolvedReturnTypeRef.coneTypeVS.toClassSymbolVS(context.session)
-        }
-
-        val types = function.valueParameterSymbols.memoryOptimizedMap { parameter ->
-            parameter.source to vsApi {
-                parameter.resolvedReturnTypeRef
-            }
-        } memoryOptimizedPlus (function.resolvedReturnTypeRef.source to function.resolvedReturnTypeRef)
+                    fun reportOn(element: KtSourceElement?, checker: FirRpcStrictModeDiagnostics.() -> KtDiagnosticFactory0) {
+                        reporter.reportOn(element, FirRpcStrictModeDiagnostics.checker(), context)
+                    }
 
-        types.forEach { (source, symbol) ->
-            checkSerializableTypes(
-                context = context,
-                typeRef = symbol,
-                serializablePropertiesProvider = serializablePropertiesProvider,
-            ) { symbol, parents ->
-                when (symbol.classId) {
-                    RpcClassId.stateFlow -> {
-                        reportOn(source) { STATE_FLOW_IN_RPC_SERVICE }
+                    val returnClassSymbol = vsApi {
+                        declaration.resolvedReturnTypeRef.coneTypeVS.toClassSymbolVS(context.session)
                     }
 
-                    RpcClassId.sharedFlow -> {
-                        reportOn(source) { SHARED_FLOW_IN_RPC_SERVICE }
+                    if (returnClassSymbol?.classId == RpcClassId.flow && declaration.isSuspend) {
+                        reportOn(declaration.source) { SUSPENDING_SERVER_STREAMING_IN_RPC_SERVICE }
                     }
 
-                    RpcClassId.flow -> {
-                        if (parents.any { it.classId == RpcClassId.flow }) {
-                            reportOn(source) { NESTED_STREAMING_IN_RPC_SERVICE }
-                        } else if (parents.isNotEmpty() && parents[0] == returnClassSymbol) {
-                            reportOn(source) { NON_TOP_LEVEL_SERVER_STREAMING_IN_RPC_SERVICE }
+                    functionParametersRecursionCheck(
+                        function = declaration,
+                        context = context,
+                    ) { source, symbol, parents ->
+                        when (symbol.classId) {
+                            RpcClassId.stateFlow -> {
+                                reportOn(source) { STATE_FLOW_IN_RPC_SERVICE }
+                            }
+
+                            RpcClassId.sharedFlow -> {
+                                reportOn(source) { SHARED_FLOW_IN_RPC_SERVICE }
+                            }
+
+                            RpcClassId.flow -> {
+                                if (parents.any { it.classId == RpcClassId.flow }) {
+                                    reportOn(source) { NESTED_STREAMING_IN_RPC_SERVICE }
+                                } else if (parents.isNotEmpty() && parents[0] == returnClassSymbol) {
+                                    reportOn(source) { NON_TOP_LEVEL_SERVER_STREAMING_IN_RPC_SERVICE }
+                                }
+                            }
                         }
                     }
                 }
-            }
-        }
 
-        if (returnClassSymbol?.classId == RpcClassId.flow && function.isSuspend) {
-            reportOn(function.source) { SUSPENDING_SERVER_STREAMING_IN_RPC_SERVICE }
-        }
-    }
-
-    private fun checkSerializableTypes(
-        context: CheckerContext,
-        typeRef: FirTypeRef,
-        serializablePropertiesProvider: FirSerializablePropertiesProvider,
-        parentContext: List<FirClassLikeSymbol<*>> = emptyList(),
-        checker: (FirClassLikeSymbol<*>, List<FirClassLikeSymbol<*>>) -> Unit,
-    ) {
-        val symbol = typeRef.toClassLikeSymbol(context.session) ?: return
-
-        checker(symbol, parentContext)
-
-        if (symbol !is FirClassSymbol<*>) {
-            return
-        }
-
-        val nextContext = parentContext memoryOptimizedPlus symbol
-
-        if (symbol in parentContext && symbol.typeParameterSymbols.isEmpty()) {
-            return
-        }
-
-        val typeParameters = extractArgumentsTypeRefAndSource(typeRef)
-            .orEmpty()
-            .withIndex()
-            .associate { (i, refSource) ->
-                symbol.typeParameterSymbols[i].toConeType() to refSource.typeRef
+                else -> {}
             }
-
-        val flowProps: List<FirTypeRef> = if (symbol.classId == RpcClassId.flow) {
-            listOf(typeParameters.values.toList()[0]!!)
-        } else {
-            emptyList()
         }
-
-        serializablePropertiesProvider.getSerializablePropertiesForClass(symbol)
-            .mapNotNull { property ->
-                val resolvedTypeRef = property.resolvedReturnTypeRef
-                if (resolvedTypeRef.toClassLikeSymbol(context.session) != null) {
-                    resolvedTypeRef
-                } else {
-                    typeParameters[property.resolvedReturnType]
-                }
-            }.memoryOptimizedPlus(flowProps)
-            .forEach { symbol ->
-                checkSerializableTypes(
-                    context = context,
-                    typeRef = symbol,
-                    serializablePropertiesProvider = serializablePropertiesProvider,
-                    parentContext = nextContext,
-                    checker = checker,
-                )
-            }
     }
 }