From 8ae8f630e905b6a72ef628cefbbca15f612e8e46 Mon Sep 17 00:00:00 2001
From: Ting-Yuan Huang <laszio@google.com>
Date: Tue, 5 Nov 2024 14:33:01 -0800
Subject: [PATCH] KSP2: fix module names for Android builds

by mimicking the rules from KGP.

(cherry picked from commit 3fac0c1a917acc9e52b35eb35fd32986544828ef)
---
 .../kotlin/com/google/devtools/ksp/gradle/KspAATask.kt   | 7 ++++++-
 .../kotlin/com/google/devtools/ksp/test/AndroidIT.kt     | 3 +++
 .../test-processor/src/main/kotlin/TestProcessor.kt      | 9 +++++++++
 .../workload/src/main/java/com/example/AClass.kt         | 2 ++
 4 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspAATask.kt b/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspAATask.kt
index 17d8cabfff..2cec471e17 100644
--- a/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspAATask.kt
+++ b/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspAATask.kt
@@ -154,7 +154,12 @@ abstract class KspAATask @Inject constructor(
                 kspAATask.kspClasspath.from(kspAADepCfg)
                 kspAATask.kspConfig.let { cfg ->
                     cfg.processorClasspath.from(processorClasspath)
-                    cfg.moduleName.value(project.name)
+                    // Ref: https://github.com/JetBrains/kotlin/blob/6535f86dfe36effeba976802ebd56a5a56071f45/libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/mpp/kotlinCompilations.kt#L92
+                    val moduleName = when (val compilationName = kotlinCompilation.name) {
+                        KotlinCompilation.MAIN_COMPILATION_NAME -> project.name
+                        else -> "${project.name}_$compilationName"
+                    }
+                    cfg.moduleName.value(moduleName)
                     val kotlinOutputDir = KspGradleSubplugin.getKspKotlinOutputDir(project, sourceSetName, target)
                     val javaOutputDir = KspGradleSubplugin.getKspJavaOutputDir(project, sourceSetName, target)
                     val filteredTasks =
diff --git a/integration-tests/src/test/kotlin/com/google/devtools/ksp/test/AndroidIT.kt b/integration-tests/src/test/kotlin/com/google/devtools/ksp/test/AndroidIT.kt
index 099f4bd921..0aed320ef5 100644
--- a/integration-tests/src/test/kotlin/com/google/devtools/ksp/test/AndroidIT.kt
+++ b/integration-tests/src/test/kotlin/com/google/devtools/ksp/test/AndroidIT.kt
@@ -32,6 +32,9 @@ class AndroidIT(useKSP2: Boolean) {
             assert("-keep class com.example.AClassBuilder { *; }" in configurationText) {
                 "Merged configuration did not contain generated proguard rules!\n$configurationText"
             }
+            val outputs = result.output.lines()
+            assert("w: [ksp] [workload_debug] Mangled name for internalFun: internalFun\$workload_debug" in outputs)
+            assert("w: [ksp] [workload_release] Mangled name for internalFun: internalFun\$workload_release" in outputs)
         }
     }
 
diff --git a/integration-tests/src/test/resources/playground/test-processor/src/main/kotlin/TestProcessor.kt b/integration-tests/src/test/resources/playground/test-processor/src/main/kotlin/TestProcessor.kt
index 2025a1ef49..57d479822d 100644
--- a/integration-tests/src/test/resources/playground/test-processor/src/main/kotlin/TestProcessor.kt
+++ b/integration-tests/src/test/resources/playground/test-processor/src/main/kotlin/TestProcessor.kt
@@ -1,5 +1,6 @@
 import com.google.devtools.ksp.KspExperimental
 import com.google.devtools.ksp.containingFile
+import com.google.devtools.ksp.getClassDeclarationByName
 import com.google.devtools.ksp.processing.*
 import com.google.devtools.ksp.symbol.*
 import java.io.OutputStream
@@ -49,6 +50,14 @@ class TestProcessor : SymbolProcessor {
             emit("TestProcessor: processing ${file.fileName}", "")
             file.accept(visitor, "")
         }
+
+        resolver.getClassDeclarationByName("com.example.AClass")?.let { aClass ->
+            aClass.declarations.single { it.simpleName.asString() == "internalFun" }.let { internalFun ->
+                val internalName = resolver.getJvmName(internalFun as KSFunctionDeclaration)
+                val moduleName = resolver.getModuleName().asString()
+                logger.warn("[$moduleName] Mangled name for internalFun: $internalName")
+            }
+        }
         invoked = true
         return emptyList()
     }
diff --git a/integration-tests/src/test/resources/playground/workload/src/main/java/com/example/AClass.kt b/integration-tests/src/test/resources/playground/workload/src/main/java/com/example/AClass.kt
index 180dbe6a2a..446e78dfda 100644
--- a/integration-tests/src/test/resources/playground/workload/src/main/java/com/example/AClass.kt
+++ b/integration-tests/src/test/resources/playground/workload/src/main/java/com/example/AClass.kt
@@ -14,4 +14,6 @@ class AClass(private val a: Int, val b: String, val c: Double, val d: HELLO) {
     class innerClass<T : HELLO>
 
     val generic = innerClass<HELLO>()
+
+    internal fun internalFun(): Int = 0
 }