From 8c99b0e64c545bd0850512595cef3c6614ee0c41 Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Tue, 25 Oct 2016 19:37:57 +0300
Subject: [PATCH] Performance fix for Import Member

---
 .../kotlin/idea/core/KotlinIndicesHelper.kt   | 49 ++++++++++++-------
 .../kotlin/idea/quickfix/ImportFix.kt         | 24 ++++++---
 2 files changed, 48 insertions(+), 25 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 7d5db799e66..2727c4e53ad 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
@@ -206,28 +206,41 @@ class KotlinIndicesHelper(
                 .toSet()
     }
 
-    fun getJvmCallablesByName(name: String): Collection<CallableDescriptor> {
+    fun processJvmCallablesByName(
+            name: String,
+            filter: (PsiMember) -> Boolean,
+            processor: (CallableDescriptor) -> Unit
+    ) {
         val javaDeclarations = PsiShortNamesCache.getInstance(project).getFieldsByName(name, scopeWithoutKotlin).asSequence() +
                                PsiShortNamesCache.getInstance(project).getMethodsByName(name, scopeWithoutKotlin).asSequence()
-        return javaDeclarations
-                .filterNot { it is KtLightElement<*,*> }
-                .mapNotNull { (it as PsiMember).getJavaMemberDescriptor(resolutionFacade) as? CallableDescriptor }
-                .filter(descriptorFilter)
-                .toSet()
+        val processed = HashSet<CallableDescriptor>()
+        for (javaDeclaration in javaDeclarations) {
+            ProgressManager.checkCanceled()
+            if (javaDeclaration is KtLightElement<*, *>) continue
+            if (!filter(javaDeclaration as PsiMember)) continue
+            val descriptor = javaDeclaration.getJavaMemberDescriptor(resolutionFacade) as? CallableDescriptor ?: continue
+            if (!processed.add(descriptor)) continue
+            if (!descriptorFilter(descriptor)) continue
+            processor(descriptor)
+        }
     }
 
-    fun getKotlinCallablesByName(name: String): Collection<CallableDescriptor> {
-        val functions = KotlinFunctionShortNameIndex.getInstance().get(name, project, scope)
-                .asSequence()
-                .map { it.descriptor as? CallableDescriptor }
-        val properties = KotlinPropertyShortNameIndex.getInstance().get(name, project, scope)
-                .asSequence()
-                .map { it.descriptor as? CallableDescriptor }
-
-        return (functions + properties)
-                .filterNotNull()
-                .filter(descriptorFilter)
-                .toSet()
+    fun processKotlinCallablesByName(
+            name: String,
+            filter: (KtCallableDeclaration) -> Boolean,
+            processor: (CallableDescriptor) -> Unit
+    ) {
+        val functions: Sequence<KtCallableDeclaration> = KotlinFunctionShortNameIndex.getInstance().get(name, project, scope).asSequence()
+        val properties: Sequence<KtCallableDeclaration> = KotlinPropertyShortNameIndex.getInstance().get(name, project, scope).asSequence()
+        val processed = HashSet<CallableDescriptor>()
+        for (declaration in functions + properties) {
+            ProgressManager.checkCanceled()
+            if (!filter(declaration)) continue
+            val descriptor = declaration.descriptor as? CallableDescriptor ?: continue
+            if (!processed.add(descriptor)) continue
+            if (!descriptorFilter(descriptor)) continue
+            processor(descriptor)
+        }
     }
 
     fun getKotlinClasses(nameFilter: (String) -> Boolean, kindFilter: (ClassKind) -> Boolean): Collection<ClassDescriptor> {
diff --git a/idea/src/org/jetbrains/kotlin/idea/quickfix/ImportFix.kt b/idea/src/org/jetbrains/kotlin/idea/quickfix/ImportFix.kt
index 8de46b9a841..c6030460da7 100644
--- a/idea/src/org/jetbrains/kotlin/idea/quickfix/ImportFix.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/quickfix/ImportFix.kt
@@ -28,6 +28,7 @@ import com.intellij.packageDependencies.DependencyValidationManager
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiErrorElement
 import com.intellij.psi.PsiFile
+import com.intellij.psi.PsiModifier
 import com.intellij.psi.util.PsiModificationTracker
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
@@ -62,7 +63,6 @@ import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelectorOrThis
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.isImportDirectiveExpression
 import org.jetbrains.kotlin.resolve.BindingContext
-import org.jetbrains.kotlin.resolve.DescriptorUtils.isTopLevelDeclaration
 import org.jetbrains.kotlin.resolve.bindingContextUtil.getDataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.callUtil.getParentCall
 import org.jetbrains.kotlin.resolve.calls.context.ContextDependency
@@ -393,14 +393,24 @@ internal class ImportMemberFix(expression: KtSimpleNameExpression) : ImportFixBa
         if (!element.isImportDirectiveExpression() && !isSelectorInQualified(element)) {
             val filterByCallType = { descriptor: DeclarationDescriptor -> callTypeAndReceiver.callType.descriptorKindFilter.accepts(descriptor) }
 
-            indicesHelper.getKotlinCallablesByName(name)
-                    .filter { it.canBeReferencedViaImport() && !isTopLevelDeclaration(it) }
-                    .filterTo(result, filterByCallType)
+            val processor = { descriptor: CallableDescriptor ->
+                if (descriptor.canBeReferencedViaImport() && filterByCallType(descriptor)) {
+                    result.add(descriptor)
+                }
+            }
+
+            indicesHelper.processKotlinCallablesByName(
+                    name,
+                    filter = { declaration -> (declaration.parent as? KtClassBody)?.parent is KtObjectDeclaration },
+                    processor = processor
+            )
 
             if (!ProjectStructureUtil.isJsKotlinModule(element.getContainingKtFile())) {
-                indicesHelper.getJvmCallablesByName(name)
-                        .filter { it.canBeReferencedViaImport() }
-                        .filterTo(result, filterByCallType)
+                indicesHelper.processJvmCallablesByName(
+                        name,
+                        filter = { it.hasModifierProperty(PsiModifier.STATIC) },
+                        processor = processor
+                )
             }
         }
         return result