Handle Kotlin source sets from Kotlin compilations more lazily

(cherry picked from commit 53065a8)
google · Jul 31, 2023 · d5f9df4 · d5f9df4
1 parent b106ee2
commit d5f9df4
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 25 deletions.
diff --git a/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KgpUtils.kt b/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KgpUtils.kt
@@ -0,0 +1,11 @@
+package com.google.devtools.ksp.gradle
+
+import org.jetbrains.kotlin.gradle.plugin.KotlinCompilation
+import org.jetbrains.kotlin.gradle.plugin.KotlinSourceSet
+import org.jetbrains.kotlin.gradle.utils.ObservableSet
+
+internal val KotlinCompilation<*>.allKotlinSourceSetsObservable
+    get() = this.allKotlinSourceSets as ObservableSet<KotlinSourceSet>
+
+internal val KotlinCompilation<*>.kotlinSourceSetsObservable
+    get() = this.kotlinSourceSets as ObservableSet<KotlinSourceSet>
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
@@ -74,9 +74,9 @@ abstract class KspAATask @Inject constructor(
                 kspAATask.kspConfig.let { cfg ->
                     cfg.processorClasspath.from(processorClasspath)
                     cfg.moduleName.value(kotlinCompilation.defaultSourceSet.name)
-                    kotlinCompilation.allKotlinSourceSets
-                        .filterNot { it == kspGeneratedSourceSet }
-                        .forEach { sourceSet ->
+                    kotlinCompilation.allKotlinSourceSetsObservable
+                        .forAll { sourceSet ->
+                            if (sourceSet == kspGeneratedSourceSet) return@forAll
                             cfg.sourceRoots.from(sourceSet.kotlin)
                             cfg.javaSourceRoots.from(sourceSet.kotlin)
                         }

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
@@ -133,7 +133,7 @@ class KspConfigurations(private val project: Project) {
             }
         } else {
             target.compilations.configureEach { compilation ->
-                compilation.kotlinSourceSets.forEach { sourceSet ->
+                compilation.kotlinSourceSetsObservable.forAll { sourceSet ->
                     createConfiguration(
                         name = getKotlinConfigurationName(compilation, sourceSet),
                         readableSetName = sourceSet.name

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
@@ -268,34 +268,36 @@ class KspGradleSubplugin @Inject internal constructor(private val registry: Tool
 
         val processorClasspath = project.configurations.maybeCreate("${kspTaskName}ProcessorClasspath")
             .extendsFrom(*nonEmptyKspConfigurations.toTypedArray()).markResolvable()
+
         fun configureAsKspTask(kspTask: KspTask, isIncremental: Boolean) {
             // depends on the processor; if the processor changes, it needs to be reprocessed.
             kspTask.dependsOn(processorClasspath.buildDependencies)
             kspTask.commandLineArgumentProviders.addAll(kspExtension.commandLineArgumentProviders)
 
-            val commonSources: List<File> = when (processingModel) {
-                "hierarchical" -> {
-                    fun unclaimedDeps(roots: Set<KotlinSourceSet>): Set<KotlinSourceSet> {
-                        val unclaimedParents =
-                            roots.flatMap { it.dependsOn }.filterNot { it in sourceSetMap }.toSet()
-                        return if (unclaimedParents.isEmpty()) {
-                            unclaimedParents
-                        } else {
-                            unclaimedParents + unclaimedDeps(unclaimedParents)
+            kspTask.options.addAll(
+                kspTask.project.provider {
+                    val commonSources: List<File> = when (processingModel) {
+                        "hierarchical" -> {
+                            fun unclaimedDeps(roots: Set<KotlinSourceSet>): Set<KotlinSourceSet> {
+                                val unclaimedParents =
+                                    roots.flatMap { it.dependsOn }.filterNot { it in sourceSetMap }.toSet()
+                                return if (unclaimedParents.isEmpty()) {
+                                    unclaimedParents
+                                } else {
+                                    unclaimedParents + unclaimedDeps(unclaimedParents)
+                                }
+                            }
+                            // Source sets that are not claimed by other compilations.
+                            // I.e., those that should be processed by this compilation.
+                            val unclaimed =
+                                kotlinCompilation.kotlinSourceSets + unclaimedDeps(kotlinCompilation.kotlinSourceSets)
+                            val commonSourceSets = kotlinCompilation.allKotlinSourceSets - unclaimed
+                            commonSourceSets.flatMap { it.kotlin.files }
                         }
+
+                        else -> emptyList()
                     }
-                    // Source sets that are not claimed by other compilations.
-                    // I.e., those that should be processed by this compilation.
-                    val unclaimed =
-                        kotlinCompilation.kotlinSourceSets + unclaimedDeps(kotlinCompilation.kotlinSourceSets)
-                    val commonSourceSets = kotlinCompilation.allKotlinSourceSets - unclaimed
-                    commonSourceSets.flatMap { it.kotlin.files }
-                }
-                else -> emptyList()
-            }
 
-            kspTask.options.addAll(
-                kspTask.project.provider {
                     getSubpluginOptions(
                         project,
                         kspExtension,
@@ -341,9 +343,11 @@ class KspGradleSubplugin @Inject internal constructor(private val registry: Tool
                     setSource(kotlinCompileTask.javaSources - kspGeneratedSourceSet.kotlin)
                 }
             } else {
-                kotlinCompilation.allKotlinSourceSets.filterNot { it == kspGeneratedSourceSet }.forEach { sourceSet ->
+                kotlinCompilation.allKotlinSourceSetsObservable.forAll { sourceSet ->
+                    if (sourceSet == kspGeneratedSourceSet) return@forAll
                     kspTask.setSource(sourceSet.kotlin)
                 }
+
                 if (kotlinCompilation is KotlinCommonCompilation) {
                     kspTask.setSource(kotlinCompilation.defaultSourceSet.kotlin)
                 }
@@ -356,6 +360,7 @@ class KspGradleSubplugin @Inject internal constructor(private val registry: Tool
                         }
                         kotlinCompilation.kotlinSourceSets.flatMap { claimedParents(it) }.map { sourceSetMap[it]!! }
                     }
+
                     else -> emptyList()
                 }
                 generated.forEach {
@@ -460,6 +465,7 @@ class KspGradleSubplugin @Inject internal constructor(private val registry: Tool
                     }
                 }
             }
+
             KotlinPlatformType.js, KotlinPlatformType.wasm -> {
                 KotlinFactories.registerKotlinJSCompileTask(project, kspTaskName, kotlinCompilation).also {
                     it.configure { kspTask ->
@@ -483,6 +489,7 @@ class KspGradleSubplugin @Inject internal constructor(private val registry: Tool
                     }
                 }
             }
+
             KotlinPlatformType.common -> {
                 KotlinFactories.registerKotlinMetadataCompileTask(project, kspTaskName, kotlinCompilation).also {
                     it.configure { kspTask ->
@@ -506,6 +513,7 @@ class KspGradleSubplugin @Inject internal constructor(private val registry: Tool
                     }
                 }
             }
+
             KotlinPlatformType.native -> {
                 KotlinFactories.registerKotlinNativeCompileTask(project, kspTaskName, kotlinCompilation).also {
                     it.configure { kspTask ->