Tracks and match source sets for MPP after the plugins are applied

Since the custom source sets may be added after the plugins are applied in both common and platform module, we need to handle `.all { ... }` source sets in both modules and add common sources to the platform module once a match is found.

Issue #KT-22510 Fixed

libraries/tools/kotlin-gradle-plugin-integration-tests/src/test/kotlin/org/jetbrains/kotlin/gradle/MultiplatformGradleIT.kt

@@ -256,4 +256,35 @@ class MultiplatformGradleIT : BaseGradleIT() {
             assertContains("$successMarker :lib:compileTestJava")
         }
     }
+
+    @Test
+    fun testCustomSourceSets() = with(Project("multiplatformProject")) {
+        setupWorkingDir()
+
+        val sourceSetName = "foo"
+        val sourceSetDeclaration = "\nsourceSets { $sourceSetName { } }"
+
+        listOf("lib", "libJvm", "libJs").forEach { module ->
+            gradleBuildScript(module).appendText(sourceSetDeclaration)
+        }
+
+        listOf(
+            "expect fun foo(): String" to "lib/src/$sourceSetName/kotlin",
+            "actual fun foo(): String = \"jvm\"" to "libJvm/src/$sourceSetName/kotlin",
+            "actual fun foo(): String = \"js\"" to "libJs/src/$sourceSetName/kotlin"
+        ).forEach { (code, path) ->
+            File(projectDir, path).run {
+                mkdirs();
+                File(this, "Foo.kt").writeText(code)
+            }
+        }
+
+        val customSourceSetCompileTasks = listOf(":lib" to "Common", ":libJs" to "2Js", ":libJvm" to "")
+            .map { (module, platform) -> "$module:compile${sourceSetName.capitalize()}Kotlin$platform" }
+
+        build(*customSourceSetCompileTasks.toTypedArray()) {
+            assertSuccessful()
+            assertTasksExecuted(customSourceSetCompileTasks)
+        }
+    }
 }

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/plugin/KotlinMultiplatformPlugin.kt

@@ -17,15 +17,12 @@
 package org.jetbrains.kotlin.gradle.plugin
 
 import com.android.build.gradle.BaseExtension
-import org.gradle.api.GradleException
-import org.gradle.api.Plugin
-import org.gradle.api.Project
+import org.gradle.api.*
 import org.gradle.api.artifacts.ProjectDependency
 import org.gradle.api.file.SourceDirectorySet
 import org.gradle.api.plugins.JavaPluginConvention
 import org.gradle.api.tasks.SourceSet
 import org.gradle.api.tasks.SourceSetContainer
-import org.jetbrains.kotlin.gradle.dsl.KotlinCompile
 import org.jetbrains.kotlin.gradle.tasks.AbstractKotlinCompile
 
 abstract class KotlinPlatformPluginBase(protected val platformName: String) : Plugin<Project> {
@@ -51,11 +48,8 @@ const val IMPLEMENT_DEPRECATION_WARNING = "The '$IMPLEMENT_CONFIG_NAME' configur
 
 open class KotlinPlatformImplementationPluginBase(platformName: String) : KotlinPlatformPluginBase(platformName) {
     private val commonProjects = arrayListOf<Project>()
-    private val platformKotlinTasksBySourceSetName = hashMapOf<String, AbstractKotlinCompile<*>>()
 
     override fun apply(project: Project) {
-        project.tasks.filterIsInstance<AbstractKotlinCompile<*>>().associateByTo(platformKotlinTasksBySourceSetName) { it.sourceSetName }
-
         val implementConfig = project.configurations.create(IMPLEMENT_CONFIG_NAME)
         val expectedByConfig = project.configurations.create(EXPECTED_BY_CONFIG_NAME)
 
@@ -107,16 +101,51 @@ open class KotlinPlatformImplementationPluginBase(platformName: String) : Kotlin
                         "dependency to non-common project $commonProject")
             }
 
-            commonProject.sourceSets.all { commonSourceSet ->
-                // todo: warn if not found
+            // Since the two projects may add source sets in arbitrary order, and both may do that after the plugin is applied,
+            // we need to handle all source sets of the two projects and connect them once we get a match:
+            // todo warn if no match found
+            matchSymmetricallyByNames(commonProject.sourceSets, namedSourceSetsContainer(platformProject)) { commonSourceSet, _ ->
                 addCommonSourceSetToPlatformSourceSet(commonSourceSet, platformProject)
             }
         }
     }
 
+    /**
+     * Applies [whenMatched] to pairs of items with the same name in [containerA] and [containerB],
+     * regardless of the order in which they are added to the containers.
+     */
+    private fun <A, B> matchSymmetricallyByNames(
+        containerA: NamedDomainObjectContainer<A>,
+        containerB: NamedDomainObjectContainer<B>,
+        whenMatched: (A, B) -> Unit
+    ) {
+        val matchedNames = mutableSetOf<String>()
+
+        fun <T, R> NamedDomainObjectContainer<T>.matchAllWith(other: NamedDomainObjectContainer<R>, match: (T, R) -> Unit) {
+            this@matchAllWith.all { item ->
+                val itemName = this@matchAllWith.namer.determineName(item)
+                if (itemName !in matchedNames) {
+                    val otherItem = other.findByName(itemName)
+                    if (otherItem != null) {
+                        matchedNames += itemName
+                        match(item, otherItem)
+                    }
+                }
+            }
+        }
+        containerA.matchAllWith(containerB) { a, b -> whenMatched(a, b) }
+        containerB.matchAllWith(containerA) { b, a -> whenMatched(a, b) }
+    }
+
+    protected open fun namedSourceSetsContainer(project: Project): NamedDomainObjectContainer<*> =
+            project.sourceSets
+
     protected open fun addCommonSourceSetToPlatformSourceSet(commonSourceSet: SourceSet, platformProject: Project) {
-        val platformTask = platformKotlinTasksBySourceSetName[commonSourceSet.name]
-        commonSourceSet.kotlin!!.srcDirs.forEach { platformTask?.source(it) }
+        val platformTask = platformProject.tasks
+            .filterIsInstance<AbstractKotlinCompile<*>>()
+            .firstOrNull { it.sourceSetName == commonSourceSet.name }
+
+        platformTask?.source(commonSourceSet.kotlin!!)
     }
 
     protected val SourceSet.kotlin: SourceDirectorySet?
@@ -148,6 +177,9 @@ open class KotlinPlatformAndroidPlugin : KotlinPlatformImplementationPluginBase(
         super.apply(project)
     }
 
+    override fun namedSourceSetsContainer(project: Project): NamedDomainObjectContainer<*> =
+        (project.extensions.getByName("android") as BaseExtension).sourceSets
+
     override fun addCommonSourceSetToPlatformSourceSet(commonSourceSet: SourceSet, platformProject: Project) {
         val androidExtension = platformProject.extensions.getByName("android") as BaseExtension
         val androidSourceSet = androidExtension.sourceSets.findByName(commonSourceSet.name) ?: return