Faster completion of non-imported extensions

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetDeclarationStub.java

@@ -17,6 +17,7 @@
 package org.jetbrains.jet.lang.psi;
 
 import com.intellij.lang.ASTNode;
+import com.intellij.psi.PsiElement;
 import com.intellij.psi.stubs.IStubElementType;
 import com.intellij.psi.stubs.StubElement;
 import org.jetbrains.annotations.NotNull;
@@ -38,4 +39,13 @@ abstract class JetDeclarationStub<T extends StubElement> extends JetModifierList
     public KDoc getDocComment() {
         return FindDocCommentPackage.findDocComment(this);
     }
+
+    @Override
+    public PsiElement getParent() {
+        T stub = getStub();
+        if (stub != null) {
+            return stub.getParentStub().getPsi();
+        }
+        return super.getParent();
+    }
 }

idea/src/org/jetbrains/jet/plugin/caches/KotlinIndicesHelper.kt

@@ -26,9 +26,7 @@ import org.jetbrains.jet.lang.psi.*
 import org.jetbrains.jet.lang.resolve.*
 import org.jetbrains.jet.lang.resolve.name.Name
 import org.jetbrains.jet.lang.psi.psiUtil.getReceiverExpression
-import org.jetbrains.jet.lang.types.JetType
 import org.jetbrains.jet.lang.types.expressions.ExpressionTypingUtils
-import org.jetbrains.jet.lang.resolve.lazy.KotlinCodeAnalyzer
 import org.jetbrains.jet.lang.resolve.scopes.JetScope
 import com.intellij.openapi.project.Project
 import java.util.HashSet
@@ -37,6 +35,8 @@ import org.jetbrains.jet.lang.resolve.bindingContextUtil.getDataFlowInfo
 import org.jetbrains.jet.lang.resolve.QualifiedExpressionResolver.LookupMode
 import org.jetbrains.jet.lang.resolve.scopes.receivers.ReceiverValue
 import org.jetbrains.jet.lang.resolve.calls.smartcasts.DataFlowInfo
+import com.intellij.psi.stubs.StringStubIndexExtension
+import org.jetbrains.jet.plugin.caches.resolve.getLazyResolveSession
 
 public class KotlinIndicesHelper(private val project: Project,
                                  private val resolveSession: ResolveSessionForBodies,
@@ -138,58 +138,81 @@ public class KotlinIndicesHelper(private val project: Project,
         val bindingContext = resolveSession.resolveToElement(expression)
         val dataFlowInfo = bindingContext.getDataFlowInfo(expression)
 
-        val sourceNames = JetTopLevelFunctionsFqnNameIndex.getInstance().getAllKeys(project).stream() +
+        val functionsIndex = JetTopLevelFunctionsFqnNameIndex.getInstance()
-                          JetTopLevelPropertiesFqnNameIndex.getInstance().getAllKeys(project).stream()
+        val propertiesIndex = JetTopLevelPropertiesFqnNameIndex.getInstance()
-        val allFqNames = sourceNames.map { FqName(it) } + JetFromJavaDescriptorHelper.getTopLevelCallableFqNames(project, scope, true).stream()
+
-        val matchingFqNames = allFqNames.filter { nameFilter(it.shortName().asString()) }.toSet()
+        val sourceFunctionNames = functionsIndex.getAllKeys(project).stream().map { FqName(it) }
+        val sourcePropertyNames = propertiesIndex.getAllKeys(project).stream().map { FqName(it) }
+        val compiledFqNames = JetFromJavaDescriptorHelper.getTopLevelCallableFqNames(project, scope, true).stream()
+
+        val result = HashSet<CallableDescriptor>()
+        result.fqNamesToSuitableExtensions(sourceFunctionNames, nameFilter, functionsIndex, expression, bindingContext, dataFlowInfo)
+        result.fqNamesToSuitableExtensions(sourcePropertyNames, nameFilter, propertiesIndex, expression, bindingContext, dataFlowInfo)
+        result.fqNamesToSuitableExtensions(compiledFqNames, nameFilter, null, expression, bindingContext, dataFlowInfo)
+        return result
+    }
+
+    private fun MutableCollection<CallableDescriptor>.fqNamesToSuitableExtensions(
+            fqNames: Stream<FqName>,
+            nameFilter: (String) -> Boolean,
+            index: StringStubIndexExtension<out JetCallableDeclaration>?,
+            expression: JetSimpleNameExpression,
+            bindingContext: BindingContext,
+            dataFlowInfo: DataFlowInfo) {
+        val matchingNames = fqNames.filter { nameFilter(it.shortName().asString()) }
 
         val receiverExpression = expression.getReceiverExpression()
         if (receiverExpression != null) {
             val isInfixCall = expression.getParent() is JetBinaryExpression
 
-            val expressionType = bindingContext.get<JetExpression, JetType>(BindingContext.EXPRESSION_TYPE, receiverExpression)
+            val expressionType = bindingContext[BindingContext.EXPRESSION_TYPE, receiverExpression]
-            if (expressionType == null || expressionType.isError()) return listOf()
+            if (expressionType == null || expressionType.isError()) return
 
-            val resolutionScope = bindingContext.get(BindingContext.RESOLUTION_SCOPE, receiverExpression) ?: return listOf()
+            val resolutionScope = bindingContext[BindingContext.RESOLUTION_SCOPE, receiverExpression] ?: return
 
             val receiverValue = ExpressionReceiver(receiverExpression, expressionType)
 
-            return matchingFqNames.flatMap {
+            matchingNames.flatMapTo(this) {
-                findSuitableExtensions(it, receiverValue, dataFlowInfo, isInfixCall, resolutionScope, moduleDescriptor, bindingContext)
+                findSuitableExtensions(it, index, receiverValue, dataFlowInfo, isInfixCall, resolutionScope, moduleDescriptor, bindingContext).stream()
             }
         }
         else {
-            val resolutionScope = bindingContext[BindingContext.RESOLUTION_SCOPE, expression] ?: return listOf()
+            val resolutionScope = bindingContext[BindingContext.RESOLUTION_SCOPE, expression] ?: return
 
-            val result = HashSet<CallableDescriptor>()
             for (receiver in resolutionScope.getImplicitReceiversHierarchy()) {
-                matchingFqNames.flatMapTo(result) {
+                matchingNames.flatMapTo(this) {
-                    findSuitableExtensions(it, receiver.getValue(), dataFlowInfo, false, resolutionScope, moduleDescriptor, bindingContext)
+                    findSuitableExtensions(it, index, receiver.getValue(), dataFlowInfo, false, resolutionScope, moduleDescriptor, bindingContext).stream()
                 }
             }
-            return result
         }
     }
 
-
     /**
      * Check that function or property with the given qualified name can be resolved in given scope and called on given receiver
      */
     private fun findSuitableExtensions(callableFQN: FqName,
+                                       index: StringStubIndexExtension<out JetCallableDeclaration>?,
                                        receiverValue: ReceiverValue,
                                        dataFlowInfo: DataFlowInfo,
                                        isInfixCall: Boolean,
-                                       scope: JetScope,
+                                       resolutionScope: JetScope,
                                        module: ModuleDescriptor,
-                                       bindingContext: BindingContext): List<CallableDescriptor> {
+                                       bindingContext: BindingContext): Collection<CallableDescriptor> {
-        val importDirective = JetPsiFactory(project).createImportDirective(callableFQN.asString())
+        val fqnString = callableFQN.asString()
-        val declarationDescriptors = analyzeImportReference(importDirective, scope, BindingTraceContext(), module)
+        val descriptors = if (index != null) {
+            index.get(fqnString, project, scope)
+                    .filter { it.getReceiverTypeReference() != null }
+                    .map { it.getLazyResolveSession().resolveToDescriptor(it) as CallableDescriptor }
+        }
+        else {
+            val importDirective = JetPsiFactory(project).createImportDirective(fqnString)
+            analyzeImportReference(importDirective, resolutionScope, BindingTraceContext(), module)
+                    .filterIsInstance(javaClass<CallableDescriptor>())
+                    .filter { it.getExtensionReceiverParameter() != null }
+        }
 
-        return declarationDescriptors
+        return descriptors.filter { visibilityFilter(it) &&
-                .filterIsInstance(javaClass<CallableDescriptor>())
+                    ExpressionTypingUtils.checkIsExtensionCallable(receiverValue, it, isInfixCall, bindingContext, dataFlowInfo) }
-                .filter { it.getExtensionReceiverParameter() != null &&
-                          visibilityFilter(it) &&
-                          ExpressionTypingUtils.checkIsExtensionCallable(receiverValue, it, isInfixCall, bindingContext, dataFlowInfo) }
     }
 
     public fun getClassDescriptors(nameFilter: (String) -> Boolean, kindFilter: (ClassKind) -> Boolean): Collection<ClassDescriptor> {

idea/src/org/jetbrains/jet/plugin/completion/AllClassesCompletion.kt

@@ -30,9 +30,11 @@ import org.jetbrains.jet.plugin.project.ProjectStructureUtil
 import org.jetbrains.jet.plugin.project.ResolveSessionForBodies
 import org.jetbrains.jet.plugin.caches.KotlinIndicesHelper
 import org.jetbrains.jet.lang.descriptors.DeclarationDescriptor
+import com.intellij.psi.search.GlobalSearchScope
 
 class AllClassesCompletion(val parameters: CompletionParameters,
                            val resolveSession: ResolveSessionForBodies,
+                           val scope: GlobalSearchScope,
                            val prefixMatcher: PrefixMatcher,
                            val kindFilter: (ClassKind) -> Boolean,
                            val visibilityFilter: (DeclarationDescriptor) -> Boolean) {
@@ -40,14 +42,11 @@ class AllClassesCompletion(val parameters: CompletionParameters,
         val builtIns = KotlinBuiltIns.getInstance().getNonPhysicalClasses().filter { kindFilter(it.getKind()) && prefixMatcher.prefixMatches(it.getName().asString()) }
         result.addDescriptorElements(builtIns, suppressAutoInsertion = true)
 
-        val file = parameters.getOriginalFile()
+        val helper = KotlinIndicesHelper(scope.getProject(), resolveSession, scope, visibilityFilter)
-        val project = file.getProject()
-        val searchScope = file.getResolveScope()
-        val helper = KotlinIndicesHelper(project, resolveSession, searchScope, visibilityFilter)
         result.addDescriptorElements(helper.getClassDescriptors({ prefixMatcher.prefixMatches(it) }, kindFilter),
                                      suppressAutoInsertion = true)
 
-        if (!ProjectStructureUtil.isJsKotlinModule(file as JetFile)) {
+        if (!ProjectStructureUtil.isJsKotlinModule(parameters.getOriginalFile() as JetFile)) {
             addAdaptedJavaCompletion(result)
         }
     }

idea/src/org/jetbrains/jet/plugin/completion/CompletionSession.kt

@@ -32,6 +32,9 @@ import org.jetbrains.jet.plugin.references.JetSimpleNameReference
 import org.jetbrains.jet.plugin.project.ResolveSessionForBodies
 import org.jetbrains.jet.plugin.caches.KotlinIndicesHelper
 import com.intellij.openapi.project.Project
+import com.intellij.psi.search.DelegatingGlobalSearchScope
+import java.util.HashSet
+import com.intellij.openapi.vfs.VirtualFile
 import org.jetbrains.jet.lang.resolve.scopes.JetScope
 
 class CompletionSessionConfiguration(
@@ -62,7 +65,12 @@ abstract class CompletionSessionBase(protected val configuration: CompletionSess
     protected val collector: LookupElementsCollector = LookupElementsCollector(prefixMatcher, resolveSession)
 
     protected val project: Project = position.getProject()
-    protected val searchScope: GlobalSearchScope = parameters.getOriginalFile().getResolveScope()
+
+    // we need to exclude the original file from scope because our resolve session is built with this file replaced by synthetic one
+    protected val searchScope: GlobalSearchScope = object : DelegatingGlobalSearchScope(parameters.getOriginalFile().getResolveScope()) {
+        override fun contains(file: VirtualFile) = super.contains(file) && file != parameters.getOriginalFile().getVirtualFile()
+    }
+
     protected val indicesHelper: KotlinIndicesHelper = KotlinIndicesHelper(project, resolveSession, searchScope) { isVisibleDescriptor(it) }
 
     protected fun isVisibleDescriptor(descriptor: DeclarationDescriptor): Boolean {
@@ -123,7 +131,7 @@ abstract class CompletionSessionBase(protected val configuration: CompletionSess
     }
 
     protected fun addAllClasses(kindFilter: (ClassKind) -> Boolean) {
-        AllClassesCompletion(parameters, resolveSession, prefixMatcher, kindFilter, { isVisibleDescriptor(it) }).collect(collector)
+        AllClassesCompletion(parameters, resolveSession, searchScope, prefixMatcher, kindFilter, { isVisibleDescriptor(it) }).collect(collector)
     }
 }