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

^KT-52726 Verification Pending
2022-09-06 16:29:30 +02:00
parent 0e6a074c9c
commit dc8a6b8def
17 changed files with 52 additions and 138 deletions
@@ -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 })
        })
 }
@@ -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))
                }
@@ -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)
 }

@@ -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) ->
-                compilations.filter { it.target.platformType != KotlinPlatformType.common }
+            val allCompilationsForSourceSets = project.multiplatformExtension.sourceSets.associateWith { sourceSet ->
+                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 ->
-            compilations.any {
+        kotlinExtension.sourceSets.filter { sourceSet ->
+            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 ->
@@ -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 }
@@ -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()
@@ -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
-        CompilationSourceSetUtil.compilationsBySourceSets(project)
-            .filterKeys { it in sourceSets }
-            .entries.associate { (sourceSet, compilations) ->
-                sourceSet.name to compilations.map { it.name }.sorted()
-            }
+        participatingSourceSets.associate { sourceSet ->
+            sourceSet.name to sourceSet.internal.compilations.map { it.name }.sorted()
+        }
    }

    private val participatingCompilations: Iterable<KotlinCompilation<*>>
-        get() {
-            val sourceSets = participatingSourceSets
-            return CompilationSourceSetUtil.compilationsBySourceSets(project).filterKeys { it in sourceSets }.values.flatten()
-        }
+        get() = participatingSourceSets.flatMap { it.internal.compilations }.toSet()

    @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
@@ -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())
            }
@@ -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 =
@@ -148,8 +148,7 @@ abstract class DefaultKotlinSourceSet @Inject constructor(
        return getDependenciesTransformation(scope)
    }

-    fun getAdditionalVisibleSourceSets(): List<KotlinSourceSet> =
-        getVisibleSourceSetsFromAssociateCompilations(project, this)
+    fun getAdditionalVisibleSourceSets(): List<KotlinSourceSet> = getVisibleSourceSetsFromAssociateCompilations(this)

    internal fun getDependenciesTransformation(scope: KotlinDependencyScope): Iterable<MetadataDependencyTransformation> {
        val metadataDependencyResolutionByModule =
@@ -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(
-    CompilationSourceSetUtil.compilationsBySourceSets(project).getValue(sourceSet)
-)
+): List<KotlinSourceSet> = getVisibleSourceSetsFromAssociateCompilations(sourceSet.internal.compilations)

 internal fun getVisibleSourceSetsFromAssociateCompilations(
    participatesInCompilations: Set<KotlinCompilation<*>>
@@ -124,26 +120,22 @@ class UnsatisfiedSourceSetVisibilityException(
        }
 }

-fun checkSourceSetVisibilityRequirements(
-    project: Project
-) = checkSourceSetVisibilityRequirements(
-    project.kotlinExtension.sourceSets,
-    CompilationSourceSetUtil.compilationsBySourceSets(project)
+fun checkSourceSetVisibilityRequirements(project: Project) = checkSourceSetVisibilityRequirements(
+    project.kotlinExtension.sourceSets
 )

 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,
@@ -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 {
-        val platforms = compilationsBySourceSets(project)[sourceSet].orEmpty()
+    private fun isMetadataCompilationSupported(sourceSet: KotlinSourceSet): Boolean {
+        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)
-            .getValue(sourceSet).filter { it.target.name != KotlinMultiplatformPlugin.METADATA_TARGET_NAME }
+        val platformCompilations = sourceSet.internal.compilations
+            .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 {
-    val compilations = compilationsBySourceSets(project)[sourceSet].orEmpty()
+internal fun isSharedNativeSourceSet(sourceSet: KotlinSourceSet): Boolean {
+    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 {
@@ -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 }
        }
@@ -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>()
@@ -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 */
@@ -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()
@@ -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
            }