Code completion: more optimal fetching of top level functions

compiler/frontend/src/org/jetbrains/jet/lang/resolve/QualifiedExpressionResolver.java

@@ -20,12 +20,12 @@ import com.google.common.base.Predicate;
 import com.google.common.collect.Collections2;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Sets;
+import kotlin.jvm.KotlinSignature;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.descriptors.*;
 import org.jetbrains.jet.lang.diagnostics.Errors;
 import org.jetbrains.jet.lang.psi.*;
-import org.jetbrains.jet.lang.psi.codeFragmentUtil.CodeFragmentUtilPackage;
 import org.jetbrains.jet.lang.resolve.name.Name;
 import org.jetbrains.jet.lang.resolve.scopes.JetScope;
 
@@ -51,6 +51,7 @@ public class QualifiedExpressionResolver {
         EVERYTHING
     }
 
+    @KotlinSignature( "fun analyseImportReference(importDirective: JetImportDirective, scope: JetScope, trace: BindingTrace, module: ModuleDescriptor): Collection<out DeclarationDescriptor>")
     @NotNull
     public Collection<? extends DeclarationDescriptor> analyseImportReference(
             @NotNull JetImportDirective importDirective,

idea/idea-analysis/src/org/jetbrains/jet/plugin/stubindex/KotlinIndexUtil.java

@@ -19,11 +19,10 @@ package org.jetbrains.jet.plugin.stubindex;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.stubs.StubIndexExtension;
 import com.intellij.psi.stubs.StubIndexKey;
-import org.jetbrains.annotations.NonNls;
 import org.jetbrains.annotations.NotNull;
 
 public class KotlinIndexUtil {
-    @NonNls
+    @NotNull
     public static <K, Psi extends PsiElement> StubIndexKey<K, Psi> createIndexKey(@NotNull Class<? extends StubIndexExtension<K, Psi>> indexClass) {
         return StubIndexKey.createIndexKey(indexClass.getCanonicalName());
     }

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

@@ -18,7 +18,6 @@ package org.jetbrains.jet.plugin.caches
 
 import com.intellij.openapi.extensions.Extensions
 import com.intellij.openapi.project.Project
-import com.intellij.openapi.util.Condition
 import com.intellij.psi.*
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.search.PsiShortNamesCache
@@ -158,7 +157,7 @@ public class JetShortNamesCache(private val project: Project) : PsiShortNamesCac
 
     public fun getTopLevelFunctionDescriptorsByName(name: String, context: JetExpression /*TODO: to be dropped*/, resolveSession: ResolveSessionForBodies, scope: GlobalSearchScope): Collection<FunctionDescriptor> {
 
-        val jetScope = resolveSession.resolveToElement(context).get<JetExpression, JetScope>(BindingContext.RESOLUTION_SCOPE, context) ?: return listOf()
+        val jetScope = resolveSession.resolveToElement(context).get(BindingContext.RESOLUTION_SCOPE, context) ?: return listOf()
 
         val result = hashSetOf<FunctionDescriptor>()
 
@@ -167,16 +166,22 @@ public class JetShortNamesCache(private val project: Project) : PsiShortNamesCac
         val topLevelFunctionFqNames = JetFromJavaDescriptorHelper.getTopLevelFunctionFqNames(project, scope, false)
         for (fqName in topLevelFunctionFqNames) {
             if (fqName.lastSegmentIs(identifier)) {
-                val importDirective = JetPsiFactory(context.getProject()).createImportDirective(ImportPath(fqName, false))
+                findCompiledTopLevelEntities(fqName, context, jetScope, resolveSession).filterIsInstanceTo(result, javaClass<FunctionDescriptor>())
-                val declarationDescriptors = QualifiedExpressionResolver().analyseImportReference(importDirective, jetScope, BindingTraceContext(), resolveSession.getModuleDescriptor())
-                for (declarationDescriptor in declarationDescriptors) {
-                    if (declarationDescriptor is FunctionDescriptor) {
-                        result.add(declarationDescriptor as FunctionDescriptor)
-                    }
-                }
             }
         }
 
+        result.addJetSourceTopLevelFunctions(name, resolveSession, scope)
+
+        return result
+    }
+
+    private fun findCompiledTopLevelEntities(fqName: FqName, context: JetExpression, jetScope: JetScope, resolveSession: ResolveSessionForBodies): Collection<DeclarationDescriptor> {
+        val importDirective = JetPsiFactory(context.getProject()).createImportDirective(ImportPath(fqName, false))
+        return QualifiedExpressionResolver().analyseImportReference(importDirective, jetScope, BindingTraceContext(), resolveSession.getModuleDescriptor())
+    }
+
+    private fun MutableCollection<in FunctionDescriptor>.addJetSourceTopLevelFunctions(name: String, resolveSession: ResolveSessionForBodies, scope: GlobalSearchScope) {
+        val identifier = Name.identifier(name)
         val affectedPackages = JetTopLevelNonExtensionFunctionShortNameIndex.getInstance().get(name, project, scope)
                 .map { it.getContainingFile() }
                 .filterIsInstance(javaClass<JetFile>())
@@ -186,36 +191,54 @@ public class JetShortNamesCache(private val project: Project) : PsiShortNamesCac
         for (affectedPackage in affectedPackages) {
             val packageDescriptor = resolveSession.getModuleDescriptor().getPackage(affectedPackage)
                     ?: error("There's a function in stub index with invalid package: $affectedPackage")
-            result.addAll(packageDescriptor.getMemberScope().getFunctions(identifier))
+            addAll(packageDescriptor.getMemberScope().getFunctions(identifier))
         }
-
-        return result
     }
 
     private fun getJetExtensionFunctionsByName(name: String, scope: GlobalSearchScope): Collection<PsiElement>
             = JetTopLevelExtensionFunctionShortNameIndex.getInstance().get(name, project, scope)
 
     // TODO: Make it work for properties
-    public fun getJetCallableExtensions(nameFilter: (String) -> Boolean, expression: JetSimpleNameExpression, resolveSession: ResolveSessionForBodies, searchScope: GlobalSearchScope): Collection<DeclarationDescriptor> {
+    public fun getJetTopLevelCallables(nameFilter: (String) -> Boolean, context: JetExpression /*TODO: to be dropped*/, resolveSession: ResolveSessionForBodies, scope: GlobalSearchScope): Collection<DeclarationDescriptor> {
+        val result = ArrayList<DeclarationDescriptor>()
+
+        JetTopLevelNonExtensionFunctionShortNameIndex.getInstance().getAllKeys(project)
+                .stream()
+                .filter(nameFilter)
+                .forEach { result.addJetSourceTopLevelFunctions(it, resolveSession, scope) }
+
+        val jetScope = resolveSession.resolveToElement(context).get(BindingContext.RESOLUTION_SCOPE, context) ?: return listOf()
+        JetFromJavaDescriptorHelper.getTopLevelFunctionFqNames(project, scope, false)
+                .filter { nameFilter(it.shortName().asString()) }
+                .flatMapTo(result) { findCompiledTopLevelEntities(it, context, jetScope, resolveSession) }
+
+        return result
+    }
+
+    // TODO: Make it work for properties
+    public fun getJetCallableExtensions(nameFilter: (String) -> Boolean,
+                                        expression: JetSimpleNameExpression,
+                                        resolveSession: ResolveSessionForBodies,
+                                        scope: GlobalSearchScope): Collection<DeclarationDescriptor> {
         val context = resolveSession.resolveToElement(expression)
         val receiverExpression = expression.getReceiverExpression() ?: return listOf()
         val expressionType = context.get<JetExpression, JetType>(BindingContext.EXPRESSION_TYPE, receiverExpression)
-        val scope = context.get<JetExpression, JetScope>(BindingContext.RESOLUTION_SCOPE, receiverExpression)
+        val jetScope = context.get(BindingContext.RESOLUTION_SCOPE, receiverExpression)
 
-        if (expressionType == null || scope == null || expressionType.isError()) {
+        if (expressionType == null || jetScope == null || expressionType.isError()) {
             return listOf()
         }
 
-        val functionFQNs = extensionFunctionsFromSourceFqNames(nameFilter, searchScope)
+        val functionFQNs = extensionFunctionsFromSourceFqNames(nameFilter, scope)
 
-        JetFromJavaDescriptorHelper.getTopLevelFunctionFqNames(project, searchScope, true)
+        JetFromJavaDescriptorHelper.getTopLevelFunctionFqNames(project, scope, true)
                 .filterTo(functionFQNs) { nameFilter(it.shortName().asString()) }
 
         // Iterate through the function with attempt to resolve found functions
-        return functionFQNs.flatMap { ExpressionTypingUtils.canFindSuitableCall(it, receiverExpression, expressionType, scope, resolveSession.getModuleDescriptor()) }
+        return functionFQNs.flatMap { ExpressionTypingUtils.canFindSuitableCall(it, receiverExpression, expressionType, jetScope, resolveSession.getModuleDescriptor()) }
     }
 
-    private fun extensionFunctionsFromSourceFqNames(nameFilter: (String) -> Boolean, searchScope: GlobalSearchScope): HashSet<FqName> {
+    private fun extensionFunctionsFromSourceFqNames(nameFilter: (String) -> Boolean, scope: GlobalSearchScope): HashSet<FqName> {
         val extensionFunctionNames = HashSet(JetTopLevelExtensionFunctionShortNameIndex.getInstance().getAllKeys(project))
 
         val functionFQNs = HashSet<FqName>()
@@ -223,7 +246,7 @@ public class JetShortNamesCache(private val project: Project) : PsiShortNamesCac
         // Collect all possible extension function qualified names
         for (name in extensionFunctionNames) {
             if (nameFilter(name)) {
-                val extensionFunctions = getJetExtensionFunctionsByName(name, searchScope)
+                val extensionFunctions = getJetExtensionFunctionsByName(name, scope)
 
                 for (extensionFunction in extensionFunctions) {
                     if (extensionFunction is JetNamedFunction) {
@@ -241,21 +264,21 @@ public class JetShortNamesCache(private val project: Project) : PsiShortNamesCac
         return functionFQNs
     }
 
-    public fun getJetClassesDescriptors(acceptedShortNameCondition: (String) -> Boolean, analyzer: KotlinCodeAnalyzer, searchScope: GlobalSearchScope): Collection<ClassDescriptor> {
+    public fun getJetClassesDescriptors(acceptedShortNameCondition: (String) -> Boolean, analyzer: KotlinCodeAnalyzer, scope: GlobalSearchScope): Collection<ClassDescriptor> {
         val classDescriptors = ArrayList<ClassDescriptor>()
 
         for (fqName in JetFullClassNameIndex.getInstance().getAllKeys(project)) {
             val classFQName = FqName(fqName)
             if (acceptedShortNameCondition(classFQName.shortName().asString())) {
-                classDescriptors.addAll(getJetClassesDescriptorsByFQName(analyzer, classFQName, searchScope))
+                classDescriptors.addAll(getJetClassesDescriptorsByFQName(analyzer, classFQName, scope))
             }
         }
 
         return classDescriptors
     }
 
-    private fun getJetClassesDescriptorsByFQName(analyzer: KotlinCodeAnalyzer, classFQName: FqName, searchScope: GlobalSearchScope): Collection<ClassDescriptor> {
+    private fun getJetClassesDescriptorsByFQName(analyzer: KotlinCodeAnalyzer, classFQName: FqName, scope: GlobalSearchScope): Collection<ClassDescriptor> {
-        val jetClassOrObjects = JetFullClassNameIndex.getInstance().get(classFQName.asString(), project, searchScope)
+        val jetClassOrObjects = JetFullClassNameIndex.getInstance().get(classFQName.asString(), project, scope)
 
         if (jetClassOrObjects.isEmpty()) {
             // This fqn is absent in caches, dead or not in scope

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

@@ -129,8 +129,7 @@ class BasicCompletionSession(configuration: CompletionSessionConfiguration,
 
         if (shouldRunTopLevelCompletion()) {
             TypesCompletion(parameters, resolveSession, prefixMatcher).addAllTypes(collector)
-            addKotlinTopLevelFunctions()
+            addKotlinTopLevelDeclarations()
-            addKotlinTopLevelObjects()
         }
 
         if (shouldRunExtensionsCompletion()) {
@@ -155,15 +154,9 @@ class BasicCompletionSession(configuration: CompletionSessionConfiguration,
         }
     }
 
-    private fun addKotlinTopLevelFunctions() {
+    private fun addKotlinTopLevelDeclarations() {
-        for (name in shortNamesCache.getAllTopLevelFunctionNames()) {
+        collector.addDescriptorElements(shortNamesCache.getJetTopLevelCallables({ prefixMatcher.prefixMatches(it) }, jetReference!!.expression, resolveSession, searchScope))
-            if (prefixMatcher.prefixMatches(name)) {
-                collector.addDescriptorElements(shortNamesCache.getTopLevelFunctionDescriptorsByName(name, jetReference!!.expression, resolveSession, searchScope))
-            }
-        }
-    }
 
-    private fun addKotlinTopLevelObjects() {
         for (name in shortNamesCache.getAllTopLevelObjectNames()) {
             if (prefixMatcher.prefixMatches(name)) {
                 collector.addDescriptorElements(shortNamesCache.getTopLevelObjectsByName(name, resolveSession, searchScope))