[Gradle][MPP] Implement InternalKotlinCompilation and make .kotlinSourceSets and .allKotlinSourceSets observable

The effect of this changes are twofold:

1) Maintaining the 'allKotlinSourceSets' closure
makes access to this field cheaper

2) The observability of changes in the kotlinSourceSets
will allow implementing a faster compilationsBySourceSet util

KT-52726

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/mpp/AbstractCompilationDetails.kt

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.gradle.plugin.mpp
+
+import org.jetbrains.kotlin.gradle.dsl.KotlinCommonOptions
+import org.jetbrains.kotlin.gradle.plugin.KotlinSourceSet
+import org.jetbrains.kotlin.gradle.plugin.sources.internal
+import org.jetbrains.kotlin.gradle.utils.MutableObservableSet
+import org.jetbrains.kotlin.gradle.utils.ObservableSet
+
+abstract class AbstractCompilationDetails<T : KotlinCommonOptions> : CompilationDetails<T> {
+    private val directlyIncludedKotlinSourceSetsImpl: MutableObservableSet<KotlinSourceSet> by lazy {
+        MutableObservableSet(defaultSourceSet)
+    }
+
+    final override val directlyIncludedKotlinSourceSets: ObservableSet<KotlinSourceSet>
+        get() = directlyIncludedKotlinSourceSetsImpl
+
+    private val kotlinSourceSetsClosureImpl: MutableObservableSet<KotlinSourceSet> by lazy {
+        MutableObservableSet<KotlinSourceSet>().also { set ->
+            defaultSourceSet.internal.withDependsOnClosure.forAll(set::add)
+        }
+    }
+
+    final override val kotlinSourceSetsClosure: ObservableSet<KotlinSourceSet>
+        get() = kotlinSourceSetsClosureImpl
+
+    final override fun source(sourceSet: KotlinSourceSet) {
+        directlyIncludedKotlinSourceSetsImpl.add(sourceSet)
+        sourceSet.internal.withDependsOnClosure.forAll { withDependsOn ->
+            kotlinSourceSetsClosureImpl.add(withDependsOn)
+        }
+
+        whenSourceSetAdded(sourceSet)
+    }
+
+    /**
+     * Called after [sourceSet] added the [sourceSet] to its respective ObservableSet's
+     */
+    protected open fun whenSourceSetAdded(sourceSet: KotlinSourceSet) = Unit
+}

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/mpp/AbstractKotlinCompilation.kt

@@ -16,14 +16,14 @@ import org.jetbrains.kotlin.gradle.dsl.KotlinCompile
 import org.jetbrains.kotlin.gradle.dsl.kotlinExtension
 import org.jetbrains.kotlin.gradle.plugin.*
 import org.jetbrains.kotlin.gradle.plugin.mpp.pm20.KotlinCompilationData
-import org.jetbrains.kotlin.gradle.plugin.sources.withDependsOnClosure
 import org.jetbrains.kotlin.gradle.tasks.locateTask
+import org.jetbrains.kotlin.gradle.utils.ObservableSet
 import org.jetbrains.kotlin.project.model.LanguageSettings
 import java.util.*
 
 abstract class AbstractKotlinCompilation<T : KotlinCommonOptions>(
     internal open val compilationDetails: CompilationDetails<T>,
-) : KotlinCompilation<T>,
+) : InternalKotlinCompilation<T>,
     HasKotlinDependencies,
     KotlinCompilationData<T> {
 
@@ -62,12 +62,11 @@ abstract class AbstractKotlinCompilation<T : KotlinCommonOptions>(
         get() = compilationDetails.compileDependencyFilesHolder.dependencyFiles
         set(value) { compilationDetails.compileDependencyFilesHolder.dependencyFiles = value }
 
-    final override val kotlinSourceSets: MutableSet<KotlinSourceSet> get() =
+    final override val kotlinSourceSets: ObservableSet<KotlinSourceSet>
-        when (val details = compilationDetails) {
+        get() = compilationDetails.directlyIncludedKotlinSourceSets
-            is DefaultCompilationDetails -> details.directlyIncludedKotlinSourceSets // mutable in that subtype
+
-            // TODO deprecate mutability of this set. We shouldn't allow mutating it directly anyway;
+    override val allKotlinSourceSets: ObservableSet<KotlinSourceSet>
-            else -> details.directlyIncludedKotlinSourceSets.toMutableSet()
+        get() = compilationDetails.kotlinSourceSetsClosure
-        }
 
     final override val defaultSourceSetName: String get() = compilationDetails.defaultSourceSetName
 
@@ -112,9 +111,6 @@ abstract class AbstractKotlinCompilation<T : KotlinCommonOptions>(
     override val owner: KotlinTarget
         get() = target
 
-    override val allKotlinSourceSets: Set<KotlinSourceSet>
-        get() = compilationDetails.directlyIncludedKotlinSourceSets.withDependsOnClosure
-
     override val relatedConfigurationNames: List<String>
         get() = compilationDetails.kotlinDependenciesHolder.relatedConfigurationNames + compileDependencyConfigurationName
 }

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/mpp/CompilationDetails.kt

@@ -50,7 +50,8 @@ interface CompilationDetails<T : KotlinCommonOptions> {
 
     fun source(sourceSet: KotlinSourceSet)
 
-    val directlyIncludedKotlinSourceSets: Set<KotlinSourceSet>
+    val directlyIncludedKotlinSourceSets: ObservableSet<KotlinSourceSet>
+    val kotlinSourceSetsClosure: ObservableSet<KotlinSourceSet>
     val defaultSourceSetName: String
 
     @Suppress("UNCHECKED_CAST")
@@ -72,7 +73,7 @@ open class DefaultCompilationDetails<T : KotlinCommonOptions>(
     final override val target: KotlinTarget,
     final override val compilationPurpose: String,
     createKotlinOptions: DefaultCompilationDetails<*>.() -> T
-) : CompilationDetails<T>, KotlinCompilationData<T> {
+) : AbstractCompilationDetails<T>(), KotlinCompilationData<T> {
 
     override val kotlinOptions: T by lazy { createKotlinOptions() }
 
@@ -104,8 +105,6 @@ open class DefaultCompilationDetails<T : KotlinCommonOptions>(
         )
     )
 
-    override val directlyIncludedKotlinSourceSets: MutableSet<KotlinSourceSet> = mutableSetOf()
-
     override val defaultSourceSetName: String
         get() = lowerCamelCaseName(
             target.disambiguationClassifier.takeIf { target !is KotlinMetadataTarget },
@@ -263,12 +262,9 @@ open class DefaultCompilationDetails<T : KotlinCommonOptions>(
         }
     }
 
-
+    override fun whenSourceSetAdded(sourceSet: KotlinSourceSet) {
-    override fun source(sourceSet: KotlinSourceSet) {
+        sourceSet.internal.withDependsOnClosure.forAll { inWithDependsOnClosure ->
-        if (directlyIncludedKotlinSourceSets.add(sourceSet)) {
+            addExactSourceSetEagerly(inWithDependsOnClosure)
-            sourceSet.internal.withDependsOnClosure.forAll { inWithDependsOnClosure ->
-                addExactSourceSetEagerly(inWithDependsOnClosure)
-            }
         }
     }
 
@@ -280,7 +276,12 @@ open class DefaultCompilationDetails<T : KotlinCommonOptions>(
             addAsCommonSources
         ) { sourceSet.kotlin }
 
+
+    private val sourceSetsAddedEagerly = hashSetOf<KotlinSourceSet>()
+
     internal fun addExactSourceSetEagerly(sourceSet: KotlinSourceSet) {
+        if (!sourceSetsAddedEagerly.add(sourceSet)) return
+
         with(target.project) {
             //TODO possibly issue with forced instantiation
             addSourcesToCompileTask(
@@ -399,7 +400,7 @@ internal open class SharedNativeCompilationDetails(
 internal open class MetadataMappedCompilationDetails<T : KotlinCommonOptions>(
     override val target: KotlinMetadataTarget,
     final override val compilationData: AbstractKotlinFragmentMetadataCompilationData<T>
-) : CompilationDetails<T> {
+) : AbstractCompilationDetails<T>() {
     override val compileDependencyFilesHolder: GradleKpmDependencyFilesHolder =
         GradleKpmDependencyFilesHolder.ofMetadataCompilationDependencies(compilationData)
 
@@ -413,7 +414,7 @@ internal open class MetadataMappedCompilationDetails<T : KotlinCommonOptions>(
     override val associateCompilations: Set<CompilationDetails<*>>
         get() = emptySet()
 
-    override fun source(sourceSet: KotlinSourceSet) {
+    override fun whenSourceSetAdded(sourceSet: KotlinSourceSet) {
         throw UnsupportedOperationException("metadata compilations have predefined sources")
     }
 
@@ -422,9 +423,6 @@ internal open class MetadataMappedCompilationDetails<T : KotlinCommonOptions>(
             it.underlyingFragment == compilationData.fragment
         }
 
-    override val directlyIncludedKotlinSourceSets: MutableSet<KotlinSourceSet>
-        get() = Collections.unmodifiableSet(hashSetOf(underlyingSourceSet))
-
     override val defaultSourceSetName: String
         get() = underlyingSourceSet.name
 }
@@ -432,7 +430,7 @@ internal open class MetadataMappedCompilationDetails<T : KotlinCommonOptions>(
 internal open class VariantMappedCompilationDetails<T : KotlinCommonOptions>(
     open val variant: GradleKpmVariantInternal,
     override val target: KotlinTarget
-) : CompilationDetails<T> {
+) : AbstractCompilationDetails<T>() {
 
     @Suppress("UNCHECKED_CAST")
     override val compilationData: KotlinCompilationData<T>
@@ -441,7 +439,7 @@ internal open class VariantMappedCompilationDetails<T : KotlinCommonOptions>(
     override val defaultSourceSetName: String
         get() = variant.unambiguousNameInProject
 
-    override fun source(sourceSet: KotlinSourceSet) {
+    override fun whenSourceSetAdded(sourceSet: KotlinSourceSet) {
         compilation.defaultSourceSet.dependsOn(sourceSet)
     }
 
@@ -462,8 +460,6 @@ internal open class VariantMappedCompilationDetails<T : KotlinCommonOptions>(
     override val kotlinDependenciesHolder: HasKotlinDependencies
         get() = variant
 
-    override val directlyIncludedKotlinSourceSets: Set<KotlinSourceSet>
-        get() = compilation.defaultSourceSet.dependsOn
 }
 
 internal open class VariantMappedCompilationDetailsWithRuntime<T : KotlinCommonOptions>(

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/mpp/InternalKotlinCompilation.kt

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.gradle.plugin.mpp
+
+import org.jetbrains.kotlin.gradle.dsl.KotlinCommonOptions
+import org.jetbrains.kotlin.gradle.plugin.KotlinCompilation
+import org.jetbrains.kotlin.gradle.plugin.KotlinSourceSet
+import org.jetbrains.kotlin.gradle.utils.ObservableSet
+
+internal interface InternalKotlinCompilation<T : KotlinCommonOptions> : KotlinCompilation<T> {
+    override val kotlinSourceSets: ObservableSet<KotlinSourceSet>
+    override val allKotlinSourceSets: ObservableSet<KotlinSourceSet>
+}
+
+internal val <T : KotlinCommonOptions> KotlinCompilation<T>.internal: InternalKotlinCompilation<T>
+    get() = (this as? InternalKotlinCompilation<T>) ?: throw IllegalArgumentException(
+        "KotlinCompilation($name) ${this::class} does not implement ${InternalKotlinCompilation::class}"
+    )

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/targets/metadata/KotlinMetadataTargetConfigurator.kt

@@ -454,7 +454,7 @@ class KotlinMetadataTargetConfigurator :
         }
 
         val resolvedMetadataFilesProviders = lazy {
-            val transformationTaskHolders = sourceSet.withDependsOnClosure.mapNotNull { hierarchySourceSet ->
+            val transformationTaskHolders = sourceSet.internal.withDependsOnClosure.mapNotNull { hierarchySourceSet ->
                 project.locateTask<TransformKotlinGranularMetadata>(transformGranularMetadataTaskName(hierarchySourceSet.name))
             }
             transformationTaskHolders.map { SourceSetResolvedMetadataProvider(it) }

libraries/tools/kotlin-gradle-plugin/src/functionalTest/kotlin/org/jetbrains/kotlin/gradle/mpp/InternalKotlinCompilationTest.kt

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+@file:Suppress("FunctionName")
+
+package org.jetbrains.kotlin.gradle.mpp
+
+import org.jetbrains.kotlin.gradle.plugin.KotlinCompilation
+import org.jetbrains.kotlin.gradle.plugin.mpp.InternalKotlinCompilation
+import org.jetbrains.kotlin.gradle.reflections
+import org.junit.Test
+import kotlin.reflect.full.isSubclassOf
+import kotlin.test.fail
+
+class InternalKotlinSourceSetTest {
+    @Test
+    fun `test - all implementations of KotlinCompilation - implement InternalKotlinCompilation`() {
+        val subtypesOfKotlinSourceSet = reflections.getSubTypesOf(KotlinCompilation::class.java)
+        subtypesOfKotlinSourceSet
+            .filter { subtype -> !subtype.isInterface }
+            .forEach { implementation ->
+                if (!implementation.kotlin.isSubclassOf(InternalKotlinCompilation::class)) {
+                    fail("$implementation does not implement ${InternalKotlinCompilation::class}")
+                }
+            }
+    }
+}