Fix unused source sets missing in compilationsBySourceSets (KT-39304)

Issue #KT-39304 Fixed

libraries/tools/kotlin-gradle-plugin-integration-tests/src/test/resources/testProject/hierarchical-mpp-published-modules/my-lib-foo/build.gradle.kts

@@ -91,6 +91,13 @@ kotlin {
         linuxX64().compilations["test"].defaultSourceSet {
             dependsOn(linuxAndJsTest)
         }
+
+        // Test KT-39304
+        val unusedSourceSet by creating {
+            dependencies {
+                implementation(kotlin("stdlib"))
+            }
+        }
     }
 }
 

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

@@ -18,6 +18,7 @@ import org.gradle.api.tasks.TaskProvider
 import org.gradle.api.tasks.TaskState
 import org.gradle.api.tasks.bundling.AbstractArchiveTask
 import org.gradle.util.ConfigureUtil
+import org.jetbrains.kotlin.fir.resolve.getOrPut
 import org.jetbrains.kotlin.gradle.dsl.*
 import org.jetbrains.kotlin.gradle.plugin.*
 import org.jetbrains.kotlin.gradle.plugin.mpp.internal.KotlinCompilationsModuleGroups
@@ -326,13 +327,17 @@ internal object CompilationSourceSetUtil {
 
                 val compilations = targets.flatMap { it.compilations }
 
-                val result = compilations
+                val result = mutableMapOf<KotlinSourceSet, MutableSet<KotlinCompilation<*>>>().apply {
-                    .flatMap { compilation -> compilation.allKotlinSourceSets.map { sourceSet -> compilation to sourceSet } }
+                    compilations.forEach { compilation ->
-                    .groupBy(
+                        compilation.allKotlinSourceSets.forEach { sourceSet ->
-                        { (_, sourceSet) -> sourceSet },
+                            getOrPut(sourceSet) { mutableSetOf() }.add(compilation)
-                        valueTransform = { (compilation, _) -> compilation }
+                        }
-                    )
+                    }
-                    .mapValues { (_, compilations) -> compilations.toSet() }
+                    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)