Cache all module dependencies to avoid O(n^2) calc complexity

Quite noticeable when there is a big module that has 100+ libraries and many modules depends on it

#KT-46622 Fixed

(cherry picked from commit b014787)

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

@@ -17,7 +17,8 @@ import org.jetbrains.kotlin.idea.configuration.IdeBuiltInsLoadingState
 import org.jetbrains.kotlin.idea.util.application.runReadAction
 import org.jetbrains.kotlin.idea.vfilefinder.KotlinStdlibIndex
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.platform.jvm.isJvm
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
+import java.util.*
 import java.util.concurrent.ConcurrentHashMap
 
 // TODO(kirpichenkov): works only for JVM (see KT-44552)
@@ -101,9 +102,30 @@ class KotlinStdlibCacheImpl(val project: Project) : KotlinStdlibCache {
 
     override fun findStdlibInModuleDependencies(module: IdeaModuleInfo): LibraryInfo? {
         val stdlibDependency = moduleStdlibDependencyCache.getOrPut(module) {
-            module.dependencies().lazyClosure { it.dependencies() }.firstOrNull {
-                it is LibraryInfo && isStdlib(it)
-            }.let { StdlibDependency(it as? LibraryInfo?) }
+
+            fun IdeaModuleInfo.asStdLibInfo() = this.safeAs<LibraryInfo>()?.takeIf { isStdlib(it) }
+
+            val stdLib: LibraryInfo? = module.asStdLibInfo() ?: run {
+                val checkedLibraryInfo = mutableSetOf<IdeaModuleInfo>()
+                val stack = ArrayDeque<IdeaModuleInfo>()
+                stack.add(module)
+
+                // bfs
+                while (stack.isNotEmpty()) {
+                    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
+                }
+                null
+            }
+            StdlibDependency(stdLib)
         }
 
         return stdlibDependency.libraryInfo

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

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.idea.core.util.CachedValue
 import org.jetbrains.kotlin.idea.core.util.getValue
 import org.jetbrains.kotlin.platform.TargetPlatform
 import org.jetbrains.kotlin.platform.isCommon
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import java.util.*
 
 internal typealias LibrariesAndSdks = Pair<List<LibraryInfo>, List<SdkInfo>>
@@ -41,55 +42,62 @@ class LibraryDependenciesCacheImpl(private val project: Project) : LibraryDepend
         )
     }
 
+    private val moduleDependenciesCache by CachedValue(project) {
+        CachedValueProvider.Result(
+            ContainerUtil.createConcurrentWeakMap<Module, LibrariesAndSdks>(),
+            ProjectRootManager.getInstance(project)
+        )
+    }
+
     override fun getLibrariesAndSdksUsedWith(libraryInfo: LibraryInfo): LibrariesAndSdks =
         cache.getOrPut(libraryInfo) { computeLibrariesAndSdksUsedWith(libraryInfo) }
 
+    private fun computeLibrariesAndSdksUsedIn(module: Module): LibrariesAndSdks {
+        val libraries = LinkedHashSet<LibraryInfo>()
+        val sdks = LinkedHashSet<SdkInfo>()
 
-    //NOTE: used LibraryRuntimeClasspathScope as reference
-    private fun computeLibrariesAndSdksUsedWith(libraryInfo: LibraryInfo): LibrariesAndSdks {
         val processedModules = HashSet<Module>()
         val condition = Condition<OrderEntry> { orderEntry ->
-            if (orderEntry is ModuleOrderEntry) {
-                val module = orderEntry.module
-                module != null && module !in processedModules
-            } else {
-                true
-            }
+            orderEntry.safeAs<ModuleOrderEntry>()?.let {
+                it.module?.run { this !in processedModules } ?: false
+            } ?: true
         }
 
+        ModuleRootManager.getInstance(module).orderEntries().recursively().satisfying(condition).process(object : RootPolicy<Unit>() {
+            override fun visitModuleSourceOrderEntry(moduleSourceOrderEntry: ModuleSourceOrderEntry, value: Unit) {
+                processedModules.add(moduleSourceOrderEntry.ownerModule)
+            }
+
+            override fun visitLibraryOrderEntry(libraryOrderEntry: LibraryOrderEntry, value: Unit) {
+                libraryOrderEntry.library.safeAs<LibraryEx>()?.takeIf { !it.isDisposed }?.let {
+                    libraries += createLibraryInfo(project, it)
+                }
+            }
+
+            override fun visitJdkOrderEntry(jdkOrderEntry: JdkOrderEntry, value: Unit) {
+                jdkOrderEntry.jdk?.let { jdk ->
+                    sdks += SdkInfo(project, jdk)
+                }
+            }
+        }, Unit)
+
+        return Pair(libraries.toList(), sdks.toList())
+    }
+
+    //NOTE: used LibraryRuntimeClasspathScope as reference
+    private fun computeLibrariesAndSdksUsedWith(libraryInfo: LibraryInfo): LibrariesAndSdks {
         val libraries = LinkedHashSet<LibraryInfo>()
         val sdks = LinkedHashSet<SdkInfo>()
 
         val platform = libraryInfo.platform
 
-        for (module in getLibraryUsageIndex().modulesLibraryIsUsedIn[libraryInfo.library]) {
-            if (!processedModules.add(module)) continue
-
-            ModuleRootManager.getInstance(module).orderEntries().recursively().satisfying(condition).process(object : RootPolicy<Unit>() {
-                override fun visitModuleSourceOrderEntry(moduleSourceOrderEntry: ModuleSourceOrderEntry, value: Unit) {
-                    processedModules.add(moduleSourceOrderEntry.ownerModule)
-                }
-
-                override fun visitLibraryOrderEntry(libraryOrderEntry: LibraryOrderEntry, value: Unit) {
-                    val otherLibrary = libraryOrderEntry.library
-                    if (otherLibrary is LibraryEx && !otherLibrary.isDisposed) {
-                        val otherLibraryInfos = createLibraryInfo(project, otherLibrary)
-                        otherLibraryInfos.firstOrNull()?.let { otherLibraryInfo ->
-                            if (compatiblePlatforms(platform, otherLibraryInfo.platform)) {
-                                libraries.addAll(otherLibraryInfos)
-                            }
-                        }
-                    }
-                }
+        for (module in getLibraryUsageIndex().getModulesLibraryIsUsedIn(libraryInfo)) {
+            val (moduleLibraries, moduleSdks) = moduleDependenciesCache.getOrPut(module) {
+                computeLibrariesAndSdksUsedIn(module)
+            }
 
-                override fun visitJdkOrderEntry(jdkOrderEntry: JdkOrderEntry, value: Unit) {
-                    val jdk = jdkOrderEntry.jdk ?: return
-                    SdkInfo(project, jdk).also { sdkInfo ->
-                        if (compatiblePlatforms(platform, sdkInfo.platform))
-                            sdks += sdkInfo
-                    }
-                }
-            }, Unit)
+            moduleLibraries.filter { compatiblePlatforms(platform, it.platform) }.forEach { libraries.add(it) }
+            moduleSdks.filter { compatiblePlatforms(platform, it.platform) }.forEach { sdks.add(it) }
         }
 
         return Pair(libraries.toList(), sdks.toList())
@@ -123,5 +131,8 @@ class LibraryDependenciesCacheImpl(private val project: Project) : LibraryDepend
                 }
             }
         }
+
+        fun getModulesLibraryIsUsedIn(libraryInfo: LibraryInfo): Collection<Module> =
+            modulesLibraryIsUsedIn[libraryInfo.library]
     }
 }