From ea24a3c60276aa7ac3f871873a9aabc61a899ac3 Mon Sep 17 00:00:00 2001
From: Simon Ogorodnik <Simon.Ogorodnik@jetbrains.com>
Date: Fri, 21 Dec 2018 16:27:27 +0300
Subject: [PATCH] KT-28898: Optimize type-alias extensions and type-alias
 expansions lookup

 #KT-28898 Fixed
---
 .../kotlin/idea/core/KotlinIndicesHelper.kt   | 32 ++++++++++++++-----
 1 file changed, 24 insertions(+), 8 deletions(-)

diff --git a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/KotlinIndicesHelper.kt b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/KotlinIndicesHelper.kt
index a0dc5214b91..48e3698ca86 100644
--- a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/KotlinIndicesHelper.kt
+++ b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/KotlinIndicesHelper.kt
@@ -27,7 +27,6 @@ import org.jetbrains.kotlin.asJava.elements.KtLightElement
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.idea.caches.KotlinShortNamesCache
 import org.jetbrains.kotlin.idea.caches.resolve.resolveImportReference
-import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptorIfAny
 import org.jetbrains.kotlin.idea.caches.resolve.unsafeResolveToDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.util.getJavaMemberDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.util.resolveToDescriptor
@@ -190,23 +189,40 @@ class KotlinIndicesHelper(
         fun searchRecursively(typeName: String) {
             ProgressManager.checkCanceled()
             index[typeName, project, scope].asSequence()
-                    .map { it.resolveToDescriptorIfAny() as? TypeAliasDescriptor }
-                    .filterNotNull()
-                    .filter { it.expandedType.constructor == typeConstructor }
-                    .filter { out.putIfAbsent(it.fqNameSafe, it) == null }
-                    .map { it.name.asString() }
-                    .forEach(::searchRecursively)
+                .filter { it in scope }
+                .flatMap { it.resolveToDescriptors<TypeAliasDescriptor>().asSequence() }
+                .filter { it.expandedType.constructor == typeConstructor }
+                .filter { out.putIfAbsent(it.fqNameSafe, it) == null }
+                .map { it.name.asString() }
+                .forEach(::searchRecursively)
         }
 
         searchRecursively(originalTypeName)
         return out.values.toSet()
     }
 
+    private fun possibleTypeAliasExpansionNames(originalTypeName: String): Set<String> {
+        val index = KotlinTypeAliasByExpansionShortNameIndex.INSTANCE
+        val out = mutableSetOf<String>()
+
+        fun searchRecursively(typeName: String) {
+            ProgressManager.checkCanceled()
+            index[typeName, project, scope].asSequence()
+                .filter { it in scope }
+                .mapNotNull { it.name }
+                .filter { out.add(it) }
+                .forEach(::searchRecursively)
+        }
+
+        searchRecursively(originalTypeName)
+        return out
+    }
+
     private fun MutableCollection<String>.addTypeNames(type: KotlinType) {
         val constructor = type.constructor
         constructor.declarationDescriptor?.name?.asString()?.let { typeName ->
             add(typeName)
-            resolveTypeAliasesUsingIndex(type, typeName).mapTo(this, { it.name.asString() })
+            addAll(possibleTypeAliasExpansionNames(typeName))
         }
         constructor.supertypes.forEach { addTypeNames(it) }
     }