From b5ed2326e77febd71a464b319d021cf98ee1bdac Mon Sep 17 00:00:00 2001
From: Anton Lakotka <anton.lakotka@jetbrains.com>
Date: Tue, 13 Feb 2024 10:24:25 +0100
Subject: [PATCH] [Gradle] Merge Android Source Sets into one K2 Fragment

Android projects has a number of build type/flavor specific source set.
These source sets are included into compilation of multiple variants.
In Android compilation model they should be compiled together and they
should see each-others declarations. With Kotlin MPP plugin applied
these android source sets gets associated with Kotlin Source Sets
counter parts. However in Kotlin Multiplatform Source sets have
a semantics about "expect/actual" refinement and with K2 it also got
stricter and visibility is checked. This is not compatible with Android
compilation model thus the solution is to merge Android Source Sets
into one K2 Fragment.

^KT-62508 Verification Pending
---
 ...nCompilationK2MultiplatformConfigurator.kt | 30 +++++++++++++++----
 1 file changed, 25 insertions(+), 5 deletions(-)

diff --git a/libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/mpp/compilationImpl/KotlinCompilationK2MultiplatformConfigurator.kt b/libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/mpp/compilationImpl/KotlinCompilationK2MultiplatformConfigurator.kt
index 73cf1be2f55..dbaf531cc5b 100644
--- a/libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/mpp/compilationImpl/KotlinCompilationK2MultiplatformConfigurator.kt
+++ b/libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/mpp/compilationImpl/KotlinCompilationK2MultiplatformConfigurator.kt
@@ -6,7 +6,9 @@
 package org.jetbrains.kotlin.gradle.plugin.mpp.compilationImpl
 
 import org.jetbrains.kotlin.gradle.dsl.usesK2
+import org.jetbrains.kotlin.gradle.plugin.KotlinSourceSet
 import org.jetbrains.kotlin.gradle.plugin.mpp.compilationImpl.factory.KotlinCompilationImplFactory
+import org.jetbrains.kotlin.gradle.plugin.sources.android.androidSourceSetInfoOrNull
 import org.jetbrains.kotlin.gradle.tasks.K2MultiplatformCompilationTask
 import org.jetbrains.kotlin.gradle.tasks.K2MultiplatformStructure
 
@@ -16,20 +18,38 @@ internal object KotlinCompilationK2MultiplatformConfigurator : KotlinCompilation
             if (compileTask.name != compilation.compileKotlinTaskName) return@configureEach
             if (compileTask !is K2MultiplatformCompilationTask) return@configureEach
 
+            /**
+             * Returns fragment name of [this]
+             * by default it is name of [KotlinSourceSet] but for android it should name of compilation's default source set.
+             * i.e. all android-specific source sets (fragments) should be combined into one.
+             * See KT-62508 for detailed explanation
+             */
+            fun KotlinSourceSet.fragmentName(): String =
+                if (androidSourceSetInfoOrNull != null) {
+                    compilation.defaultSourceSet.name
+                } else {
+                    name
+                }
+
             compileTask.multiplatformStructure.refinesEdges.set(compilation.project.provider {
                 if (!compileTask.compilerOptions.usesK2.get()) return@provider emptyList()
                 compilation.allKotlinSourceSets.flatMap { sourceSet ->
-                    sourceSet.dependsOn.map { dependsOn ->
-                        K2MultiplatformStructure.RefinesEdge(sourceSet.name, dependsOn.name)
+                    sourceSet.dependsOn.mapNotNull { dependsOn ->
+                        val from = sourceSet.fragmentName()
+                        val to = dependsOn.fragmentName()
+                        if (from == to) return@mapNotNull null
+                        K2MultiplatformStructure.RefinesEdge(from, to)
                     }
                 }
             })
 
             compileTask.multiplatformStructure.fragments.set(compilation.project.provider {
                 if (!compileTask.compilerOptions.usesK2.get()) return@provider emptyList()
-                compilation.allKotlinSourceSets.map { sourceSet ->
-                    K2MultiplatformStructure.Fragment(sourceSet.name, sourceSet.kotlin.asFileTree)
-                }
+                compilation.allKotlinSourceSets
+                    .groupBy(keySelector = { it.fragmentName() }) { sourceSet -> sourceSet.kotlin.asFileTree }
+                    .map { (fragmentName, sourceFiles) ->
+                        K2MultiplatformStructure.Fragment(fragmentName, sourceFiles.reduce { acc, fileTree -> acc + fileTree })
+                    }
             })
         }
     }