Add Android/AGP multiplatform support

3 files changed, 72 insertions, 53 deletions

diff --git a/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/AndroidPluginIntegration.kt b/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/AndroidPluginIntegration.kt
index bfe98483..8bbfda29 100644
--- a/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/AndroidPluginIntegration.kt
+++ b/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/AndroidPluginIntegration.kt
@@ -16,17 +16,15 @@
  */
 package com.google.devtools.ksp.gradle
 
-import com.android.build.api.dsl.AndroidSourceSet
 import com.android.build.api.dsl.CommonExtension
 import com.android.build.gradle.BaseExtension
-import com.google.devtools.ksp.gradle.KspGradleSubplugin.Companion.KSP_MAIN_CONFIGURATION_NAME
+import com.android.build.gradle.api.AndroidSourceSet
+import org.gradle.api.Named
 import org.gradle.api.Project
 import org.gradle.api.file.FileCollection
-import org.gradle.api.tasks.SourceSet
 import org.gradle.api.tasks.TaskProvider
 import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinJvmAndroidCompilation
 import java.io.File
-import java.util.Locale
 
 /**
  * This helper class handles communication with the android plugin.
@@ -36,46 +34,37 @@ import java.util.Locale
  * plugin APIs directly without checking its existence (we have tests covering that case).
  */
 @Suppress("UnstableApiUsage") // some android APIs are unsable.
-class AndroidPluginIntegration(
-    private val kspGradleSubplugin: KspGradleSubplugin
-) {
+object AndroidPluginIntegration {
 
     private val agpPluginIds = listOf("com.android.application", "com.android.library", "com.android.dynamic-feature")
 
-    fun applyIfAndroidProject(project: Project) {
+    fun findSourceSets(project: Project, onSourceSet: (String) -> Unit) {
         agpPluginIds.forEach { agpPluginId ->
             project.pluginManager.withPlugin(agpPluginId) {
                 // for android apps, we need a configuration per source set
-                decorateAndroidExtension(project)
+                decorateAndroidExtension(project, onSourceSet)
             }
         }
     }
 
-    @OptIn(ExperimentalStdlibApi::class)
-    private val AndroidSourceSet.kspConfigurationName: String
-        get() {
-            return if (name == SourceSet.MAIN_SOURCE_SET_NAME) {
-                KSP_MAIN_CONFIGURATION_NAME
-            } else {
-                "$KSP_MAIN_CONFIGURATION_NAME${name.capitalize(Locale.US)}"
-            }
-        }
-
-    private fun decorateAndroidExtension(project: Project) {
+    private fun decorateAndroidExtension(project: Project, onSourceSet: (String) -> Unit) {
         val sourceSets = when (val androidExt = project.extensions.getByName("android")) {
             is BaseExtension -> androidExt.sourceSets
             is CommonExtension<*, *, *, *> -> androidExt.sourceSets
             else -> throw RuntimeException("Unsupported Android Gradle plugin version.")
         }
-
-        @Suppress("UnstableApiUsage")
-        kspGradleSubplugin.run {
-            sourceSets.createKspConfigurations(project) { androidSourceSet ->
-                listOf(androidSourceSet.kspConfigurationName)
-            }
+        sourceSets.configureEach {
+            onSourceSet(it.name)
         }
     }
 
+    fun getCompilationSourceSets(kotlinCompilation: KotlinJvmAndroidCompilation): List<String> {
+        return kotlinCompilation.androidVariant
+            .sourceSets
+            .filterIsInstance(AndroidSourceSet::class.java)
+            .map { it.name }
+    }
+
     fun registerGeneratedJavaSources(
         project: Project,
         kotlinCompilation: KotlinJvmAndroidCompilation,
diff --git a/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspConfigurations.kt b/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspConfigurations.kt
index 8932056c..45b23d6a 100644
--- a/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspConfigurations.kt
+++ b/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspConfigurations.kt
@@ -4,7 +4,7 @@ import org.gradle.api.Project
 import org.gradle.api.artifacts.Configuration
 import org.jetbrains.kotlin.gradle.dsl.*
 import org.jetbrains.kotlin.gradle.plugin.*
-import java.util.*
+import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinJvmAndroidCompilation
 
 /**
  * Creates and retrieves ksp-related configurations.
@@ -19,10 +19,16 @@ class KspConfigurations(private val project: Project) {
     private val rootMainConfiguration = project.configurations.create(ROOT)
 
     // Stores all saved configurations for quick access.
-    private val configurations = mutableMapOf<KotlinSourceSet, Configuration>()
+    private val kotlinConfigurations = mutableMapOf<KotlinSourceSet, Configuration>()
+    private val androidConfigurations = mutableMapOf<String, Configuration>()
 
     @OptIn(ExperimentalStdlibApi::class)
-    private fun addConfiguration(owner: KotlinSourceSet, parent: Configuration?, name: String): Configuration {
+    private fun <T: Any> saveConfiguration(
+        owner: T,
+        parent: Configuration?,
+        name: String,
+        cache: MutableMap<T, Configuration>
+    ): Configuration {
         val configName = ROOT + name.replaceFirstChar { it.uppercase() }
         val existingConfig = project.configurations.findByName(configName)
         if (existingConfig != null && configName != ROOT) {
@@ -30,11 +36,19 @@ class KspConfigurations(private val project: Project) {
         }
 
         val config = existingConfig ?: project.configurations.create(configName)
-        if (parent != null) config.extendsFrom(parent)
-        configurations[owner] = config
+        if (parent != null && parent.name != configName) {
+            config.extendsFrom(parent)
+        }
+        cache[owner] = config
         return config
     }
 
+    private fun saveKotlinConfiguration(owner: KotlinSourceSet, parent: Configuration?, name: String) =
+        saveConfiguration(owner, parent, name, kotlinConfigurations)
+
+    private fun saveAndroidConfiguration(owner: String, parent: Configuration?, name: String) =
+        saveConfiguration(owner, parent, name, androidConfigurations)
+
     init {
         project.plugins.withType(KotlinBasePluginWrapper::class.java).configureEach {
             // 1.6.0: decorateKotlinProject(project.kotlinExtension)
@@ -49,18 +63,33 @@ class KspConfigurations(private val project: Project) {
         }
     }
 
+    /**
+     * Decorate the source sets belonging to [target].
+     * The end goal is to have one KSP configuration per source set. Examples:
+     * - in a kotlin-jvm project, we want "ksp" (applied to main set) and "kspTest" (applied to test set)
+     * - in a kotlin-multiplatform project, we want "ksp<Target>" and "ksp<Target>Test" for each target.
+     * This is done by reading [KotlinCompilation.kotlinSourceSets], which contains appropriately named sets.
+     *
+     * For Android, we prefer to use AndroidSourceSets from AGP rather than [KotlinSourceSet]s like all other
+     * targets. There are very slight differences between the two - this could be re-evaluated in the future,
+     * because Kotlin Plugin does already create [KotlinSourceSet]s out of AndroidSourceSets
+     * ( https://kotlinlang.org/docs/mpp-configure-compilations.html#compilation-of-the-source-set-hierarchy ).
+     */
     private fun decorateKotlinTarget(target: KotlinTarget) {
         if (target.platformType == KotlinPlatformType.androidJvm) {
-            /**
-             * TODO: Android might need special handling. Discuss. Tricky points:
-             * 1) KotlinSourceSets are defined in terms of AGP Variants - a resolved, compilable entity.
-             *    Using them would be consistent with other targets and simple.
-             * 2) AGP AndroidSourceSets represent a hierarchy: we have "test", "debug", but also "testDebug"
-             *    which depends on the other two. Not clear if this dependency should be reflected in the
-             *    configurations.
-             * 3) Need to find a way to retrieve the correct configurations in applyToCompilation
-             */
-            Unit
+            AndroidPluginIntegration.findSourceSets(target.project) { setName ->
+                val isMain = setName.endsWith("main", ignoreCase = true)
+                val nameWithoutMain = when {
+                    isMain -> setName.substring(0, setName.length - 4)
+                    else -> setName
+                }
+                val nameWithTargetPrefix = when {
+                    target.name.isEmpty() -> nameWithoutMain
+                    else -> target.name + nameWithoutMain.replaceFirstChar { it.uppercase() }
+                }
+                val parent = if (isMain) rootMainConfiguration else null
+                saveAndroidConfiguration(setName, parent, nameWithTargetPrefix)
+            }
         } else {
             // We could add target-specific configurations here (kspJvm, parent of kspJvmMain & kspJvmTest)
             // but we decided that kspJvm should actually mean kspJvmMain, which in turn is not created.
@@ -84,9 +113,9 @@ class KspConfigurations(private val project: Project) {
         if (isMainCompilation && isDefaultSourceSet) {
             // Use target name instead of sourceSet name, to avoid creating "kspMain" or "kspJvmMain".
             // Note: on single-platform, target name is conveniently set to "" so this resolves to "ksp".
-            addConfiguration(sourceSet, parent, compilation.target.name)
+            saveKotlinConfiguration(sourceSet, parent, compilation.target.name)
         } else {
-            addConfiguration(sourceSet, parent, sourceSet.name)
+            saveKotlinConfiguration(sourceSet, parent, sourceSet.name)
         }
     }
 
@@ -100,7 +129,12 @@ class KspConfigurations(private val project: Project) {
      *    to share code between targets. They do not currently have their own ksp configuration.
      */
     fun find(compilation: KotlinCompilation<*>): Set<Configuration> {
-        val sourceSets = compilation.kotlinSourceSets
-        return sourceSets.mapNotNull { configurations[it] }.toSet()
+        val kotlinSourceSets = compilation.kotlinSourceSets
+        val kotlinConfigurations = kotlinSourceSets.mapNotNull { kotlinConfigurations[it] }
+        val androidConfigurations = if (compilation.platformType == KotlinPlatformType.androidJvm) {
+            val androidSourceSets = AndroidPluginIntegration.getCompilationSourceSets(compilation as KotlinJvmAndroidCompilation)
+            androidSourceSets.mapNotNull { androidConfigurations[it] }
+        } else emptyList()
+        return (kotlinConfigurations + androidConfigurations).toSet()
     }
 }
diff --git a/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspSubplugin.kt b/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspSubplugin.kt
index 133643ee..613196c1 100644
--- a/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspSubplugin.kt
+++ b/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspSubplugin.kt
@@ -132,14 +132,10 @@ class KspGradleSubplugin @Inject internal constructor(private val registry: Tool
     }
 
     private lateinit var kspConfigurations: KspConfigurations
-    private val androidIntegration by lazy {
-        AndroidPluginIntegration(this)
-    }
 
-    override fun apply(project: Project) {
-        project.extensions.create("ksp", KspExtension::class.java)
-        kspConfigurations = KspConfigurations(project)
-        androidIntegration.applyIfAndroidProject(project)
+    override fun apply(target: Project) {
+        target.extensions.create("ksp", KspExtension::class.java)
+        kspConfigurations = KspConfigurations(target)
         registry.register(KspModelBuilder())
     }
 
@@ -342,7 +338,7 @@ class KspGradleSubplugin @Inject internal constructor(private val registry: Tool
             }
         }
         if (kotlinCompilation is KotlinJvmAndroidCompilation) {
-            androidIntegration.registerGeneratedJavaSources(
+            AndroidPluginIntegration.registerGeneratedJavaSources(
                 project = project,
                 kotlinCompilation = kotlinCompilation,
                 kspTaskProvider = kspTaskProvider as TaskProvider<KspTaskJvm>,