Optimize SdkInfo lookup for a huge monorepo with no SdkInfo

DFS of SdkInfo on moduleInfo and reuse cached libraryInfo dependencies
fix traverse order

#KT-47065 Fixed
Relates to #KT-46622

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/KotlinStdlibCache.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.idea.caches.project
 
 import com.intellij.openapi.components.ServiceManager
+import com.intellij.openapi.progress.ProgressManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.OrderRootType
 import com.intellij.openapi.vfs.VirtualFile
@@ -112,16 +113,16 @@ class KotlinStdlibCacheImpl(val project: Project) : KotlinStdlibCache {
 
                 // bfs
                 while (stack.isNotEmpty()) {
+                    ProgressManager.checkCanceled()
+
                     val poll = stack.poll()
                     if (!checkedLibraryInfo.add(poll)) continue
 
-                    val dependencies = poll.dependencies().filter { !checkedLibraryInfo.contains(it) }
-                    dependencies
-                        .filterIsInstance<LibraryInfo>()
-                        .firstOrNull { isStdlib(it) }
-                        ?.let { return@run it }
-
-                    stack += dependencies
+                    stack += poll.dependencies().also { dependencies ->
+                        dependencies
+                            .firstOrNull { it is LibraryInfo && isStdlib(it) }
+                            ?.let { return@run it as LibraryInfo }
+                    }
                 }
                 null
             }

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/LibraryDependenciesCache.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.idea.caches.project
 import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.module.Module
 import com.intellij.openapi.module.ModuleManager
+import com.intellij.openapi.progress.ProgressManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.*
 import com.intellij.openapi.roots.impl.libraries.LibraryEx
@@ -94,6 +95,7 @@ class LibraryDependenciesCacheImpl(private val project: Project) : LibraryDepend
         val platform = libraryInfo.platform
 
         for (module in getLibraryUsageIndex().getModulesLibraryIsUsedIn(libraryInfo)) {
+            ProgressManager.checkCanceled()
             val (moduleLibraries, moduleSdks) = moduleDependenciesCache.getOrPut(module) {
                 computeLibrariesAndSdksUsedIn(module)
             }

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/SdkInfoCache.kt

@@ -6,12 +6,14 @@
 package org.jetbrains.kotlin.idea.caches.project
 
 import com.intellij.openapi.components.ServiceManager
-import com.intellij.openapi.diagnostic.Logger
+import com.intellij.openapi.progress.ProgressManager
 import com.intellij.openapi.project.Project
-import com.intellij.util.containers.ConcurrentWeakKeySoftValueHashMap
-import com.intellij.util.containers.ContainerUtil
+import com.jetbrains.rd.util.concurrentMapOf
 import org.jetbrains.kotlin.analyzer.ModuleInfo
 import org.jetbrains.kotlin.caches.project.cacheInvalidatingOnRootModifications
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
+import java.util.ArrayDeque
 
 /**
  * Maintains and caches mapping ModuleInfo -> SdkInfo *form its dependencies*
@@ -35,18 +37,105 @@ interface SdkInfoCache {
 }
 
 class SdkInfoCacheImpl(private val project: Project) : SdkInfoCache {
-    private val cache: MutableMap<ModuleInfo, SdkInfo?>
-        // Note: it's ok to use thread-unsafe map, because 'doFindSdk' is pure
-        get() = project.cacheInvalidatingOnRootModifications { mutableMapOf() }
+    @JvmInline
+    value class SdkDependency(val sdk: SdkInfo?)
+
+    private val cache: MutableMap<ModuleInfo, SdkDependency>
+        get() = project.cacheInvalidatingOnRootModifications { concurrentMapOf() }
 
     override fun findOrGetCachedSdk(moduleInfo: ModuleInfo): SdkInfo? {
-        if (!cache.containsKey(moduleInfo)) {
-            cache[moduleInfo] = doFindSdk(moduleInfo)
+        // get an operate on the fixed instance of a cache to avoid case of roots modification in the middle of lookup
+        val instance = cache
+
+        if (!instance.containsKey(moduleInfo)) {
+            findSdk(instance, moduleInfo)
         }
 
-        return cache[moduleInfo]
+        return instance[moduleInfo]?.sdk
     }
 
-    private fun doFindSdk(moduleInfo: ModuleInfo): SdkInfo? =
-        moduleInfo.dependencies().lazyClosure { it.dependencies() }.firstOrNull { it is SdkInfo } as SdkInfo?
+    private fun findSdk(cache: MutableMap<ModuleInfo, SdkDependency>, moduleInfo: ModuleInfo) {
+        moduleInfo.safeAs<SdkDependency>()?.let {
+            cache[moduleInfo] = it
+            return
+        }
+
+        val libraryDependenciesCache = LibraryDependenciesCache.getInstance(this.project)
+        val visitedModuleInfos = mutableSetOf<ModuleInfo>()
+
+        // graphs is a stack of paths is used to implement DFS without recursion
+        // it depends on a number of libs, that could be > 10k for a huge monorepos
+        val graphs = ArrayDeque<List<ModuleInfo>>().also {
+            // initial graph item
+            it.add(listOf(moduleInfo))
+        }
+
+        val (path, sdkInfo) = run {
+            while (graphs.isNotEmpty()) {
+                ProgressManager.checkCanceled()
+                // graph of DFS from the root i.e from `moduleInfo`
+
+                // note: traverse of graph goes in a depth over the most left edge:
+                // - poll() retrieves and removes the head of the queue
+                // - push(element) inserts the element at the front of the deque
+                val graph = graphs.poll()
+
+                val last = graph.last()
+                // the result could be immediately returned when cache already has it
+                val cached = cache[last]
+                if (cached != null) {
+                    cached.sdk?.let { return@run graph to cached }
+                    continue
+                }
+
+                if (!visitedModuleInfos.add(last)) continue
+
+                val dependencies = run deps@{
+                    if (last is LibraryInfo) {
+                        // use a special case for LibraryInfo to reuse values from a library dependencies cache
+                        val (libraries, sdks) = libraryDependenciesCache.getLibrariesAndSdksUsedWith(last)
+                        sdks.firstOrNull()?.let {
+                            return@run graph to SdkDependency(it)
+                        }
+                        libraries
+                    } else {
+                        last.dependencies()
+                            .also { dependencies ->
+                                dependencies.firstIsInstanceOrNull<SdkInfo>()?.let {
+                                    return@run graph to SdkDependency(it)
+                                }
+                            }
+                    }
+                }
+
+                dependencies.forEach { dependency ->
+                    val sdkDependency = cache[dependency]
+                    if (sdkDependency != null) {
+                        sdkDependency.sdk?.let {
+                            // sdk is found when some dependency is already resolved
+                            return@run (graph + dependency) to sdkDependency
+                        }
+                    } else {
+                        // otherwise add a new graph of (existed graph + dependency) as candidates for DFS lookup
+                        if (!visitedModuleInfos.contains(dependency)) {
+                            graphs.push(graph + dependency)
+                        }
+                    }
+                }
+            }
+            return@run null to noSdkDependency
+        }
+        // when sdk is found: mark all graph elements could be resolved to the same sdk
+        path?.let {
+            it.forEach { info -> cache[info] = sdkInfo }
+
+            visitedModuleInfos.removeAll(it)
+        }
+        // mark all visited modules (apart of found path) as dead ends
+        visitedModuleInfos.forEach { info -> cache[info] = noSdkDependency }
+    }
+
+    companion object {
+        private val noSdkDependency = SdkDependency(null)
+    }
 }