Performance optimization in completion: do not analyze all extensions available from imports - use index

core/descriptors/src/org/jetbrains/kotlin/resolve/scopes/KtScope.kt

@@ -176,9 +176,10 @@ public class DescriptorKindFilter(
         public val ALL_KINDS_MASK: Int = 0x1F
         public val CLASSIFIERS_MASK: Int = NON_SINGLETON_CLASSIFIERS_MASK or SINGLETON_CLASSIFIERS_MASK
         public val VALUES_MASK: Int = SINGLETON_CLASSIFIERS_MASK or FUNCTIONS_MASK or VARIABLES_MASK
+        public val CALLABLES_MASK: Int = FUNCTIONS_MASK or VARIABLES_MASK
 
         public val ALL: DescriptorKindFilter = DescriptorKindFilter(ALL_KINDS_MASK)
-        public val CALLABLES: DescriptorKindFilter = DescriptorKindFilter(FUNCTIONS_MASK or VARIABLES_MASK)
+        public val CALLABLES: DescriptorKindFilter = DescriptorKindFilter(CALLABLES_MASK)
         public val NON_SINGLETON_CLASSIFIERS: DescriptorKindFilter = DescriptorKindFilter(NON_SINGLETON_CLASSIFIERS_MASK)
         public val SINGLETON_CLASSIFIERS: DescriptorKindFilter = DescriptorKindFilter(SINGLETON_CLASSIFIERS_MASK)
         public val CLASSIFIERS: DescriptorKindFilter = DescriptorKindFilter(CLASSIFIERS_MASK)

idea/ide-common/src/org/jetbrains/kotlin/idea/codeInsight/ReferenceVariantsHelper.kt

@@ -304,7 +304,7 @@ class ReferenceVariantsHelper(
         if (kindFilter.excludes.contains(DescriptorKindExclude.Extensions)) return
 
         fun process(extension: CallableDescriptor) {
-            if (nameFilter(extension.name) && kindFilter.accepts(extension)) {
+            if (kindFilter.accepts(extension) && nameFilter(extension.name)) {
                 addAll(extension.substituteExtensionIfCallable(receiverTypes, callType))
             }
         }

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/BasicCompletionSession.kt

@@ -72,12 +72,18 @@ class BasicCompletionSession(configuration: CompletionSessionConfiguration,
 
     private val completionKind = detectCompletionKind()
 
-    override val descriptorKindFilter = if (isNoQualifierContext()) {
+    override val descriptorKindFilter: DescriptorKindFilter? by lazy {
-        // it's an optimization because obtaining top-level packages from scope is very slow, we obtains them in other way
+        var filter = completionKind.descriptorKindFilter ?: return@lazy null
-        completionKind.descriptorKindFilter?.exclude(DescriptorKindExclude.TopLevelPackages)
+
-    }
+        if (filter.kindMask.and(DescriptorKindFilter.PACKAGES_MASK) != 0) {
-    else {
+            filter = filter.exclude(DescriptorKindExclude.TopLevelPackages)
-        completionKind.descriptorKindFilter
+        }
+
+        if (callTypeAndReceiver.shouldCompleteCallableExtensions()) {
+            filter = filter.exclude(topLevelExtensionsExclude) // completed via indices
+        }
+
+        filter
     }
 
     private val parameterNameAndTypeCompletion = if (shouldCompleteParameterNameAndType())
@@ -208,7 +214,7 @@ class BasicCompletionSession(configuration: CompletionSessionConfiguration,
             completeKeywords()
 
             // getting root packages from scope is very slow so we do this in alternative way
-            if (isNoQualifierContext() && (descriptorKindFilter?.kindMask ?: 0).and(DescriptorKindFilter.PACKAGES_MASK) != 0) {
+            if (callTypeAndReceiver.receiver == null && (descriptorKindFilter?.kindMask ?: 0).and(DescriptorKindFilter.PACKAGES_MASK) != 0) {
                 //TODO: move this code somewhere else?
                 val packageNames = PackageIndexUtil.getSubPackageFqNames(FqName.ROOT, originalSearchScope, project, prefixMatcher.asNameFilter())
                         .toMutableSet()
@@ -225,13 +231,16 @@ class BasicCompletionSession(configuration: CompletionSessionConfiguration,
                 packageNames.forEach { collector.addElement(lookupElementFactory.createLookupElementForPackage(it)) }
             }
 
-            if (completionKind != CompletionKind.KEYWORDS_ONLY) {
+            flushToResultSet()
-                flushToResultSet()
 
-                if (!configuration.completeNonImportedDeclarations && isNoQualifierContext()) {
+            if (completionKind == CompletionKind.ALL && callTypeAndReceiver.shouldCompleteCallableExtensions()) {
-                    collector.advertiseSecondCompletion()
+                for (extension in getCallableTopLevelExtensions()) {
+                    collector.addDescriptorElements(extension, !isImportableDescriptorImported(extension))
                 }
+                flushToResultSet()
+            }
 
+            if (completionKind != CompletionKind.KEYWORDS_ONLY) {
                 completeNonImported()
 
                 if (position.getContainingFile() is KtCodeFragment) {
@@ -367,23 +376,22 @@ class BasicCompletionSession(configuration: CompletionSessionConfiguration,
     }
 
     private fun completeNonImported() {
-        if (completionKind == CompletionKind.ALL) {
+        if (shouldCompleteTopLevelCallablesFromIndex()) {
-            collector.addDescriptorElements(getTopLevelExtensions(), notImported = true)
+            collector.addDescriptorElements(getTopLevelCallables(), notImported = true)
         }
 
-        flushToResultSet()
+        val classKindFilter: ((ClassKind) -> Boolean)?
-
+        when (callTypeAndReceiver) {
-        if (shouldRunTopLevelCompletion()) {
+            is CallTypeAndReceiver.ANNOTATION -> classKindFilter = { it == ClassKind.ANNOTATION_CLASS }
-            if (completionKind == CompletionKind.ALL) {
+            is CallTypeAndReceiver.DEFAULT, is CallTypeAndReceiver.TYPE -> classKindFilter = { it != ClassKind.ENUM_ENTRY }
-                val classKindFilter: ((ClassKind) -> Boolean)?
+            else -> classKindFilter = null
-                when (callTypeAndReceiver) {
+        }
-                    is CallTypeAndReceiver.ANNOTATION -> classKindFilter = { it == ClassKind.ANNOTATION_CLASS }
+        if (classKindFilter != null) {
-                    is CallTypeAndReceiver.DEFAULT, is CallTypeAndReceiver.TYPE -> classKindFilter = { it != ClassKind.ENUM_ENTRY }
+            if (configuration.completeNonImportedDeclarations) {
-                    else -> classKindFilter = null
+                addClassesFromIndex(classKindFilter)
-                }
+            }
-                classKindFilter?.let { addAllClasses(it) }
+            else {
-
+                collector.advertiseSecondCompletion()
-                collector.addDescriptorElements(getTopLevelCallables(), notImported = true)
             }
         }
 
@@ -402,7 +410,7 @@ class BasicCompletionSession(configuration: CompletionSessionConfiguration,
             superClasses += classDescriptor.builtIns.any
         }
 
-        if (!isNoQualifierContext()) {
+        if (callTypeAndReceiver.receiver != null) {
             val referenceVariantsSet = referenceVariants.toSet()
             superClasses = superClasses.filter { it in referenceVariantsSet }
         }

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/CompletionSession.kt

@@ -29,6 +29,7 @@ import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
 import org.jetbrains.kotlin.idea.caches.resolve.getResolveScope
+import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
 import org.jetbrains.kotlin.idea.codeInsight.DescriptorToSourceUtilsIde
 import org.jetbrains.kotlin.idea.codeInsight.ReferenceVariantsHelper
 import org.jetbrains.kotlin.idea.core.ImportableFqNameClassifier
@@ -41,9 +42,12 @@ import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.idea.util.*
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.isExtensionDeclaration
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.descriptorUtil.isExtension
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
+import org.jetbrains.kotlin.resolve.scopes.DescriptorKindExclude
 import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.typeUtil.makeNotNullable
@@ -152,6 +156,20 @@ abstract class CompletionSession(protected val configuration: CompletionSessionC
     protected val toFromOriginalFileMapper: ToFromOriginalFileMapper
             = ToFromOriginalFileMapper(parameters.originalFile as KtFile, position.containingFile as KtFile, parameters.offset)
 
+    // excludes top-level extensions except for ones declared in the current file - those that are fetched from indices
+    protected val topLevelExtensionsExclude = object : DescriptorKindExclude() {
+        val extensionsFromThisFile = file.declarations
+                .filter { it.isExtensionDeclaration() }
+                .map { it.resolveToDescriptor() }
+                .toSet()
+
+        override fun excludes(descriptor: DeclarationDescriptor)
+                = descriptor.isExtension && descriptor.containingDeclaration is PackageFragmentDescriptor && descriptor !in extensionsFromThisFile
+
+        override val fullyExcludedDescriptorKinds: Int
+            get() = 0
+    }
+
     private fun isVisibleDescriptor(descriptor: DeclarationDescriptor): Boolean {
         if (!configuration.completeJavaClassesNotToBeUsed && descriptor is ClassDescriptor) {
             val classification = descriptor.importableFqName?.let { importableFqNameClassifier.classify(it, isPackage = false) }
@@ -206,7 +224,7 @@ abstract class CompletionSession(protected val configuration: CompletionSessionC
 
     protected abstract val expectedInfos: Collection<ExpectedInfo>
 
-    private val importableFqNameClassifier = ImportableFqNameClassifier(file)
+    protected val importableFqNameClassifier = ImportableFqNameClassifier(file)
 
     protected open fun createSorter(): CompletionSorter {
         var sorter = CompletionSorter.defaultSorter(parameters, prefixMatcher)!!
@@ -282,12 +300,11 @@ abstract class CompletionSession(protected val configuration: CompletionSessionC
         }
     }
 
-    protected fun shouldRunTopLevelCompletion(): Boolean
+    protected fun shouldCompleteTopLevelCallablesFromIndex(): Boolean {
-            = configuration.completeNonImportedDeclarations && isNoQualifierContext()
+        if (!configuration.completeNonImportedDeclarations) return false
-
+        if ((descriptorKindFilter?.kindMask ?: 0).and(DescriptorKindFilter.CALLABLES_MASK) == 0) return false
-    protected fun isNoQualifierContext(): Boolean {
+        if (callTypeAndReceiver is CallTypeAndReceiver.IMPORT_DIRECTIVE) return false
-        val parent = position.getParent()
+        return callTypeAndReceiver.receiver == null
-        return parent is KtSimpleNameExpression && !KtPsiUtil.isSelectorInQualified(parent)
     }
 
     protected fun getTopLevelCallables(): Collection<DeclarationDescriptor> {
@@ -295,9 +312,14 @@ abstract class CompletionSession(protected val configuration: CompletionSessionC
                 .filterShadowedNonImported()
     }
 
-    protected fun getTopLevelExtensions(): Collection<CallableDescriptor> {
+    protected fun CallTypeAndReceiver<*, *>.shouldCompleteCallableExtensions(): Boolean {
+        return callType.descriptorKindFilter.kindMask.and(DescriptorKindFilter.CALLABLES_MASK) != 0
+               && this !is CallTypeAndReceiver.IMPORT_DIRECTIVE
+    }
+
+    protected fun getCallableTopLevelExtensions(): Collection<CallableDescriptor> {
         return indicesHelper.getCallableTopLevelExtensions({ prefixMatcher.prefixMatches(it) }, callTypeAndReceiver, expression!!, bindingContext)
-                .filterShadowedNonImported()
+                .filterShadowedNonImported() //TODO: not correct!
     }
 
     private fun Collection<CallableDescriptor>.filterShadowedNonImported(): Collection<CallableDescriptor> {
@@ -305,7 +327,7 @@ abstract class CompletionSession(protected val configuration: CompletionSessionC
         return if (filter != null) filter.filterNonImported(this, referenceVariants) else this
     }
 
-    protected fun addAllClasses(kindFilter: (ClassKind) -> Boolean) {
+    protected fun addClassesFromIndex(kindFilter: (ClassKind) -> Boolean) {
         AllClassesCompletion(parameters, indicesHelper, prefixMatcher, resolutionFacade, kindFilter)
                 .collect(
                         { descriptor -> collector.addDescriptorElements(descriptor, notImported = true) },
@@ -344,4 +366,10 @@ abstract class CompletionSession(protected val configuration: CompletionSessionC
 
         return callTypeAndReceiver to receiverTypes
     }
+
+    protected fun isImportableDescriptorImported(descriptor: DeclarationDescriptor): Boolean {
+        val classification = importableFqNameClassifier.classify(descriptor.importableFqName!!, false)
+        return classification != ImportableFqNameClassifier.Classification.notImported
+               && classification != ImportableFqNameClassifier.Classification.hasImportFromSamePackage
+    }
 }

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/smart/SmartCompletionSession.kt

@@ -20,7 +20,6 @@ import com.intellij.codeInsight.completion.CompletionParameters
 import com.intellij.codeInsight.completion.CompletionResultSet
 import com.intellij.codeInsight.completion.CompletionSorter
 import com.intellij.psi.impl.source.tree.LeafPsiElement
-import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.idea.completion.*
 import org.jetbrains.kotlin.idea.util.CallTypeAndReceiver
 import org.jetbrains.kotlin.load.java.descriptors.SamConstructorDescriptorKindExclude
@@ -34,17 +33,20 @@ import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
 class SmartCompletionSession(configuration: CompletionSessionConfiguration, parameters: CompletionParameters, resultSet: CompletionResultSet)
 : CompletionSession(configuration, parameters, resultSet) {
 
-    // we do not include SAM-constructors because they are handled separately and adding them requires iterating of java classes
+    override val descriptorKindFilter: DescriptorKindFilter by lazy {
-    override val descriptorKindFilter: DescriptorKindFilter
+        // we do not include SAM-constructors because they are handled separately and adding them requires iterating of java classes
-        get() {
+        var filter = DescriptorKindFilter.VALUES exclude SamConstructorDescriptorKindExclude
-            var filter = DescriptorKindFilter.VALUES exclude SamConstructorDescriptorKindExclude
+
-            if (smartCompletion?.expectedInfos?.filterFunctionExpected()?.isNotEmpty() ?: false) {
+        filter = filter exclude topLevelExtensionsExclude // handled via indices
-                // if function type is expected we need classes to obtain their constructors
+
-                filter = filter.withKinds(DescriptorKindFilter.NON_SINGLETON_CLASSIFIERS_MASK)
+        if (smartCompletion?.expectedInfos?.filterFunctionExpected()?.isNotEmpty() ?: false) {
-            }
+            // if function type is expected we need classes to obtain their constructors
-            return filter
+            filter = filter.withKinds(DescriptorKindFilter.NON_SINGLETON_CLASSIFIERS_MASK)
         }
 
+        filter
+    }
+
     private val smartCompletion by lazy(LazyThreadSafetyMode.NONE) {
         expression?.let {
             SmartCompletion(it, resolutionFacade, bindingContext, moduleDescriptor, isVisibleFilter,
@@ -74,9 +76,22 @@ class SmartCompletionSession(configuration: CompletionSessionConfiguration, para
                     referenceVariants.forEach { collector.addElements(filter(it)) }
                     flushToResultSet()
 
-                    processNonImported { collector.addElements(filter(it), notImported = true) }
+                    if (callTypeAndReceiver.shouldCompleteCallableExtensions()) {
+                        for (extension in getCallableTopLevelExtensions()) {
+                            val elements = filter(extension)
+                            if (elements.isNotEmpty()) {
+                                collector.addElements(elements, !isImportableDescriptorImported(extension))
+                            }
+                        }
+                    }
+
                     flushToResultSet()
 
+                    if (shouldCompleteTopLevelCallablesFromIndex()) {
+                        getTopLevelCallables().forEach { collector.addElements(filter(it), notImported = true) }
+                        flushToResultSet()
+                    }
+
                     if (position.getContainingFile() is KtCodeFragment) {
                         getRuntimeReceiverTypeReferenceVariants().forEach {
                             collector.addElements(filter(it).map { it.withReceiverCast() })
@@ -123,16 +138,6 @@ class SmartCompletionSession(configuration: CompletionSessionConfiguration, para
         }
     }
 
-    private fun processNonImported(processor: (DeclarationDescriptor) -> Unit) {
-        getTopLevelExtensions().forEach(processor)
-
-        flushToResultSet()
-
-        if (shouldRunTopLevelCompletion()) {
-            getTopLevelCallables().forEach(processor)
-        }
-    }
-
     override fun createSorter(): CompletionSorter {
         return super.createSorter()
                 .weighBefore(KindWeigher.toString(), NameSimilarityWeigher, SmartCompletionPriorityWeigher)