From f7ed3139ab1857580ec77139c9f47753e1bd01c2 Mon Sep 17 00:00:00 2001
From: Sergey Igushkin <sergey.igushkin@jetbrains.com>
Date: Tue, 2 Jun 2020 00:48:55 +0300
Subject: [PATCH] Fix unused source sets missing in compilationsBySourceSets
 (KT-39304)

Issue #KT-39304 Fixed
---
 .../my-lib-foo/build.gradle.kts               |  7 +++++++
 .../gradle/plugin/mpp/kotlinCompilations.kt   | 19 ++++++++++++-------
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/libraries/tools/kotlin-gradle-plugin-integration-tests/src/test/resources/testProject/hierarchical-mpp-published-modules/my-lib-foo/build.gradle.kts b/libraries/tools/kotlin-gradle-plugin-integration-tests/src/test/resources/testProject/hierarchical-mpp-published-modules/my-lib-foo/build.gradle.kts
index e8dfda51a79..9c923b74e77 100644
--- a/libraries/tools/kotlin-gradle-plugin-integration-tests/src/test/resources/testProject/hierarchical-mpp-published-modules/my-lib-foo/build.gradle.kts
+++ b/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"))
+            }
+        }
     }
 }
 
diff --git a/libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/plugin/mpp/kotlinCompilations.kt b/libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/plugin/mpp/kotlinCompilations.kt
index 779b24c6e3d..dc01fb29276 100644
--- a/libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/plugin/mpp/kotlinCompilations.kt
+++ b/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
-                    .flatMap { compilation -> compilation.allKotlinSourceSets.map { sourceSet -> compilation to sourceSet } }
-                    .groupBy(
-                        { (_, sourceSet) -> sourceSet },
-                        valueTransform = { (compilation, _) -> compilation }
-                    )
-                    .mapValues { (_, compilations) -> compilations.toSet() }
+                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)