Support for @rpc annotation

Author: Mr3zee

compiler-plugin/compiler-plugin-cli/src/main/latest/kotlinx/rpc/codegen/RPCCompilerPlugin.kt

@@ -34,5 +34,5 @@ fun CompilerPluginRegistrar.ExtensionStorage.registerRpcExtensions(configuration
     VersionSpecificApi.INSTANCE = VersionSpecificApiImpl
 
     IrGenerationExtension.registerExtension(RPCIrExtension(configuration))
-    FirExtensionRegistrarAdapter.registerExtension(FirRPCExtensionRegistrar(configuration))
+    FirExtensionRegistrarAdapter.registerExtension(FirRpcExtensionRegistrar(configuration))
 }

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

@@ -8,8 +8,9 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 
-object ClassDeclarations {
+object RpcClassId {
     val rpcInterface = ClassId(FqName("kotlinx.rpc"), Name.identifier("RPC"))
+    val rpcAnnotation = ClassId(FqName("kotlinx.rpc"), Name.identifier("Rpc"))
 
     val serializableAnnotation = ClassId(FqName("kotlinx.serialization"), Name.identifier("Serializable"))
     val contextualAnnotation = ClassId(FqName("kotlinx.serialization"), Name.identifier("Contextual"))

compiler-plugin/compiler-plugin-k2/src/main/latest/kotlinx/rpc/codegen/FirRPCExtensionRegistrar.kt

@@ -8,12 +8,16 @@ import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
+import org.jetbrains.kotlin.fir.analysis.extensions.FirAdditionalCheckersExtension.Factory as CFactory
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension.Factory as GFactory
+import org.jetbrains.kotlin.fir.extensions.FirSupertypeGenerationExtension.Factory as SFactory
 
-class FirRPCExtensionRegistrar(private val configuration: CompilerConfiguration) : FirExtensionRegistrar() {
+class FirRpcExtensionRegistrar(private val configuration: CompilerConfiguration) : FirExtensionRegistrar() {
     override fun ExtensionRegistrarContext.configurePlugin() {
         val logger = configuration.get(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY, MessageCollector.NONE)
 
-        +GFactory { FirRPCServiceGenerator(it, logger) }
+        +CFactory { FirRpcCheckers(it)  }
+        +GFactory { FirRpcServiceGenerator(it, logger) }
+        +SFactory { FirRpcSupertypeGenerator(it, logger) }
     }
 }

compiler-plugin/compiler-plugin-k2/src/main/latest/kotlinx/rpc/codegen/FirRPCServiceGenerator.kt

@@ -4,7 +4,7 @@
 
 package kotlinx.rpc.codegen
 
-import kotlinx.rpc.codegen.common.ClassDeclarations
+import kotlinx.rpc.codegen.common.RpcClassId
 import kotlinx.rpc.codegen.common.RpcNames
 import kotlinx.rpc.codegen.common.rpcMethodClassName
 import kotlinx.rpc.codegen.common.rpcMethodName
@@ -19,9 +19,7 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirFunction
 import org.jetbrains.kotlin.fir.declarations.utils.isInterface
 import org.jetbrains.kotlin.fir.declarations.utils.visibility
-import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
-import org.jetbrains.kotlin.fir.extensions.MemberGenerationContext
-import org.jetbrains.kotlin.fir.extensions.NestedClassGenerationContext
+import org.jetbrains.kotlin.fir.extensions.*
 import org.jetbrains.kotlin.fir.moduleData
 import org.jetbrains.kotlin.fir.plugin.*
 import org.jetbrains.kotlin.fir.symbols.SymbolInternals
@@ -44,22 +42,7 @@ import org.jetbrains.kotlinx.serialization.compiler.resolve.SerializationPackage
  *
  * ## General idea
  *
- * [getNestedClassifiersNames] should return a set of [Name]s to generate for.
- * For these names [generateNestedClassLikeDeclaration] can generate some nested classes,
- * which is what we need.
- *
- * But the catch is that we cannot say for sure, if we need to generate a class
- * while in [getNestedClassifiersNames], but if we do not return anything
- * [generateNestedClassLikeDeclaration] will not be called.
- * We need to generate a class if only the current declaration is an RPC interface
- * (inherits kotlinx.rpc.RPC). There is no resolved supertypes in [getNestedClassifiersNames],
- * But, if the potentially generated class is not referenced anywhere,
- * then [generateNestedClassLikeDeclaration] will already have supertypes resolved,
- * so we can use this info to check the actual supertypes for RPC interface.
- *
- * So we always return a class name that may be generated.
- * And then, in [generateNestedClassLikeDeclaration] we do the actual check with the resolved supertypes
- * and generate a class if needed, otherwise returning null.
+ * Detect `@Rpc` annotation - generate stub classes.
  *
  * ## Usage of kotlinx.serialization plugin
  *
@@ -68,38 +51,38 @@ import org.jetbrains.kotlinx.serialization.compiler.resolve.SerializationPackage
  * We generate classes that are marked `@Serializable`.
  * In that case, the serialization plugin will not be able to pick up those classes and process them accordingly.
  *
- * That's why we have an instance of this plugin, which we call only on our generated classes - [serializationExtension]
+ * That is why we have an instance of this plugin, which we call only on our generated classes - [serializationExtension]
  *
  * Not all the methods that we would like to call are public, so we access them via reflection:
  * - [generateCompanionDeclaration]
  * - [generateSerializerImplClass]
  *
  * This is basically copying the behavior of the actual plugin but isolated only for our generated classes.
  */
-class FirRPCServiceGenerator(
+class FirRpcServiceGenerator(
     session: FirSession,
     @Suppress("unused")
     private val logger: MessageCollector,
 ) : FirDeclarationGenerationExtension(session) {
     private val serializationExtension = SerializationFirResolveExtension(session)
     private val isJvmOrMetadata = !session.moduleData.platform.run { isJs() || isWasm() || isNative() }
 
+    override fun FirDeclarationPredicateRegistrar.registerPredicates() {
+        register(FirRpcPredicates.rpc)
+    }
+
     /**
      * Generates nested classifiers.
      *
      * They can be of three kinds:
      * - Nested Service Stub class.
      * In that case [classSymbol] will not have any RPC-generated [FirClassSymbol.origin].
-     * Rge only check we do - is we check that the declaration is an interface,
-     * and return [RpcNames.SERVICE_STUB_NAME].
-     * We cannot be sure if the declaration is actually an RPC service,
-     * because superTypes are not resolved during that stage.
-     * We postpone this check until [generateNestedClassLikeDeclaration].
+     * The only check we do - presence of the `@Rpc` annotation and return [RpcNames.SERVICE_STUB_NAME].
      *
      * - Companion object of the service stub and method classes.
      * If we generate this companion object, we will have [FirClassSymbol.origin]
      * of [classSymbol] be set to [RPCGeneratedStubKey],
-     * because we are inside the previously generated service stub class.
+     * because we're inside the previously generated service stub class.
      * The same goes for method classes too.
      * So we return [SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT]
      * and a list of method class names.
@@ -141,7 +124,7 @@ class FirRPCServiceGenerator(
                         SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT
             }
 
-            classSymbol.isInterface -> {
+            classSymbol.isInterface && session.predicateBasedProvider.matches(FirRpcPredicates.rpc, classSymbol) -> {
                 setOf(RpcNames.SERVICE_STUB_NAME)
             }
 
@@ -158,7 +141,7 @@ class FirRPCServiceGenerator(
     override fun generateNestedClassLikeDeclaration(
         owner: FirClassSymbol<*>,
         name: Name,
-        context: NestedClassGenerationContext
+        context: NestedClassGenerationContext,
     ): FirClassLikeSymbol<*>? {
         val rpcServiceStubKey = owner.generatedRpcServiceStubKey
         return when {
@@ -214,7 +197,7 @@ class FirRPCServiceGenerator(
             modality = Modality.FINAL
         }
 
-        rpcMethodClass.addAnnotation(ClassDeclarations.serializableAnnotation, session)
+        rpcMethodClass.addAnnotation(RpcClassId.serializableAnnotation, session)
 
         /**
          * Required to pass isSerializableObjectAndNeedsFactory check
@@ -264,18 +247,10 @@ class FirRPCServiceGenerator(
     }
 
     /**
-     * Checks whether the [owner] class is actually an RPC service
-     * (the supertypes are resolved at this stage,
-     * as the [RpcNames.SERVICE_STUB_NAME] is not references anywhere)
-     *
-     * If the [owner] is an RPC service - generates its service stub.
+     * Generates [owner]'s service stub.
      * Scrapes the functions from the [owner] to generate method classes.
      */
     private fun generateRpcServiceStubClass(owner: FirClassSymbol<*>): FirRegularClassSymbol? {
-        owner.resolvedSuperTypes.find {
-            it.classId == ClassDeclarations.rpcInterface
-        } ?: return null
-
         @OptIn(SymbolInternals::class)
         val functions = owner.fir.declarations
             .filterIsInstance<FirFunction>()
@@ -306,7 +281,7 @@ class FirRPCServiceGenerator(
     }
 
     /**
-     * If the method does not have any parameters, it is an object,
+     * If the method doesn't have any parameters, it is an object,
      * and its only callable names are constructor, and the ones provided by the [serializationExtension].
      * Otherwise, the callable names are the names of the method parameters and the constructor.
      */
@@ -322,7 +297,7 @@ class FirRPCServiceGenerator(
             rpcMethodClassKey.rpcMethod.valueParameterSymbols.map { it.name }.toSet()
         } + SpecialNames.INIT
 
-        // ^ init is necessary either way, as serialization does not add it for a serializable object
+        // ^ init is necessary either way, as serialization doesn't add it for a serializable object
     }
 
     override fun generateConstructors(context: MemberGenerationContext): List<FirConstructorSymbol> {
@@ -361,7 +336,7 @@ class FirRPCServiceGenerator(
 
     override fun generateProperties(
         callableId: CallableId,
-        context: MemberGenerationContext?
+        context: MemberGenerationContext?,
     ): List<FirPropertySymbol> {
         context ?: return emptyList()
 
@@ -399,14 +374,14 @@ class FirRPCServiceGenerator(
             returnType = valueParam.resolvedReturnType,
         ).apply {
             if (valueParam.resolvedReturnType.requiresContextual()) {
-                addAnnotation(ClassDeclarations.contextualAnnotation, session)
+                addAnnotation(RpcClassId.contextualAnnotation, session)
             }
         }.symbol.let(::listOf)
     }
 
     private fun ConeKotlinType.requiresContextual(): Boolean {
         return when (classId) {
-            ClassDeclarations.flow, ClassDeclarations.sharedFlow, ClassDeclarations.stateFlow -> true
+            RpcClassId.flow, RpcClassId.sharedFlow, RpcClassId.stateFlow -> true
             else -> false
         }
     }
@@ -416,7 +391,7 @@ class FirRPCServiceGenerator(
      */
     override fun generateFunctions(
         callableId: CallableId,
-        context: MemberGenerationContext?
+        context: MemberGenerationContext?,
     ): List<FirNamedFunctionSymbol> {
         val owner = context?.owner ?: return emptyList()
 

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

@@ -0,0 +1,19 @@
+/*
+ * Copyright 2023-2024 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.rpc.codegen
+
+import kotlinx.rpc.codegen.checkers.FirRpcDeclarationCheckers
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.analysis.checkers.declaration.DeclarationCheckers
+import org.jetbrains.kotlin.fir.analysis.extensions.FirAdditionalCheckersExtension
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
+
+class FirRpcCheckers(session: FirSession) : FirAdditionalCheckersExtension(session) {
+    override fun FirDeclarationPredicateRegistrar.registerPredicates() {
+        register(FirRpcPredicates.rpc)
+    }
+
+    override val declarationCheckers: DeclarationCheckers = FirRpcDeclarationCheckers
+}

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

@@ -0,0 +1,14 @@
+/*
+ * Copyright 2023-2024 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.rpc.codegen
+
+import kotlinx.rpc.codegen.common.RpcClassId
+import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
+
+object FirRpcPredicates {
+    internal val rpc = DeclarationPredicate.create {
+        annotated(RpcClassId.rpcAnnotation.asSingleFqName()) // @Rpc
+    }
+}

compiler-plugin/compiler-plugin-k2/src/main/latest/kotlinx/rpc/codegen/FirRpcSupertypeGenerator.kt

@@ -0,0 +1,46 @@
+/*
+ * Copyright 2023-2024 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.rpc.codegen
+
+import kotlinx.rpc.codegen.common.RpcClassId
+import org.jetbrains.kotlin.cli.common.messages.MessageCollector
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
+import org.jetbrains.kotlin.fir.extensions.FirSupertypeGenerationExtension
+import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.classId
+import org.jetbrains.kotlin.fir.types.constructClassLikeType
+import org.jetbrains.kotlin.fir.types.toFirResolvedTypeRef
+
+class FirRpcSupertypeGenerator(
+    session: FirSession,
+    @Suppress("unused") private val logger: MessageCollector,
+) : FirSupertypeGenerationExtension(session) {
+    override fun FirDeclarationPredicateRegistrar.registerPredicates() {
+        register(FirRpcPredicates.rpc)
+    }
+
+    override fun needTransformSupertypes(declaration: FirClassLikeDeclaration): Boolean {
+        return session.predicateBasedProvider.matches(FirRpcPredicates.rpc, declaration)
+    }
+
+    override fun computeAdditionalSupertypes(
+        classLikeDeclaration: FirClassLikeDeclaration,
+        resolvedSupertypes: List<FirResolvedTypeRef>,
+        typeResolver: TypeResolveService,
+    ): List<FirResolvedTypeRef> {
+        if (resolvedSupertypes.any { it.type.classId == RpcClassId.rpcInterface }) {
+            return emptyList()
+        }
+
+        return listOf(
+            RpcClassId.rpcInterface
+                .constructClassLikeType(emptyArray(), isNullable = false)
+                .toFirResolvedTypeRef()
+        )
+    }
+}

compiler-plugin/compiler-plugin-k2/src/main/latest/kotlinx/rpc/codegen/FirRpcUtils.kt

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2023-2024 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.rpc.codegen
+
+import kotlinx.rpc.codegen.common.RpcClassId
+import org.jetbrains.kotlin.KtSourceElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.getAnnotationByClassId
+import org.jetbrains.kotlin.fir.expressions.FirAnnotation
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
+import org.jetbrains.kotlin.fir.types.classId
+
+fun FirClassSymbol<*>.isRpc(): Boolean = resolvedSuperTypes.any {
+    it.classId == RpcClassId.rpcInterface
+}
+
+fun FirBasedSymbol<*>.rpcAnnotationSource(session: FirSession): KtSourceElement? {
+    return rpcAnnotation(session)?.source
+}
+
+fun FirBasedSymbol<*>.rpcAnnotation(session: FirSession): FirAnnotation? {
+    return resolvedCompilerAnnotationsWithClassIds.rpcAnnotation(session)
+}
+
+fun List<FirAnnotation>.rpcAnnotation(session: FirSession): FirAnnotation? {
+    return getAnnotationByClassId(RpcClassId.rpcAnnotation, session)
+}

compiler-plugin/compiler-plugin-k2/src/main/latest/kotlinx/rpc/codegen/checkers/FirRpcAnnotationChecker.kt

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2023-2024 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.FirRpcDiagnostics
+import kotlinx.rpc.codegen.isRpc
+import kotlinx.rpc.codegen.rpcAnnotationSource
+import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
+import org.jetbrains.kotlin.diagnostics.reportOn
+import org.jetbrains.kotlin.fir.analysis.checkers.MppCheckerKind
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirRegularClassChecker
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.declarations.utils.isInterface
+import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
+
+object FirRpcAnnotationChecker : FirRegularClassChecker(MppCheckerKind.Common) {
+    override fun check(
+        declaration: FirRegularClass,
+        context: CheckerContext,
+        reporter: DiagnosticReporter,
+    ) {
+        val rpcAnnotated = context.session.predicateBasedProvider.matches(FirRpcPredicates.rpc, declaration)
+
+        if (!declaration.isInterface && rpcAnnotated) {
+            reporter.reportOn(
+                source = declaration.symbol.rpcAnnotationSource(context.session),
+                factory = FirRpcDiagnostics.WRONG_RPC_ANNOTATION_TARGET,
+                context = context,
+            )
+        }
+
+        if (declaration.symbol.isRpc() && !rpcAnnotated) {
+            reporter.reportOn(
+                source = declaration.symbol.rpcAnnotationSource(context.session),
+                factory = FirRpcDiagnostics.MISSING_RPC_ANNOTATION,
+                context = context,
+            )
+        }
+    }
+}