Strict mode

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

@@ -6,7 +6,7 @@ package kotlinx.rpc.codegen
 
 import kotlinx.rpc.codegen.extension.RpcIrExtension
 import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
-import org.jetbrains.kotlin.compiler.plugin.CliOption
+import org.jetbrains.kotlin.compiler.plugin.AbstractCliOption
 import org.jetbrains.kotlin.compiler.plugin.CommandLineProcessor
 import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar
 import org.jetbrains.kotlin.compiler.plugin.ExperimentalCompilerApi
@@ -15,9 +15,25 @@ import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrarAdapter
 
 @OptIn(ExperimentalCompilerApi::class)
 class RpcCommandLineProcessor : CommandLineProcessor {
-    override val pluginId = "kotlinx.rpc.compiler-plugin"
-
-    override val pluginOptions = emptyList<CliOption>()
+    override val pluginId = "kotlinx-rpc"
+
+    override val pluginOptions = listOf(
+        StrictModeCliOptions.STATE_FLOW,
+        StrictModeCliOptions.SHARED_FLOW,
+        StrictModeCliOptions.NESTED_FLOW,
+        StrictModeCliOptions.STREAM_SCOPED_FUNCTIONS,
+        StrictModeCliOptions.SUSPENDING_SERVER_STREAMING,
+        StrictModeCliOptions.NOT_TOP_LEVEL_SERVER_FLOW,
+        StrictModeCliOptions.FIELDS,
+    )
+
+    override fun processOption(
+        option: AbstractCliOption,
+        value: String,
+        configuration: CompilerConfiguration,
+    ) {
+        option.processAsStrictModeOption(value, configuration)
+    }
 }
 
 @OptIn(ExperimentalCompilerApi::class)

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

@@ -4,6 +4,7 @@
 
 package kotlinx.rpc.codegen.common
 
+import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -21,6 +22,16 @@ object RpcClassId {
     val stateFlow = ClassId(FqName("kotlinx.coroutines.flow"), Name.identifier("StateFlow"))
 }
 
+object RpcCallableId {
+    val streamScoped = CallableId(FqName("kotlinx.rpc.krpc"), Name.identifier("streamScoped"))
+    val withStreamScope = CallableId(FqName("kotlinx.rpc.krpc"), Name.identifier("withStreamScope"))
+    val StreamScope = CallableId(FqName("kotlinx.rpc.krpc"), Name.identifier("StreamScope"))
+    val invokeOnStreamScopeCompletion = CallableId(
+        FqName("kotlinx.rpc.krpc"),
+        Name.identifier("invokeOnStreamScopeCompletion"),
+    )
+}
+
 object RpcNames {
     val SERVICE_STUB_NAME: Name = Name.identifier("\$rpcServiceStub")
 

compiler-plugin/compiler-plugin-k2/build.gradle.kts

@@ -2,6 +2,7 @@
  * Copyright 2023-2024 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
  */
 
+import com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar
 import org.jetbrains.kotlin.gradle.dsl.ExplicitApiMode
 import util.enableContextReceivers
 import util.otherwise
@@ -10,6 +11,49 @@ import util.whenForIde
 plugins {
     alias(libs.plugins.conventions.jvm)
     alias(libs.plugins.compiler.specific.module)
+    alias(libs.plugins.shadow.jar)
+}
+
+/*
+This is a hack to solve the next problem:
+
+There is PsiElement class that is used in compiler and plugin.
+
+If we use kotlin-compiler-embeddable.jar to compile the module,
+PsiElement is org.jetbrains.kotlin.com.jetbrains.intellij.PsiElement.
+And it works ok in the user projects!
+
+But.
+
+If we run tests, which use kotlin-compiler.jar, we run into ClassNotFoundException.
+Because the class it has in th classpath is com.jetbrains.intellij.PsiElement
+
+- Alright, we can use kotlin-compiler.jar to compile the plugin.
+- No, because the same error now will occur in the user projects,
+  but it is com.jetbrains.intellij.PsiElement that is not found.
+
+- Ok, we can use kotlin-compiler-embeddable.jar to compile tests.
+- Then we ran into java.lang.VerifyError: Bad type on operand stack, which I have no idea how to fix.
+
+This solution replaces org.jetbrains.kotlin.com.jetbrains.intellij.PsiElement usages in plugin
+with com.jetbrains.intellij.PsiElement only for the tests, fixing both use cases.
+It is basically a reverse engineering of what Kotlin does for the embedded jar.
+ */
+val shadowJar = tasks.named<ShadowJar>("shadowJar") {
+    configurations = listOf(project.configurations.compileClasspath.get())
+    relocate("org.jetbrains.kotlin.com.intellij.psi", "com.intellij.psi")
+
+    exclude("javaslang/**")
+    exclude("kotlin/**")
+    exclude("messages/**")
+    exclude("misc/**")
+    exclude("org/**")
+
+    archiveFileName.set("plugin-k2-for-tests.jar")
+}
+
+tasks.jar {
+    finalizedBy(shadowJar)
 }
 
 kotlin {

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

@@ -7,6 +7,7 @@ package kotlinx.rpc.codegen
 import kotlinx.rpc.codegen.checkers.FirCheckedAnnotationHelper
 import kotlinx.rpc.codegen.checkers.FirRpcDeclarationCheckers
 import kotlinx.rpc.codegen.checkers.FirRpcExpressionCheckers
+import kotlinx.rpc.codegen.checkers.diagnostics.FirRpcStrictModeDiagnostics
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.checkers.declaration.DeclarationCheckers
 import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExpressionCheckers
@@ -16,19 +17,29 @@ import org.jetbrains.kotlin.fir.caches.firCachesFactory
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 
-class FirRpcCheckers(session: FirSession, serializationIsPresent: Boolean) : FirAdditionalCheckersExtension(session) {
+class FirRpcAdditionalCheckers(
+    session: FirSession,
+    serializationIsPresent: Boolean,
+    modes: StrictModeAggregator,
+) : FirAdditionalCheckersExtension(session) {
     override fun FirDeclarationPredicateRegistrar.registerPredicates() {
         register(FirRpcPredicates.rpc)
         register(FirRpcPredicates.checkedAnnotationMeta)
     }
 
-    private val ctx = FirCheckersContext(session, serializationIsPresent)
+    private val ctx = FirCheckersContext(session, serializationIsPresent, modes)
 
     override val declarationCheckers: DeclarationCheckers = FirRpcDeclarationCheckers(ctx)
     override val expressionCheckers: ExpressionCheckers = FirRpcExpressionCheckers(ctx)
 }
 
-class FirCheckersContext(private val session: FirSession, val serializationIsPresent: Boolean) {
+class FirCheckersContext(
+    private val session: FirSession,
+    val serializationIsPresent: Boolean,
+    modes: StrictModeAggregator,
+) {
+    val strictModeDiagnostics = FirRpcStrictModeDiagnostics(modes)
+
     val typeParametersCache = session.firCachesFactory.createCache { typeParameter: FirTypeParameterSymbol ->
         FirCheckedAnnotationHelper.checkedAnnotations(session, typeParameter)
     }

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

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.fir.extensions.FirSupertypeGenerationExtension.Facto
 class FirRpcExtensionRegistrar(private val configuration: CompilerConfiguration) : FirExtensionRegistrar() {
     override fun ExtensionRegistrarContext.configurePlugin() {
         val logger = configuration.get(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY, MessageCollector.NONE)
+        val modes = configuration.strictModeAggregator()
 
         val serializationIsPresent = try {
             Class.forName("org.jetbrains.kotlinx.serialization.compiler.fir.SerializationFirResolveExtension")
@@ -29,7 +30,7 @@ class FirRpcExtensionRegistrar(private val configuration: CompilerConfiguration)
             false
         }
 
-        +CFactory { FirRpcCheckers(it, serializationIsPresent) }
+        +CFactory { FirRpcAdditionalCheckers(it, serializationIsPresent, modes) }
 
         +SFactory { FirRpcSupertypeGenerator(it, logger) }
     }

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

@@ -187,9 +187,10 @@ class FirRpcServiceGenerator(
         owner: FirClassSymbol<*>,
         name: Name,
         rpcServiceStubKey: RpcGeneratedStubKey,
-    ): FirClassLikeSymbol<*> {
+    ): FirClassLikeSymbol<*>? {
         val methodName = name.rpcMethodName
-        val rpcMethod = rpcServiceStubKey.functions.single { it.name == methodName }
+        val rpcMethod = rpcServiceStubKey.functions.singleOrNull { it.name == methodName }
+            ?: return null
         val rpcMethodClassKey = RpcGeneratedRpcMethodClassKey(rpcMethod)
         val classKind = if (rpcMethodClassKey.isObject) ClassKind.OBJECT else ClassKind.CLASS
 

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

@@ -0,0 +1,142 @@
+/*
+ * 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 org.jetbrains.kotlin.compiler.plugin.AbstractCliOption
+import org.jetbrains.kotlin.compiler.plugin.CliOption
+import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.config.CompilerConfigurationKey
+import kotlin.text.lowercase
+
+enum class StrictMode {
+    NONE, WARNING, ERROR;
+
+    companion object {
+        fun fromCli(value: String): StrictMode? {
+            return when (value.lowercase()) {
+                "none" -> NONE
+                "warning" -> WARNING
+                "error" -> ERROR
+                else -> null
+            }
+        }
+    }
+}
+
+data class StrictModeAggregator(
+    val stateFlow: StrictMode,
+    val sharedFlow: StrictMode,
+    val nestedFlow: StrictMode,
+    val streamScopedFunctions: StrictMode,
+    val suspendingServerStreaming: StrictMode,
+    val notTopLevelServerFlow: StrictMode,
+    val fields: StrictMode,
+)
+
+object StrictModeConfigurationKeys {
+    val STATE_FLOW = CompilerConfigurationKey.create<StrictMode>("state flow rpc mode")
+    val SHARED_FLOW = CompilerConfigurationKey.create<StrictMode>("shared flow rpc mode")
+    val NESTED_FLOW = CompilerConfigurationKey.create<StrictMode>("nested flow rpc mode")
+    val STREAM_SCOPED_FUNCTIONS = CompilerConfigurationKey.create<StrictMode>("stream scoped rpc mode")
+    val SUSPENDING_SERVER_STREAMING = CompilerConfigurationKey.create<StrictMode>(
+        "suspending server streaming rpc mode"
+    )
+    val NOT_TOP_LEVEL_SERVER_FLOW = CompilerConfigurationKey.create<StrictMode>("not top level server flow rpc mode")
+    val FIELDS = CompilerConfigurationKey.create<StrictMode>("fields rpc mode")
+}
+
+fun CompilerConfiguration.strictModeAggregator(): StrictModeAggregator {
+    return StrictModeAggregator(
+        stateFlow = get(StrictModeConfigurationKeys.STATE_FLOW, StrictMode.WARNING),
+        sharedFlow = get(StrictModeConfigurationKeys.SHARED_FLOW, StrictMode.WARNING),
+        nestedFlow = get(StrictModeConfigurationKeys.NESTED_FLOW, StrictMode.WARNING),
+        streamScopedFunctions = get(StrictModeConfigurationKeys.STREAM_SCOPED_FUNCTIONS, StrictMode.NONE),
+        suspendingServerStreaming = get(StrictModeConfigurationKeys.SUSPENDING_SERVER_STREAMING, StrictMode.NONE),
+        notTopLevelServerFlow = get(StrictModeConfigurationKeys.NOT_TOP_LEVEL_SERVER_FLOW, StrictMode.WARNING),
+        fields = get(StrictModeConfigurationKeys.FIELDS, StrictMode.WARNING),
+    )
+}
+
+object StrictModeCliOptions {
+    val STATE_FLOW = CliOption(
+        optionName = "strict-stateFlow",
+        valueDescription = VALUE_DESCRIPTION,
+        description = description("StateFlow"),
+        required = false,
+        allowMultipleOccurrences = false,
+    )
+
+    val SHARED_FLOW = CliOption(
+        optionName = "strict-sharedFlow",
+        valueDescription = VALUE_DESCRIPTION,
+        description = description("SharedFlow"),
+        required = false,
+        allowMultipleOccurrences = false,
+    )
+
+    val NESTED_FLOW = CliOption(
+        optionName = "strict-nested-flow",
+        valueDescription = VALUE_DESCRIPTION,
+        description = description("Nested flows"),
+        required = false,
+        allowMultipleOccurrences = false,
+    )
+
+    val STREAM_SCOPED_FUNCTIONS = CliOption(
+        optionName = "strict-stream-scope",
+        valueDescription = VALUE_DESCRIPTION,
+        description = description("Stream Scopes"),
+        required = false,
+        allowMultipleOccurrences = false,
+    )
+
+    val SUSPENDING_SERVER_STREAMING = CliOption(
+        optionName = "strict-suspending-server-streaming",
+        valueDescription = VALUE_DESCRIPTION,
+        description = description("suspending server streaming methods"),
+        required = false,
+        allowMultipleOccurrences = false,
+    )
+
+    val NOT_TOP_LEVEL_SERVER_FLOW = CliOption(
+        optionName = "strict-not-top-level-server-flow",
+        valueDescription = VALUE_DESCRIPTION,
+        description = description("not top-level server streaming declarations"),
+        required = false,
+        allowMultipleOccurrences = false,
+    )
+
+    val FIELDS = CliOption(
+        optionName = "strict-fields",
+        valueDescription = VALUE_DESCRIPTION,
+        description = description("fields"),
+        required = false,
+        allowMultipleOccurrences = false,
+    )
+
+    const val VALUE_DESCRIPTION = "none, warning or error"
+
+    fun description(entity: String): String {
+        return "Diagnostic level for $entity in @Rpc services."
+    }
+
+    val configurationMapper = mapOf(
+        STATE_FLOW to StrictModeConfigurationKeys.STATE_FLOW,
+        SHARED_FLOW to StrictModeConfigurationKeys.SHARED_FLOW,
+        NESTED_FLOW to StrictModeConfigurationKeys.NESTED_FLOW,
+        STREAM_SCOPED_FUNCTIONS to StrictModeConfigurationKeys.STREAM_SCOPED_FUNCTIONS,
+        SUSPENDING_SERVER_STREAMING to StrictModeConfigurationKeys.SUSPENDING_SERVER_STREAMING,
+        NOT_TOP_LEVEL_SERVER_FLOW to StrictModeConfigurationKeys.NOT_TOP_LEVEL_SERVER_FLOW,
+        FIELDS to StrictModeConfigurationKeys.FIELDS,
+    )
+}
+
+fun AbstractCliOption.processAsStrictModeOption(value: String, configuration: CompilerConfiguration): Boolean {
+    val key = StrictModeCliOptions.configurationMapper[this] ?: return false
+    val mode = StrictMode.fromCli(value) ?: return false
+
+    configuration.put(key, mode)
+    return true
+}

compiler-plugin/compiler-plugin-k2/src/main/core/kotlinx/rpc/codegen/checkers/FirRpcDeclarationCheckers.kt → compiler-plugin/compiler-plugin-k2/src/main/core/kotlinx/rpc/codegen/checkers/FirRpcCheckers.kt

@@ -14,8 +14,10 @@ import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExpressionCheckers
 import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirFunctionCallChecker
 
 class FirRpcDeclarationCheckers(ctx: FirCheckersContext) : DeclarationCheckers() {
-    override val regularClassCheckers: Set<FirRegularClassChecker> = setOf(
+    override val regularClassCheckers: Set<FirRegularClassChecker> = setOfNotNull(
         FirRpcAnnotationChecker(ctx),
+        if (ctx.serializationIsPresent) FirRpcStrictModeClassChecker(ctx) else null,
+        FirRpcServiceDeclarationChecker(ctx),
     )
 
     override val classCheckers: Set<FirClassChecker> = setOf(
@@ -34,5 +36,6 @@ class FirRpcDeclarationCheckers(ctx: FirCheckersContext) : DeclarationCheckers()
 class FirRpcExpressionCheckers(ctx: FirCheckersContext) : ExpressionCheckers() {
     override val functionCallCheckers: Set<FirFunctionCallChecker> = setOf(
         FirCheckedAnnotationFunctionCallChecker(ctx),
+        FirRpcStrictModeExpressionChecker(ctx),
     )
 }