KRPC-162 Internal vs Public Codegen separation in gRPC (#285)

Mr3zee · Mr3zee · commit 277b4b7eb86e · 2025-08-01T16:12:31.000+02:00
diff --git a/protobuf-plugin/src/main/kotlin/kotlinx/rpc/protobuf/CodeGenerator.kt b/protobuf-plugin/src/main/kotlin/kotlinx/rpc/protobuf/CodeGenerator.kt
@@ -255,13 +255,14 @@ class FileGenerator(
     codeGenerationParameters: CodeGenerationParameters,
     var filename: String? = null,
     var packageName: String? = null,
+    var packagePath: String? = packageName,
     var fileOptIns: List<String> = emptyList(),
     logger: Logger = NOPLogger.NOP_LOGGER,
 ) : CodeGenerator(codeGenerationParameters, "", logger = logger) {
     private val imports = mutableListOf<String>()
 
     fun importPackage(name: String) {
-        if (name != packageName) {
+        if (name != packageName && name.isNotBlank()) {
             imports.add("$name.*")
         }
     }
@@ -304,4 +305,5 @@ fun file(
     packageName: String? = null,
     logger: Logger = NOPLogger.NOP_LOGGER,
     block: FileGenerator.() -> Unit,
-): FileGenerator = FileGenerator(codeGenerationParameters, name, packageName, emptyList(), logger).apply(block)
+): FileGenerator = FileGenerator(codeGenerationParameters, name, packageName, packageName, emptyList(), logger)
+    .apply(block)
diff --git a/protobuf-plugin/src/main/kotlin/kotlinx/rpc/protobuf/ModelToKotlinGenerator.kt b/protobuf-plugin/src/main/kotlin/kotlinx/rpc/protobuf/ModelToKotlinGenerator.kt
@@ -8,74 +8,89 @@ import kotlinx.rpc.protobuf.CodeGenerator.DeclarationType
 import kotlinx.rpc.protobuf.model.*
 import org.slf4j.Logger
 
+private const val RPC_INTERNAL_PACKAGE_SUFFIX = "_rpc_internal"
+
 class ModelToKotlinGenerator(
     private val model: Model,
     private val logger: Logger,
     private val codeGenerationParameters: CodeGenerationParameters,
 ) {
     fun generateKotlinFiles(): List<FileGenerator> {
-        return model.files.map { it.generateKotlinFile() }
+        return model.files.flatMap { it.generateKotlinFiles() }
+    }
+
+    private fun FileDeclaration.generateKotlinFiles(): List<FileGenerator> {
+        return listOf(
+            generatePublicKotlinFile(),
+            generateInternalKotlinFile(),
+        )
     }
 
-    private fun FileDeclaration.generateKotlinFile(): FileGenerator {
+    private fun FileDeclaration.generatePublicKotlinFile(): FileGenerator {
         return file(codeGenerationParameters, logger = logger) {
             filename = name.simpleName
             packageName = name.packageName
-            fileOptIns = listOf("ExperimentalRpcApi::class", "InternalRpcApi::class")
+            packagePath = name.packageName
 
             dependencies.forEach { dependency ->
                 importPackage(dependency.name.packageName)
             }
 
-            generateDeclaredEntities(this@generateKotlinFile)
+            generatePublicDeclaredEntities(this@generatePublicKotlinFile)
 
-            additionalImports.forEach {
+            additionalPublicImports.forEach {
                 import(it)
             }
+        }
+    }
+
+    private fun FileDeclaration.generateInternalKotlinFile(): FileGenerator {
+        return file(codeGenerationParameters, logger = logger) {
+            filename = name.simpleName
+            packageName = name.packageName
+            packagePath = name.packageName.packageNameSuffixed(RPC_INTERNAL_PACKAGE_SUFFIX)
+
+            fileOptIns = listOf("ExperimentalRpcApi::class", "InternalRpcApi::class")
+
+            dependencies.forEach { dependency ->
+                importPackage(dependency.name.packageName)
+            }
+
+            generateInternalDeclaredEntities(this@generateInternalKotlinFile)
+
             import("kotlinx.rpc.internal.utils.*")
         }
     }
 
-    private val additionalImports = mutableSetOf<String>()
+    private val additionalPublicImports = mutableSetOf<String>()
 
-    private fun CodeGenerator.generateDeclaredEntities(fileDeclaration: FileDeclaration) {
-        fileDeclaration.messageDeclarations.forEach { generateMessage(it) }
+    private fun CodeGenerator.generatePublicDeclaredEntities(fileDeclaration: FileDeclaration) {
+        fileDeclaration.messageDeclarations.forEach { generatePublicMessage(it) }
         // KRPC-141 Enum Types
 //        fileDeclaration.enumDeclarations.forEach { generateEnum(it) }
-        fileDeclaration.serviceDeclarations.forEach { generateService(it) }
+        fileDeclaration.serviceDeclarations.forEach { generatePublicService(it) }
     }
 
-    @Suppress("detekt.CyclomaticComplexMethod")
-    private fun CodeGenerator.generateMessage(declaration: MessageDeclaration) {
-        val fields = declaration.actualFields.map { it.generateFieldDeclaration() to it.type.defaultValue }
-
-        val isInterfaceMode = parameters.messageMode == RpcProtobufPlugin.MessageMode.Interface
-
-        val (declarationType, modifiers) = when {
-            isInterfaceMode -> {
-                DeclarationType.Interface to ""
-            }
-
-            fields.isEmpty() -> {
-                DeclarationType.Object to ""
-            }
+    private fun CodeGenerator.generateInternalDeclaredEntities(fileDeclaration: FileDeclaration) {
+        fileDeclaration.messageDeclarations.forEach { generateInternalMessage(it) }
+        // KRPC-141 Enum Types
+//        fileDeclaration.enumDeclarations.forEach { generateEnum(it) }
+        fileDeclaration.serviceDeclarations.forEach { generateInternalService(it) }
+    }
 
-            else -> {
-                DeclarationType.Class to "data"
-            }
-        }
+    private fun MessageDeclaration.fields() = actualFields.map {
+        it.transformToFieldDeclaration() to it.type.defaultValue
+    }
 
+    @Suppress("detekt.CyclomaticComplexMethod")
+    private fun CodeGenerator.generatePublicMessage(declaration: MessageDeclaration) {
         clazz(
             name = declaration.name.simpleName,
-            modifiers = modifiers,
-            constructorArgs = if (isInterfaceMode) emptyList() else fields.map { "val ${it.first}" to it.second },
-            declarationType = declarationType,
+            declarationType = DeclarationType.Interface,
         ) {
-            if (isInterfaceMode) {
-                fields.forEach {
-                    code("val ${it.first}")
-                    newLine()
-                }
+            declaration.fields().forEach {
+                code("val ${it.first}")
+                newLine()
             }
 
             // KRPC-147 OneOf Types
@@ -93,32 +108,31 @@ class ModelToKotlinGenerator(
 //                generateEnum(enum)
 //            }
 
-            if (isInterfaceMode) {
-                clazz("", modifiers = "companion", declarationType = DeclarationType.Object)
-            }
+            clazz("", modifiers = "companion", declarationType = DeclarationType.Object)
         }
+    }
 
-        if (isInterfaceMode) {
-            clazz(
-                name = "${declaration.name.simpleName}Builder",
-                declarationType = DeclarationType.Class,
-                superTypes = listOf(declaration.name.simpleName),
-            ) {
-                fields.forEach {
-                    code("override var ${it.first} = ${it.second}")
-                    newLine()
-                }
+    @Suppress("detekt.CyclomaticComplexMethod")
+    private fun CodeGenerator.generateInternalMessage(declaration: MessageDeclaration) {
+        clazz(
+            name = "${declaration.name.simpleName}Builder",
+            declarationType = DeclarationType.Class,
+            superTypes = listOf(declaration.name.simpleName),
+        ) {
+            declaration.fields().forEach {
+                code("override var ${it.first} = ${it.second}")
+                newLine()
             }
+        }
 
-            function(
-                name = "invoke",
-                modifiers = "operator",
-                args = "body: ${declaration.name.simpleName}Builder.() -> Unit",
-                contextReceiver = "${declaration.name.simpleName}.Companion",
-                returnType = declaration.name.simpleName,
-            ) {
-                code("return ${declaration.name.simpleName}Builder().apply(body)")
-            }
+        function(
+            name = "invoke",
+            modifiers = "operator",
+            args = "body: ${declaration.name.simpleName}Builder.() -> Unit",
+            contextReceiver = "${declaration.name.simpleName}.Companion",
+            returnType = declaration.name.simpleName,
+        ) {
+            code("return ${declaration.name.simpleName}Builder().apply(body)")
         }
 
         val platformType = "${declaration.outerClassName.simpleName}.${declaration.name.simpleName}"
@@ -175,7 +189,7 @@ class ModelToKotlinGenerator(
         }
     }
 
-    private fun FieldDeclaration.generateFieldDeclaration(): String {
+    private fun FieldDeclaration.transformToFieldDeclaration(): String {
         return "${name}: ${typeFqName()}"
     }
 
@@ -219,7 +233,7 @@ class ModelToKotlinGenerator(
                     superTypes = listOf(interfaceName),
                 )
 
-                additionalImports.add("kotlin.jvm.JvmInline")
+                additionalPublicImports.add("kotlin.jvm.JvmInline")
             }
         }
     }
@@ -247,7 +261,7 @@ class ModelToKotlinGenerator(
     }
 
     @Suppress("detekt.LongMethod")
-    private fun CodeGenerator.generateService(service: ServiceDeclaration) {
+    private fun CodeGenerator.generatePublicService(service: ServiceDeclaration) {
         code("@kotlinx.rpc.grpc.annotations.Grpc")
         clazz(service.name.simpleName, declarationType = DeclarationType.Interface) {
             service.methods.forEach { method ->
@@ -262,9 +276,9 @@ class ModelToKotlinGenerator(
                 )
             }
         }
+    }
 
-        newLine()
-
+    private fun CodeGenerator.generateInternalService(service: ServiceDeclaration) {
         code("@Suppress(\"unused\", \"all\")")
         clazz(
             modifiers = "private",
@@ -373,3 +387,7 @@ class ModelToKotlinGenerator(
         return "${outerClassName.simpleName}.${name.simpleName.removePrefix(name.parentNameAsPrefix)}"
     }
 }
+
+internal fun String.packageNameSuffixed(suffix: String): String {
+    return if (isEmpty()) suffix else "$this.$suffix"
+}
diff --git a/protobuf-plugin/src/main/kotlin/kotlinx/rpc/protobuf/RpcProtobufPlugin.kt b/protobuf-plugin/src/main/kotlin/kotlinx/rpc/protobuf/RpcProtobufPlugin.kt
@@ -14,6 +14,7 @@ import com.google.protobuf.compiler.PluginProtos.CodeGeneratorResponse.Feature
 import org.slf4j.Logger
 import org.slf4j.LoggerFactory
 import org.slf4j.helpers.NOPLogger
+import java.io.File
 
 class RpcProtobufPlugin {
     companion object {
@@ -68,7 +69,9 @@ class RpcProtobufPlugin {
             .map { file ->
                 CodeGeneratorResponse.File.newBuilder()
                     .apply {
-                        val dir = file.packageName?.replace('.', '/')?.plus("/") ?: ""
+                        val dir = file.packagePath
+                            ?.replace('.', File.separatorChar)?.plus(File.separatorChar)
+                            ?: ""
 
                         // some filename already contain package (true for Google's default .proto files)
                         val filename = file.filename?.removePrefix(dir) ?: error("File name can not be null")
