[Gradle][MPP] Replace slow CompilationSourceSetUtil with InternalKotlinSourceSet api

^KT-52726 Verification Pending

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

@@ -384,7 +384,7 @@ internal open class SharedNativeCompilationDetails(
     override val friendArtifacts: FileCollection
         get() = super.friendArtifacts.plus(run {
             val project = target.project
-            val friendSourceSets = getVisibleSourceSetsFromAssociateCompilations(project, defaultSourceSet).toMutableSet().apply {
+            val friendSourceSets = getVisibleSourceSetsFromAssociateCompilations(defaultSourceSet).toMutableSet().apply {
                 // TODO: implement proper dependsOn/refines compiler args for Kotlin/Native and pass the dependsOn klibs separately;
                 //       But for now, those dependencies don't have any special semantics, so passing all them as friends works, too
                 addAll(defaultSourceSet.internal.dependsOnClosure)
@@ -564,7 +564,7 @@ internal class MetadataCompilationDetails(target: KotlinTarget, name: String) :
     override val friendArtifacts: FileCollection
         get() = super.friendArtifacts.plus(run {
             val project = target.project
-            val friendSourceSets = getVisibleSourceSetsFromAssociateCompilations(target.project, defaultSourceSet)
+            val friendSourceSets = getVisibleSourceSetsFromAssociateCompilations(defaultSourceSet)
             project.files(friendSourceSets.mapNotNull { target.compilations.findByName(it.name)?.output?.classesDirs })
         })
 }

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

@@ -131,7 +131,6 @@ class GradleProjectModuleBuilder(private val addInferredSourceSetVisibilityAsExp
 
         val extension = project.multiplatformExtensionOrNull
             ?: project.kotlinExtension
-            ?: return emptyList()
 
         val targets = when (extension) {
             is KotlinMultiplatformExtension -> extension.targets.filter { it.name != KotlinMultiplatformPlugin.METADATA_TARGET_NAME }
@@ -219,7 +218,7 @@ class GradleProjectModuleBuilder(private val addInferredSourceSetVisibilityAsExp
         if (addInferredSourceSetVisibilityAsExplicit) {
             project.kotlinExtension.sourceSets.forEach { sourceSet ->
                 val fragment = fragmentByName(sourceSet.name)
-                getVisibleSourceSetsFromAssociateCompilations(project, sourceSet).forEach { dependency ->
+                getVisibleSourceSetsFromAssociateCompilations(sourceSet).forEach { dependency ->
                     val dependencyFragment = fragmentByName(dependency.name)
                     fragment.declaredModuleDependencies.add(KpmModuleDependency(dependencyFragment.containingModule.moduleIdentifier))
                 }

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

@@ -368,7 +368,7 @@ internal fun requestedDependencies(
             project.configurations.sourceSetDependencyConfigurationByScope(sourceSet, scope).incoming.dependencies
         }
 
-    val otherContributingSourceSets = dependsOnClosureWithInterCompilationDependencies(project, sourceSet)
+    val otherContributingSourceSets = dependsOnClosureWithInterCompilationDependencies(sourceSet)
     return listOf(sourceSet, *otherContributingSourceSets.toTypedArray()).flatMap(::collectScopedDependenciesFromSourceSet)
 }
 

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

@@ -29,6 +29,7 @@ import org.jetbrains.kotlin.gradle.plugin.mpp.pm20.util.copyAttributes
 import org.jetbrains.kotlin.gradle.plugin.sources.DefaultLanguageSettingsBuilder
 import org.jetbrains.kotlin.gradle.plugin.sources.SourceSetMetadataStorageForIde
 import org.jetbrains.kotlin.gradle.plugin.sources.checkSourceSetVisibilityRequirements
+import org.jetbrains.kotlin.gradle.plugin.sources.internal
 import org.jetbrains.kotlin.gradle.plugin.statistics.KotlinBuildStatsService
 import org.jetbrains.kotlin.gradle.scripting.internal.ScriptingGradleSubplugin
 import org.jetbrains.kotlin.gradle.targets.js.ir.KotlinJsIrTargetPreset
@@ -115,8 +116,8 @@ class KotlinMultiplatformPlugin : Plugin<Project> {
             project.multiplatformExtension.metadata().compilations.getByName(KotlinCompilation.MAIN_COMPILATION_NAME)
 
         val primaryCompilationsBySourceSet by lazy { // don't evaluate eagerly: Android targets are not created at this point
-            val allCompilationsForSourceSets = CompilationSourceSetUtil.compilationsBySourceSets(project).mapValues { (_, compilations) ->
+            val allCompilationsForSourceSets = project.multiplatformExtension.sourceSets.associateWith { sourceSet ->
-                compilations.filter { it.target.platformType != KotlinPlatformType.common }
+                sourceSet.internal.compilations.filter { compilation -> compilation.target.platformType != KotlinPlatformType.common }
             }
 
             allCompilationsForSourceSets.mapValues { (_, compilations) -> // choose one primary compilation
@@ -347,14 +348,14 @@ internal fun sourcesJarTaskNamed(
 
 internal fun Project.setupGeneralKotlinExtensionParameters() {
     val sourceSetsInMainCompilation by lazy {
-        CompilationSourceSetUtil.compilationsBySourceSets(project).filterValues { compilations ->
+        kotlinExtension.sourceSets.filter { sourceSet ->
-            compilations.any {
+            sourceSet.internal.compilations.any {
                 // kotlin main compilation
                 it.isMain()
                         // android compilation which is NOT in tested variant
                         || (it as? KotlinJvmAndroidCompilation)?.let { getTestedVariantData(it.androidVariant) == null } == true
             }
-        }.keys
+        }
     }
 
     kotlinExtension.sourceSets.all { sourceSet ->

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

@@ -168,14 +168,14 @@ private fun buildKotlinProjectStructureMetadata(extension: KotlinMultiplatformEx
              * published as API dependencies of the metadata module to get into the resolution result, see
              * [KotlinMetadataTargetConfigurator.exportDependenciesForPublishing].
              */
-            val isNativeSharedSourceSet = isSharedNativeSourceSet(project, sourceSet)
+            val isNativeSharedSourceSet = isSharedNativeSourceSet(sourceSet)
             val scopes = listOfNotNull(
                 KotlinDependencyScope.API_SCOPE,
                 KotlinDependencyScope.IMPLEMENTATION_SCOPE.takeIf { isNativeSharedSourceSet }
             )
             val sourceSetsToIncludeDependencies =
                 if (isNativeSharedSourceSet)
-                    dependsOnClosureWithInterCompilationDependencies(project, sourceSet).plus(sourceSet)
+                    dependsOnClosureWithInterCompilationDependencies(sourceSet).plus(sourceSet)
                 else listOf(sourceSet)
             val sourceSetExportedDependencies = scopes.flatMap { scope ->
                 sourceSetsToIncludeDependencies.flatMap { hierarchySourceSet ->
@@ -185,7 +185,7 @@ private fun buildKotlinProjectStructureMetadata(extension: KotlinMultiplatformEx
             sourceSet.name to sourceSetExportedDependencies.map { ModuleIds.fromDependency(it) }.toSet()
         },
         sourceSetCInteropMetadataDirectory = sourceSetsWithMetadataCompilations.keys
-            .filter { isSharedNativeSourceSet(project, it) }
+            .filter { isSharedNativeSourceSet(it) }
             .associate { sourceSet -> sourceSet.name to cinteropMetadataDirectoryPath(sourceSet.name) },
         hostSpecificSourceSets = getHostSpecificSourceSets(project)
             .filter { it in sourceSetsWithMetadataCompilations }.map { it.name }

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

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.gradle.plugin.KotlinPlatformType
 import org.jetbrains.kotlin.gradle.plugin.KotlinSourceSet
 import org.jetbrains.kotlin.gradle.plugin.mpp.pm20.toSingleKpmModuleIdentifier
 import org.jetbrains.kotlin.gradle.plugin.sources.KotlinDependencyScope
+import org.jetbrains.kotlin.gradle.plugin.sources.internal
 import java.io.File
 
 internal data class SourceSetVisibilityResult(
@@ -54,7 +55,7 @@ internal class SourceSetVisibilityProvider(
         dependencyProjectStructureMetadata: KotlinProjectStructureMetadata,
         resolvedToOtherProject: Project?
     ): SourceSetVisibilityResult {
-        val compilations = CompilationSourceSetUtil.compilationsBySourceSets(project).getValue(visibleFrom)
+        val compilations = visibleFrom.internal.compilations
 
         val component = resolvedRootMppDependency ?: resolvedMetadataDependency
         val mppModuleIdentifier = component.toSingleKpmModuleIdentifier()

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

@@ -58,19 +58,13 @@ open class TransformKotlinGranularMetadata
     @Suppress("unused") // Gradle input
     @get:Input
     internal val inputSourceSetsAndCompilations: Map<String, Iterable<String>> by project.provider {
-        val sourceSets = participatingSourceSets
+        participatingSourceSets.associate { sourceSet ->
-        CompilationSourceSetUtil.compilationsBySourceSets(project)
+            sourceSet.name to sourceSet.internal.compilations.map { it.name }.sorted()
-            .filterKeys { it in sourceSets }
+        }
-            .entries.associate { (sourceSet, compilations) ->
-                sourceSet.name to compilations.map { it.name }.sorted()
-            }
     }
 
     private val participatingCompilations: Iterable<KotlinCompilation<*>>
-        get() {
+        get() = participatingSourceSets.flatMap { it.internal.compilations }.toSet()
-            val sourceSets = participatingSourceSets
-            return CompilationSourceSetUtil.compilationsBySourceSets(project).filterKeys { it in sourceSets }.values.flatten()
-        }
 
     @Suppress("unused") // Gradle input
     @get:Input
@@ -89,7 +83,7 @@ open class TransformKotlinGranularMetadata
             kotlinSourceSet,
             listOf(API_SCOPE, IMPLEMENTATION_SCOPE, COMPILE_ONLY_SCOPE),
             lazy {
-                dependsOnClosureWithInterCompilationDependencies(project, kotlinSourceSet).map {
+                dependsOnClosureWithInterCompilationDependencies(kotlinSourceSet).map {
                     project.tasks.withType(TransformKotlinGranularMetadata::class.java)
                         .getByName(KotlinMetadataTargetConfigurator.transformGranularMetadataTaskName(it.name))
                         .transformation

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

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.gradle.dsl.kotlinExtension
 import org.jetbrains.kotlin.gradle.plugin.KotlinSourceSet
 import org.jetbrains.kotlin.gradle.plugin.extraProperties
 import org.jetbrains.kotlin.gradle.plugin.sources.android.androidSourceSetInfoOrNull
+import org.jetbrains.kotlin.gradle.plugin.sources.internal
 import org.jetbrains.kotlin.gradle.plugin.whenEvaluated
 
 object UnusedSourceSetsChecker {
@@ -39,11 +40,10 @@ object UnusedSourceSetsChecker {
 
     fun checkSourceSets(project: Project) {
         project.whenEvaluated {
-            val compilationsBySourceSet = CompilationSourceSetUtil.compilationsBySourceSets(project)
             val unusedSourceSets = project.kotlinExtension.sourceSets
                 // Ignoring Android source sets
                 .filter { it.androidSourceSetInfoOrNull == null }
-                .filter { compilationsBySourceSet[it]?.isEmpty() ?: true }
+                .filter { it.internal.compilations.isEmpty() }
             if (unusedSourceSets.isNotEmpty()) {
                 reportUnusedSourceSets(project, unusedSourceSets.toSet())
             }

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

@@ -86,75 +86,6 @@ internal fun KotlinCompilation<*>.disambiguateName(simpleName: String): String {
     )
 }
 
-private typealias CompilationsBySourceSet = Map<KotlinSourceSet, Set<KotlinCompilation<*>>>
-
-internal object CompilationSourceSetUtil {
-    private const val EXT_NAME = "kotlin.compilations.bySourceSets"
-
-    @Suppress("UNCHECKED_CAST")
-    private fun getOrCreateProperty(
-        project: Project,
-        initialize: Property<CompilationsBySourceSet>.() -> Unit
-    ): Property<CompilationsBySourceSet> {
-        val ext = project.extensions.getByType(ExtraPropertiesExtension::class.java)
-        if (!ext.has(EXT_NAME)) {
-            ext.set(EXT_NAME, project.objects.property(Any::class.java as Class<CompilationsBySourceSet>).also(initialize))
-        }
-        return ext.get(EXT_NAME) as Property<CompilationsBySourceSet>
-    }
-
-    // FIXME: the results include the compilations of the metadata target; the callers should care about filtering them out
-    //        if they need only the platform compilations
-    // TODO: create a separate util function: `platformCompilationsBySourceSets`
-    // TODO: visit all call sites, check if they handle the metadata compilations correctly
-    fun compilationsBySourceSets(project: Project): CompilationsBySourceSet {
-        val compilationNamesBySourceSetName = getOrCreateProperty(project) {
-            var shouldFinalizeValue = false
-
-            set(project.provider {
-                val kotlinExtension = project.kotlinExtension
-                val targets = when (kotlinExtension) {
-                    is KotlinMultiplatformExtension -> kotlinExtension.targets
-                    is KotlinSingleTargetExtension<*> -> listOf(kotlinExtension.target)
-                    else -> emptyList()
-                }
-
-                val compilations = targets.flatMap { it.compilations }
-
-                val result = mutableMapOf<KotlinSourceSet, MutableSet<KotlinCompilation<*>>>().apply {
-                    compilations.forEach { compilation ->
-                        compilation.allKotlinSourceSets.forEach { sourceSet ->
-                            getOrPut(sourceSet) { mutableSetOf() }.add(compilation)
-                        }
-                    }
-                    kotlinExtension.sourceSets.forEach { sourceSet ->
-                        // For source sets not taking part in any compilation, keep an empty set to avoid errors on access by key
-                        getOrPut(sourceSet) { mutableSetOf() }
-                    }
-                }
-
-                if (shouldFinalizeValue) {
-                    set(result)
-                }
-
-                return@provider result
-            })
-
-            project.gradle.taskGraph.whenReady { shouldFinalizeValue = true }
-
-            // In case the value is first queried after the task graph has been calculated, finalize the value as soon as a task executes:
-            object : TaskExecutionListener {
-                override fun beforeExecute(task: Task) = Unit
-                override fun afterExecute(task: Task, state: TaskState) {
-                    shouldFinalizeValue = true
-                }
-            }
-        }
-
-        return compilationNamesBySourceSetName.get()
-    }
-}
-
 private val invalidModuleNameCharactersRegex = """[\\/\r\n\t]""".toRegex()
 
 internal fun filterModuleName(moduleName: String): String =

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/sources/DefaultKotlinSourceSet.kt

@@ -148,8 +148,7 @@ abstract class DefaultKotlinSourceSet @Inject constructor(
         return getDependenciesTransformation(scope)
     }
 
-    fun getAdditionalVisibleSourceSets(): List<KotlinSourceSet> =
+    fun getAdditionalVisibleSourceSets(): List<KotlinSourceSet> = getVisibleSourceSetsFromAssociateCompilations(this)
-        getVisibleSourceSetsFromAssociateCompilations(project, this)
 
     internal fun getDependenciesTransformation(scope: KotlinDependencyScope): Iterable<MetadataDependencyTransformation> {
         val metadataDependencyResolutionByModule =

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/sources/SourceSetsVisibilityInference.kt

@@ -10,18 +10,14 @@ import org.gradle.api.Project
 import org.jetbrains.kotlin.gradle.dsl.kotlinExtension
 import org.jetbrains.kotlin.gradle.plugin.KotlinCompilation
 import org.jetbrains.kotlin.gradle.plugin.KotlinSourceSet
-import org.jetbrains.kotlin.gradle.plugin.mpp.CompilationSourceSetUtil
 import org.jetbrains.kotlin.gradle.plugin.mpp.associateWithClosure
 
 fun getSourceSetsFromAssociatedCompilations(fromCompilation: KotlinCompilation<*>): Map<KotlinCompilation<*>, Set<KotlinSourceSet>> =
     fromCompilation.associateWithClosure.associate { it to it.allKotlinSourceSets }
 
 fun getVisibleSourceSetsFromAssociateCompilations(
-    project: Project,
     sourceSet: KotlinSourceSet
-): List<KotlinSourceSet> = getVisibleSourceSetsFromAssociateCompilations(
+): List<KotlinSourceSet> = getVisibleSourceSetsFromAssociateCompilations(sourceSet.internal.compilations)
-    CompilationSourceSetUtil.compilationsBySourceSets(project).getValue(sourceSet)
-)
 
 internal fun getVisibleSourceSetsFromAssociateCompilations(
     participatesInCompilations: Set<KotlinCompilation<*>>
@@ -124,26 +120,22 @@ class UnsatisfiedSourceSetVisibilityException(
         }
 }
 
-fun checkSourceSetVisibilityRequirements(
+fun checkSourceSetVisibilityRequirements(project: Project) = checkSourceSetVisibilityRequirements(
-    project: Project
+    project.kotlinExtension.sourceSets
-) = checkSourceSetVisibilityRequirements(
-    project.kotlinExtension.sourceSets,
-    CompilationSourceSetUtil.compilationsBySourceSets(project)
 )
 
 internal fun checkSourceSetVisibilityRequirements(
     sourceSets: Iterable<KotlinSourceSet>,
-    compilationsBySourceSet: Map<KotlinSourceSet, Set<KotlinCompilation<*>>>
 ) {
     sourceSets.forEach { sourceSet ->
         val requiredVisibility = sourceSet.requiresVisibilityOf
         val inferredVisibility =
-            getVisibleSourceSetsFromAssociateCompilations(compilationsBySourceSet[sourceSet].orEmpty())
+            getVisibleSourceSetsFromAssociateCompilations(sourceSet.internal.compilations)
 
         val requiredButNotVisible = requiredVisibility - inferredVisibility - sourceSet.internal.withDependsOnClosure
 
         if (requiredButNotVisible.isNotEmpty()) {
-            val compilations = compilationsBySourceSet.getValue(sourceSet)
+            val compilations = sourceSet.internal.compilations
 
             throw UnsatisfiedSourceSetVisibilityException(
                 sourceSet,

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

@@ -20,7 +20,6 @@ import org.gradle.api.tasks.bundling.Zip
 import org.jetbrains.kotlin.gradle.dsl.*
 import org.jetbrains.kotlin.gradle.plugin.*
 import org.jetbrains.kotlin.gradle.plugin.mpp.*
-import org.jetbrains.kotlin.gradle.plugin.mpp.CompilationSourceSetUtil.compilationsBySourceSets
 import org.jetbrains.kotlin.gradle.plugin.mpp.pm20.*
 import org.jetbrains.kotlin.gradle.plugin.mpp.pm20.hasKpmModel
 import org.jetbrains.kotlin.gradle.plugin.sources.*
@@ -189,7 +188,7 @@ class KotlinMetadataTargetConfigurator :
                 createMetadataCompilation(target, sourceSet, allMetadataJar, sourceSet in hostSpecificSourceSets)
             }
             .onEach { (sourceSet, compilation) ->
-                if (!isMetadataCompilationSupported(target.project, sourceSet)) {
+                if (!isMetadataCompilationSupported(sourceSet)) {
                     compilation.compileKotlinTaskProvider.configure { it.enabled = false }
                 }
             }
@@ -204,8 +203,8 @@ class KotlinMetadataTargetConfigurator :
         }
     }
 
-    private fun isMetadataCompilationSupported(project: Project, sourceSet: KotlinSourceSet): Boolean {
+    private fun isMetadataCompilationSupported(sourceSet: KotlinSourceSet): Boolean {
-        val platforms = compilationsBySourceSets(project)[sourceSet].orEmpty()
+        val platforms = sourceSet.internal.compilations
             .filter { it.target !is KotlinMetadataTarget }
             .map { it.target.platformType }.distinct()
 
@@ -300,10 +299,10 @@ class KotlinMetadataTargetConfigurator :
 
         val compilationName = sourceSet.name
 
-        val platformCompilations = compilationsBySourceSets(project)
+        val platformCompilations = sourceSet.internal.compilations
-            .getValue(sourceSet).filter { it.target.name != KotlinMultiplatformPlugin.METADATA_TARGET_NAME }
+            .filter { it.target.name != KotlinMultiplatformPlugin.METADATA_TARGET_NAME }
 
-        val isNativeSourceSet = isSharedNativeSourceSet(project, sourceSet)
+        val isNativeSourceSet = isSharedNativeSourceSet(sourceSet)
 
         val compilationFactory: KotlinCompilationFactory<out AbstractKotlinCompilation<*>> = when {
             isNativeSourceSet -> KotlinSharedNativeCompilationFactory(
@@ -368,7 +367,7 @@ class KotlinMetadataTargetConfigurator :
                 sourceSet,
                 listOf(scope),
                 lazy {
-                    dependsOnClosureWithInterCompilationDependencies(project, sourceSet).filterIsInstance<DefaultKotlinSourceSet>()
+                    dependsOnClosureWithInterCompilationDependencies(sourceSet).filterIsInstance<DefaultKotlinSourceSet>()
                         .map { checkNotNull(it.dependencyTransformations[scope]) }
                 }
             )
@@ -566,16 +565,17 @@ internal fun createMetadataDependencyTransformationClasspath(
     )
 }
 
-internal fun isSharedNativeSourceSet(project: Project, sourceSet: KotlinSourceSet): Boolean {
+internal fun isSharedNativeSourceSet(sourceSet: KotlinSourceSet): Boolean {
-    val compilations = compilationsBySourceSets(project)[sourceSet].orEmpty()
+    val compilations = sourceSet.internal.compilations
     return compilations.isNotEmpty() && compilations.all {
         it.platformType == KotlinPlatformType.common || it.platformType == KotlinPlatformType.native
     }
 }
 
-internal fun dependsOnClosureWithInterCompilationDependencies(project: Project, sourceSet: KotlinSourceSet): Set<KotlinSourceSet> =
+internal fun dependsOnClosureWithInterCompilationDependencies(sourceSet: KotlinSourceSet): Set<KotlinSourceSet> =
     sourceSet.internal.dependsOnClosure.toMutableSet().apply {
-        addAll(getVisibleSourceSetsFromAssociateCompilations(project, sourceSet))
+        addAll(getVisibleSourceSetsFromAssociateCompilations(sourceSet))
+
     }
 
 /**
@@ -588,7 +588,7 @@ internal fun getCommonSourceSetsForMetadataCompilation(project: Project): Set<Ko
         return setOf(project.multiplatformExtension.sourceSets.getByName(KotlinSourceSet.COMMON_MAIN_SOURCE_SET_NAME))
 
     val compilationsBySourceSet: Map<KotlinSourceSet, Set<KotlinCompilation<*>>> =
-        compilationsBySourceSets(project)
+        project.kotlinExtension.sourceSets.associateWith { it.internal.compilations }
 
     val sourceSetsUsedInMultipleTargets = compilationsBySourceSet.filterValues { compilations ->
         compilations.map { it.target.platformType }.distinct().run {

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/targets/native/KotlinNativeTarget.kt

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.gradle.plugin.*
 import org.jetbrains.kotlin.gradle.plugin.mpp.pm20.*
 import org.jetbrains.kotlin.gradle.plugin.mpp.pm20.isNativeShared
 import org.jetbrains.kotlin.gradle.plugin.mpp.pm20.variantsContainingFragment
+import org.jetbrains.kotlin.gradle.plugin.sources.internal
 import org.jetbrains.kotlin.gradle.targets.metadata.*
 import org.jetbrains.kotlin.gradle.targets.metadata.filesWithUnpackedArchives
 import org.jetbrains.kotlin.gradle.targets.metadata.isKotlinGranularMetadataEnabled
@@ -175,18 +176,14 @@ internal fun getHostSpecificFragments(
 )
 
 internal fun getHostSpecificSourceSets(project: Project): Set<KotlinSourceSet> {
-    val compilationsBySourceSet = CompilationSourceSetUtil.compilationsBySourceSets(project).mapValues { (_, compilations) ->
-        compilations.filter { it !is KotlinMetadataCompilation<*> }
-    }
-
     return getHostSpecificElements(
         project.kotlinExtension.sourceSets,
         isNativeShared = { sourceSet ->
-            val compilations = compilationsBySourceSet[sourceSet].orEmpty()
+            val compilations = sourceSet.internal.compilations
             compilations.isNotEmpty() && compilations.all { it.platformType == KotlinPlatformType.native }
         },
         getKonanTargets = { sourceSet ->
-            compilationsBySourceSet[sourceSet].orEmpty()
+            sourceSet.internal.compilations
                 .filterIsInstance<KotlinNativeCompilation>()
                 .mapTo(mutableSetOf()) { it.konanTarget }
         }

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/targets/native/internal/CInteropCommonizerDependent.kt

@@ -8,10 +8,10 @@ package org.jetbrains.kotlin.gradle.targets.native.internal
 import org.gradle.api.Project
 import org.jetbrains.kotlin.commonizer.SharedCommonizerTarget
 import org.jetbrains.kotlin.gradle.plugin.KotlinSourceSet
-import org.jetbrains.kotlin.gradle.plugin.mpp.CompilationSourceSetUtil.compilationsBySourceSets
 import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinNativeCompilation
 import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinSharedNativeCompilation
 import org.jetbrains.kotlin.gradle.plugin.mpp.associateWithClosure
+import org.jetbrains.kotlin.gradle.plugin.sources.internal
 import org.jetbrains.kotlin.gradle.targets.native.internal.CInteropIdentifier.Scope
 import org.jetbrains.kotlin.tooling.core.UnsafeApi
 
@@ -83,7 +83,7 @@ internal fun CInteropCommonizerDependent.Factory.from(compilation: KotlinSharedN
 internal fun CInteropCommonizerDependent.Factory.from(project: Project, sourceSet: KotlinSourceSet): CInteropCommonizerDependent? {
     return from(
         target = project.getCommonizerTarget(sourceSet) as? SharedCommonizerTarget ?: return null,
-        compilations = (compilationsBySourceSets(project)[sourceSet] ?: return null)
+        compilations = sourceSet.internal.compilations
             .filterIsInstance<KotlinNativeCompilation>().toSet()
     )
 }
@@ -93,7 +93,7 @@ internal fun CInteropCommonizerDependent.Factory.fromAssociateCompilations(
 ): CInteropCommonizerDependent? {
     return from(
         target = project.getCommonizerTarget(sourceSet) as? SharedCommonizerTarget ?: return null,
-        compilations = (compilationsBySourceSets(project)[sourceSet] ?: return null)
+        compilations = sourceSet.internal.compilations
             .filterIsInstance<KotlinNativeCompilation>()
             .flatMap { compilation -> compilation.associateWithClosure }
             .filterIsInstance<KotlinNativeCompilation>()

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/targets/native/internal/CInteropPropagatedDependencies.kt

@@ -8,11 +8,11 @@ package org.jetbrains.kotlin.gradle.targets.native.internal
 import org.gradle.api.Project
 import org.gradle.api.file.FileCollection
 import org.jetbrains.kotlin.gradle.dsl.multiplatformExtensionOrNull
-import org.jetbrains.kotlin.gradle.plugin.mpp.CompilationSourceSetUtil
 import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinMetadataCompilation
 import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinNativeCompilation
 import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinSharedNativeCompilation
 import org.jetbrains.kotlin.gradle.plugin.sources.DefaultKotlinSourceSet
+import org.jetbrains.kotlin.gradle.plugin.sources.internal
 import org.jetbrains.kotlin.gradle.tasks.CInteropProcess
 import org.jetbrains.kotlin.gradle.utils.filesProvider
 import java.io.File
@@ -61,7 +61,7 @@ private fun Project.getPropagatedCInteropDependenciesOrEmpty(sourceSet: DefaultK
         compatibility metadata variant will still register
         a 'KotlinMetadataCompilation for 'commonMain' which is irrelevant here
         */
-        val compilations = CompilationSourceSetUtil.compilationsBySourceSets(this)[sourceSet].orEmpty()
+        val compilations = sourceSet.internal.compilations
             .filter { compilation -> compilation !is KotlinMetadataCompilation }
 
         /* Participating in multiple compilations? -> can't propagate -> should be commonized */

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/targets/native/internal/GetCommonizerTargetOfSourceSet.kt

@@ -10,11 +10,11 @@ import org.jetbrains.kotlin.commonizer.CommonizerTarget
 import org.jetbrains.kotlin.commonizer.SharedCommonizerTarget
 import org.jetbrains.kotlin.commonizer.allLeaves
 import org.jetbrains.kotlin.gradle.plugin.KotlinSourceSet
-import org.jetbrains.kotlin.gradle.plugin.mpp.CompilationSourceSetUtil.compilationsBySourceSets
 import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinMetadataCompilation
+import org.jetbrains.kotlin.gradle.plugin.sources.internal
 
 internal fun Project.getCommonizerTarget(sourceSet: KotlinSourceSet): CommonizerTarget? {
-    val allCompilationLeafTargets = compilationsBySourceSets(this)[sourceSet].orEmpty()
+    val allCompilationLeafTargets = sourceSet.internal.compilations
         .filter { compilation -> compilation !is KotlinMetadataCompilation }
         .map { compilation -> getCommonizerTarget(compilation) ?: return null }
         .allLeaves()

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/targets/native/internal/KotlinNativePlatformDependencies.kt

@@ -69,7 +69,7 @@ internal val KotlinMultiplatformExtension.nativeRootSourceSets: Collection<Kotli
     get() {
         val nativeSourceSets = sourceSets.filter { sourceSet -> project.getCommonizerTarget(sourceSet) != null }
         return nativeSourceSets.filter { sourceSet ->
-            val allVisibleSourceSets = sourceSet.dependsOn + getVisibleSourceSetsFromAssociateCompilations(project, sourceSet)
+            val allVisibleSourceSets = sourceSet.dependsOn + getVisibleSourceSetsFromAssociateCompilations(sourceSet)
             allVisibleSourceSets.none { dependency ->
                 dependency in nativeSourceSets
             }