Make compilationsBySourceSets safe to use during evaluation

This function used to only calculate the mapping once and store it for
the project, so it was unsafe to use early. Fix that by evaluating the
value each time and then freezing it at the point of either task graph
being ready, or any task being executed.

Also, store the map in the project's `ext` rather than in a static weak
hash map, which is safer wrt memory leaks.

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

@@ -7,10 +7,15 @@ package org.jetbrains.kotlin.gradle.plugin.mpp
 
 import groovy.lang.Closure
 import org.gradle.api.Project
+import org.gradle.api.Task
 import org.gradle.api.attributes.AttributeContainer
+import org.gradle.api.execution.TaskExecutionListener
 import org.gradle.api.file.FileCollection
 import org.gradle.api.plugins.BasePluginConvention
+import org.gradle.api.plugins.ExtraPropertiesExtension
+import org.gradle.api.provider.Property
 import org.gradle.api.tasks.TaskProvider
+import org.gradle.api.tasks.TaskState
 import org.gradle.util.ConfigureUtil
 import org.jetbrains.kotlin.gradle.dsl.*
 import org.jetbrains.kotlin.gradle.plugin.*
@@ -218,32 +223,28 @@ internal fun KotlinCompilation<*>.disambiguateName(simpleName: String): String {
     )
 }
 
+ private typealias CompilationsBySourceSet = Map<KotlinSourceSet, Set<KotlinCompilation<*>>>
+
 internal object CompilationSourceSetUtil {
-    // Store only names in the cache to avoid memory leak through indirect references to the project
+    private const val EXT_NAME = "kotlin.compilations.bySourceSets"
-    private data class TargetCompilationName(val targetName: String, val compilationName: String) {
+
-        fun toCompilation(project: Project): KotlinCompilation<*>? {
+    @Suppress("UNCHECKED_CAST")
-            val kotlinExtension = project.kotlinExtension
+    private fun getOrCreateProperty(
-            val target = when (kotlinExtension) {
+        project: Project,
-                is KotlinMultiplatformExtension -> kotlinExtension.targets.findByName(targetName)
+        initialize: Property<CompilationsBySourceSet>.() -> Unit
-                is KotlinSingleTargetExtension -> kotlinExtension.target.takeIf { it.name == targetName }
+    ): Property<CompilationsBySourceSet> {
-                else -> null
+        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 target?.compilations?.getByName(compilationName)
+        return ext.get(EXT_NAME) as Property<CompilationsBySourceSet>
     }
 
-        companion object {
+    fun compilationsBySourceSets(project: Project): CompilationsBySourceSet {
-            fun from(compilation: KotlinCompilation<*>) = TargetCompilationName(compilation.target.name, compilation.name)
+        val compilationNamesBySourceSetName = getOrCreateProperty(project) {
-        }
+            var shouldFinalizeValue = false
-    }
-
-    private val compilationsBySourceSetCache = WeakHashMap<Project, Map<String, Set<TargetCompilationName>>>()
-
-    /** Evaluates once per project. Don't access until all source set dependsOn relationships are built and all source sets are added
-     * to the relevant compilations. */
-    fun compilationsBySourceSets(project: Project): Map<KotlinSourceSet, Set<KotlinCompilation<*>>> {
-        val compilationNamesBySourceSetName = compilationsBySourceSetCache.computeIfAbsent(project) { _ ->
-            check(project.state.executed) { "Should only be computed after the project is evaluated" }
 
+            set(project.provider {
                 val kotlinExtension = project.kotlinExtension
                 val targets = when (kotlinExtension) {
                     is KotlinMultiplatformExtension -> kotlinExtension.targets
@@ -253,21 +254,34 @@ internal object CompilationSourceSetUtil {
 
                 val compilations = targets.flatMap { it.compilations }
 
-            compilations
+                val result = compilations
                     .flatMap { compilation -> compilation.allKotlinSourceSets.map { sourceSet -> compilation to sourceSet } }
                     .groupBy(
-                    { (_, sourceSet) -> sourceSet.name },
+                        { (_, sourceSet) -> sourceSet },
-                    valueTransform = { (compilation, _) -> TargetCompilationName.from(compilation) }
+                        valueTransform = { (compilation, _) -> compilation }
                     )
                     .mapValues { (_, compilations) -> compilations.toSet() }
+
+                if (shouldFinalizeValue) {
+                    set(result)
                 }
 
-        return compilationNamesBySourceSetName.entries.associate { (sourceSetName, compilationNames) ->
+                return@provider result
-            project.kotlinExtension.sourceSets.getByName(sourceSetName).to(
+            })
-                compilationNames.map { checkNotNull(it.toCompilation(project)) }.toSet()
+
-            )
+            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()
+    }
 
     fun sourceSetsInMultipleCompilations(project: Project) =
         compilationsBySourceSets(project).mapNotNullTo(mutableSetOf()) { (sourceSet, compilations) ->