1 files changed, 4 insertions, 42 deletions

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 94d8abcf..133643ee 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
@@ -131,49 +131,21 @@ class KspGradleSubplugin @Inject internal constructor(private val registry: Tool
         }
     }
 
+    private lateinit var kspConfigurations: KspConfigurations
     private val androidIntegration by lazy {
         AndroidPluginIntegration(this)
     }
 
-    @OptIn(ExperimentalStdlibApi::class)
-    private val KotlinSourceSet.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 KotlinSourceSet.kspConfiguration(project: Project): Configuration? {
-        return project.configurations.findByName(kspConfigurationName)
-    }
-
     override fun apply(project: Project) {
         project.extensions.create("ksp", KspExtension::class.java)
-        // Always include the main `ksp` configuration.
-        // TODO: multiplatform
-        project.configurations.create(KSP_MAIN_CONFIGURATION_NAME)
-        project.plugins.withType(KotlinPluginWrapper::class.java) {
-            // kotlin extension has the compilation target that we need to look for to create configurations
-            decorateKotlinExtension(project)
-        }
+        kspConfigurations = KspConfigurations(project)
         androidIntegration.applyIfAndroidProject(project)
         registry.register(KspModelBuilder())
     }
 
-    private fun decorateKotlinExtension(project: Project) {
-        project.extensions.configure(KotlinSingleTargetExtension::class.java) { kotlinExtension ->
-            kotlinExtension.target.compilations.createKspConfigurations(project) { kotlinCompilation ->
-                kotlinCompilation.kotlinSourceSets.map {
-                    it.kspConfigurationName
-                }
-            }
-        }
-    }
-
     /**
      * Creates a KSP configuration for each element in the object container.
+     * TODO: remove after revisiting AndroidPluginIntegration.
      */
     internal fun <T> NamedDomainObjectContainer<T>.createKspConfigurations(
         project: Project,
@@ -216,17 +188,7 @@ class KspGradleSubplugin @Inject internal constructor(private val registry: Tool
             project.locateTask(kotlinCompilation.compileKotlinTaskName) ?: return project.provider { emptyList() }
         val javaCompile = findJavaTaskForKotlinCompilation(kotlinCompilation)?.get()
         val kspExtension = project.extensions.getByType(KspExtension::class.java)
-        val kspConfigurations = LinkedHashSet<Configuration>()
-        kotlinCompilation.allKotlinSourceSets.forEach {
-            it.kspConfiguration(project)?.let {
-                kspConfigurations.add(it)
-            }
-        }
-        // Always include the main `ksp` configuration.
-        // TODO: multiplatform
-        project.configurations.findByName(KSP_MAIN_CONFIGURATION_NAME)?.let {
-            kspConfigurations.add(it)
-        }
+        val kspConfigurations = kspConfigurations.find(kotlinCompilation)
         val nonEmptyKspConfigurations = kspConfigurations.filter { it.dependencies.isNotEmpty() }
         if (nonEmptyKspConfigurations.isEmpty()) {
             return project.provider { emptyList() }