From 6deb679e5c5b534bdfdfcb6df5a8fc4364365a8e Mon Sep 17 00:00:00 2001
From: Pavel Kirpichenkov <pavel.kirpichenkov@jetbrains.com>
Date: Tue, 16 Jan 2024 20:22:38 +0200
Subject: [PATCH] [LL] Fix incorrect KMP groups in case of reordered
 dependencies

KTIJ-27569
---
 .../api/fir/providers/KmpModuleSorter.kt      | 21 +++++++++-----
 .../low/level/api/fir/KmpModuleSorterTest.kt  | 28 +++++++++++++++++++
 2 files changed, 42 insertions(+), 7 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/KmpModuleSorter.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/KmpModuleSorter.kt
index 4edea4ff0e8..8959046188c 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/KmpModuleSorter.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/KmpModuleSorter.kt
@@ -39,10 +39,9 @@ class KmpModuleSorter private constructor(private val modules: List<KtModule>) {
     private fun groupModules() {
         for ((index, module) in modules.withIndex()) {
             originalPositions[module] = index
-            val group = findOrCreateKmpGroup(module, groupsByModules)
+            val group = findOrCreateRootKmpGroup(module)
             group.addModule(module)
             groupsByModules.putIfAbsent(module, group)
-            module.directDependsOnDependencies.forEach { dependency -> groupsByModules.putIfAbsent(dependency, group) }
         }
     }
 
@@ -60,13 +59,21 @@ class KmpModuleSorter private constructor(private val modules: List<KtModule>) {
         return sortedModules.toList().filterNotNull()
     }
 
-    private fun findOrCreateKmpGroup(
+    // The group of the root source set, e.g., commonMain
+    private fun findOrCreateRootKmpGroup(
         module: KtModule,
-        groups: MutableMap<KtModule, KmpGroup>,
     ): KmpGroup {
-        groups[module]?.let { return it }
-        module.directDependsOnDependencies.firstNotNullOfOrNull { dependency -> groups[dependency] }?.let { return it }
-        return KmpGroup()
+        groupsByModules[module]?.let { return it }
+
+        val root = if (module.directDependsOnDependencies.isEmpty()) module
+        else module.transitiveDependsOnDependencies.singleOrNull { dependency -> dependency.directDependsOnDependencies.isEmpty() }
+
+        if (root == null) {
+            hasErrors = true
+            return KmpGroup()
+        }
+
+        return groupsByModules.getOrPut(root) { KmpGroup() }
     }
 
     private fun getOriginalPositionOrSetCorrupted(module: KtModule): Int {
diff --git a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/KmpModuleSorterTest.kt b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/KmpModuleSorterTest.kt
index cf7b7eae4df..e90344f5780 100644
--- a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/KmpModuleSorterTest.kt
+++ b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/KmpModuleSorterTest.kt
@@ -76,6 +76,34 @@ class KmpModuleSorterTest : TestWithMockProject() {
         assertEquals(listOf(p1Platform, p1Intermediate, p1Common, p2Intermediate, p2Common), this.buildDependenciesToTest(p2Platform))
     }
 
+    @Test
+    fun testPartsOfTheGroupAreMergedCorrectly1() {
+        val m1 = createKtModule("m1")
+        val m2 = createKtModule("m2", directDependsOnDependencies = listOf(m1))
+        val m3 = createKtModule("m3", directDependsOnDependencies = listOf(m2))
+        val m4 = createKtModule("m4", directDependsOnDependencies = listOf(m3))
+
+        val c = createKtModule("c", directRegularDependencies = listOf(m1, m3, m4, m2))
+        val d = createKtModule("d", directRegularDependencies = listOf(m2, m4, m3, m1))
+        assertEquals(listOf(m4, m3, m2, m1), buildDependenciesToTest(c))
+        assertEquals(listOf(m4, m3, m2, m1), buildDependenciesToTest(d))
+    }
+
+    @Test
+    fun testPartsOfTheGroupAreMergedCorrectly2() {
+        val m1 = createKtModule("m1")
+        val m2 = createKtModule("m2", directDependsOnDependencies = listOf(m1))
+        val m3 = createKtModule("m3", directDependsOnDependencies = listOf(m2))
+        val m4 = createKtModule("m4", directDependsOnDependencies = listOf(m3))
+        val m5 = createKtModule("m5", directDependsOnDependencies = listOf(m4))
+        val m6 = createKtModule("m6", directDependsOnDependencies = listOf(m5))
+
+        val c = createKtModule("c", directRegularDependencies = listOf(m3, m4, m6, m5, m2, m1))
+        val d = createKtModule("d", directRegularDependencies = listOf(m1, m2, m5, m6, m4, m3))
+        assertEquals(listOf(m6, m5, m4, m3, m2, m1), buildDependenciesToTest(c))
+        assertEquals(listOf(m6, m5, m4, m3, m2, m1), buildDependenciesToTest(d))
+    }
+
     // See [org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirAbstractSessionFactory#collectDependencySymbolProviders]
     private fun buildDependenciesToTest(module: KtModule): List<KtModule> {
         val dependenciesToSort = buildSet {