From 98e3813af9e6379f560fc1599e0463d3a2820425 Mon Sep 17 00:00:00 2001
From: Pavel Kirpichenkov <pavel.kirpichenkov@jetbrains.com>
Date: Thu, 11 Jan 2024 13:24:52 +0200
Subject: [PATCH] [LL] Add error recovery to KmpModuleSorter

KTIJ-27569
---
 .../api/fir/providers/KmpModuleSorter.kt      | 34 +++++++++++--------
 1 file changed, 19 insertions(+), 15 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 40411fe9ec9..4edea4ff0e8 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
@@ -26,9 +26,14 @@ class KmpModuleSorter private constructor(private val modules: List<KtModule>) {
     private val groupsByModules = mutableMapOf<KtModule, KmpGroup>()
     private val originalPositions = mutableMapOf<KtModule, Int>()
 
+    // Signals corrupted state. When set, original positions will be used for recovery
+    var hasErrors = false
+        private set
+
     private fun sort(): List<KtModule> {
         groupModules()
-        return sortModules()
+        val sorted = sortModules()
+        return if (hasErrors) modules else sorted
     }
 
     private fun groupModules() {
@@ -64,6 +69,12 @@ class KmpModuleSorter private constructor(private val modules: List<KtModule>) {
         return KmpGroup()
     }
 
+    private fun getOriginalPositionOrSetCorrupted(module: KtModule): Int {
+        val originalPosition = originalPositions[module]
+        if (originalPosition == null) hasErrors = true
+        return originalPosition ?: 0
+    }
+
     /**
      * Modules, corresponding to source sets of the same KMP project.
      */
@@ -71,7 +82,7 @@ class KmpModuleSorter private constructor(private val modules: List<KtModule>) {
         private val modules = mutableListOf<KtModule>()
         private val sortedModules by lazy {
             topologicalSort(modules) { directDependsOnDependencies }.also {
-                check(it.size == modules.size) { "The number of sorted modules doesn't match the number of registered modules" }
+                if (it.size != modules.size) hasErrors = true
             }
         }
 
@@ -83,11 +94,12 @@ class KmpModuleSorter private constructor(private val modules: List<KtModule>) {
 
         fun getUpdatedIndexOf(module: KtModule): Int {
             check(module in modules)
-            if (modules.size == 1) return originalPositions[module]
-                ?: error("Can't find position for module: $module")
+            if (modules.size == 1) return getOriginalPositionOrSetCorrupted(module)
 
-            return oldReplacedModulesBySortedModules[module]?.let { originalPositions[it] }
-                ?: error("Can't find position for module: $module")
+            val newPosition = oldReplacedModulesBySortedModules[module]?.let { getOriginalPositionOrSetCorrupted(it) }
+            if (newPosition == null) hasErrors = true
+
+            return newPosition ?: 0
         }
 
         // N.B.: evaluating debug text before all modules are registered will corrupt the group
@@ -101,15 +113,7 @@ class KmpModuleSorter private constructor(private val modules: List<KtModule>) {
     companion object {
         fun order(modules: List<KtModule>): List<KtModule> {
             if (modules.size < 2) return modules
-
-            val sorter = KmpModuleSorter(modules)
-            val sortedModules = sorter.sort()
-
-            check(sortedModules.size == modules.size) {
-                "The resulting number of modules (${sortedModules.size}) doesn't match the number of input modules ($modules.size)"
-            }
-
-            return sortedModules
+            return KmpModuleSorter(modules).sort()
         }
     }
 }