Completion: all elements are first collected and then flushed to result set

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

@@ -49,16 +49,26 @@ class CompletionSession(public val parameters: CompletionParameters,
             .withPrefixMatcher(CompletionUtil.findJavaIdentifierPrefix(parameters))
             .addKotlinSorting(parameters)
 
-    public val result: CompletionResultSetWrapper = CompletionResultSetWrapper(this.resultSet, resolveSession, { isVisibleDescriptor(it) })
+    private val prefixMatcher = this.resultSet.getPrefixMatcher()
 
-    public fun completeForReference() {
+    private val collector: LookupElementsCollector = LookupElementsCollector(prefixMatcher, resolveSession, { isVisibleDescriptor(it) })
+
+    public fun completeBasic(): Boolean {
         assert(parameters.getCompletionType() == CompletionType.BASIC)
 
+        fillCollectorBasic()
+
+        collector.flushToResultSet(resultSet)
+
+        return !collector.isEmpty
+    }
+
+    private fun fillCollectorBasic() {
         if (isOnlyKeywordCompletion(position)) return
 
         if (shouldRunOnlyTypeCompletion()) {
             if (parameters.getInvocationCount() >= 2) {
-                TypesCompletion(parameters, resolveSession).addAllTypes(result)
+                TypesCompletion(parameters, resolveSession, prefixMatcher).addAllTypes(collector)
             }
             else {
                 addReferenceVariants { isPartOfTypeDeclaration(it) }
@@ -70,7 +80,7 @@ class CompletionSession(public val parameters: CompletionParameters,
 
         addReferenceVariants { true }
 
-        val prefix = result.prefixMatcher.getPrefix()
+        val prefix = prefixMatcher.getPrefix()
 
         // Try to avoid computing not-imported descriptors for empty prefix
         if (prefix.isEmpty()) {
@@ -80,7 +90,7 @@ class CompletionSession(public val parameters: CompletionParameters,
         }
 
         if (shouldRunTopLevelCompletion()) {
-            TypesCompletion(parameters, resolveSession).addAllTypes(result)
+            TypesCompletion(parameters, resolveSession, prefixMatcher).addAllTypes(collector)
             addJetTopLevelFunctions()
             addJetTopLevelObjects()
         }
@@ -88,14 +98,17 @@ class CompletionSession(public val parameters: CompletionParameters,
         if (shouldRunExtensionsCompletion()) {
             addJetExtensions()
         }
+
     }
 
+    //TODO: no collector is needed here
     public fun completeSmart() {
         assert(parameters.getCompletionType() == CompletionType.SMART)
 
         val descriptors = TipsManager.getReferenceVariants(jetReference.expression, bindingContext)
         val completion = SmartCompletion(jetReference.expression, resolveSession, { isVisibleDescriptor(it) }, parameters.getOriginalFile() as JetFile)
-        completion.buildLookupElements(descriptors)?.forEach { result.addElement(it) }
+        completion.buildLookupElements(descriptors)?.forEach { resultSet.addElement(it) }
+        collector.flushToResultSet(resultSet)
     }
 
     private fun isOnlyKeywordCompletion(position: PsiElement)
@@ -104,7 +117,7 @@ class CompletionSession(public val parameters: CompletionParameters,
     private fun addJetExtensions() {
         val project = position.getProject()
         val namesCache = JetShortNamesCache.getKotlinInstance(project)
-        result.addDescriptorElements(namesCache.getJetCallableExtensions({ result.prefixMatcher.prefixMatches(it!!) }, jetReference.expression, resolveSession, GlobalSearchScope.allScope(project)))
+        collector.addDescriptorElements(namesCache.getJetCallableExtensions({ prefixMatcher.prefixMatches(it!!) }, jetReference.expression, resolveSession, GlobalSearchScope.allScope(project)))
     }
 
     private fun isPartOfTypeDeclaration(descriptor: DeclarationDescriptor): Boolean {
@@ -122,7 +135,7 @@ class CompletionSession(public val parameters: CompletionParameters,
     }
 
     private fun addJetTopLevelFunctions() {
-        val actualPrefix = result.prefixMatcher.getPrefix()
+        val actualPrefix = prefixMatcher.getPrefix()
         val project = position.getProject()
         val namesCache = JetShortNamesCache.getKotlinInstance(project)
         val scope = GlobalSearchScope.allScope(project)
@@ -131,7 +144,7 @@ class CompletionSession(public val parameters: CompletionParameters,
         // TODO: Fix complete extension not only on contains
         for (name in functionNames) {
             if (name.contains(actualPrefix)) {
-                result.addDescriptorElements(namesCache.getTopLevelFunctionDescriptorsByName(name, jetReference.expression, resolveSession, scope))
+                collector.addDescriptorElements(namesCache.getTopLevelFunctionDescriptorsByName(name, jetReference.expression, resolveSession, scope))
             }
         }
     }
@@ -143,8 +156,8 @@ class CompletionSession(public val parameters: CompletionParameters,
         val objectNames = namesCache.getAllTopLevelObjectNames()
 
         for (name in objectNames) {
-            if (result.prefixMatcher.prefixMatches(name)) {
+            if (prefixMatcher.prefixMatches(name)) {
-                result.addDescriptorElements(namesCache.getTopLevelObjectsByName(name, jetReference.expression, resolveSession, scope))
+                collector.addDescriptorElements(namesCache.getTopLevelObjectsByName(name, jetReference.expression, resolveSession, scope))
             }
         }
     }
@@ -175,12 +188,12 @@ class CompletionSession(public val parameters: CompletionParameters,
     }
 
     private fun shouldRunExtensionsCompletion(): Boolean {
-        return parameters.getInvocationCount() > 1 || result.prefixMatcher.getPrefix().length >= 3
+        return parameters.getInvocationCount() > 1 || prefixMatcher.getPrefix().length >= 3
     }
 
     private fun addReferenceVariants(filterCondition: (DeclarationDescriptor) -> Boolean) {
         val descriptors = TipsManager.getReferenceVariants(jetReference.expression, bindingContext)
-        result.addDescriptorElements(descriptors.filter { filterCondition(it) })
+        collector.addDescriptorElements(descriptors.filter { filterCondition(it) })
     }
 
     private fun isVisibleDescriptor(descriptor: DeclarationDescriptor): Boolean {

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

@@ -120,12 +120,11 @@ public class JetCompletionContributor : CompletionContributor() {
 
                     var session = CompletionSession(parameters, result, jetReference, position)
                     if (parameters.getCompletionType() == CompletionType.BASIC) {
-                        session.completeForReference()
+                        val somethingAdded = session.completeBasic()
-
+                        if (!somethingAdded && session.parameters.getInvocationCount() < 2) {
-                        if (!session.result.isSomethingAdded && session.parameters.getInvocationCount() < 2) {
                             // Rerun completion if nothing was found
                             session = CompletionSession(parameters.withInvocationCount(2), result, jetReference, position)
-                            session.completeForReference()
+                            session.completeBasic()
                         }
                     }
                     else {

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

@@ -16,7 +16,6 @@
 
 package org.jetbrains.jet.plugin.completion
 
-import com.intellij.codeInsight.completion.CompletionResultSet
 import com.intellij.codeInsight.completion.InsertionContext
 import com.intellij.codeInsight.lookup.LookupElement
 import com.intellij.codeInsight.lookup.LookupElementDecorator
@@ -27,14 +26,20 @@ import org.jetbrains.jet.lang.types.lang.KotlinBuiltIns
 import org.jetbrains.jet.plugin.completion.handlers.*
 import org.jetbrains.jet.plugin.project.ResolveSessionForBodies
 import com.intellij.codeInsight.completion.PrefixMatcher
+import java.util.ArrayList
+import com.intellij.codeInsight.completion.CompletionResultSet
 
-class CompletionResultSetWrapper(private val resultSet: CompletionResultSet,
+class LookupElementsCollector(private val prefixMatcher: PrefixMatcher,
-                                 private val resolveSession: ResolveSessionForBodies,
+                              private val resolveSession: ResolveSessionForBodies,
-                                 private val descriptorFilter: (DeclarationDescriptor) -> Boolean) {
+                              private val descriptorFilter: (DeclarationDescriptor) -> Boolean) {
-    public var isSomethingAdded: Boolean = false
+    private val elements = ArrayList<LookupElement>()
-        private set
 
-    public val prefixMatcher: PrefixMatcher = resultSet.getPrefixMatcher()
+    public fun flushToResultSet(resultSet: CompletionResultSet) {
+        resultSet.addAllElements(elements)
+    }
+
+    public val isEmpty: Boolean
+        get() = elements.isEmpty()
 
     public fun addDescriptorElements(descriptors: Iterable<DeclarationDescriptor>) {
         for (descriptor in descriptors) {
@@ -74,8 +79,7 @@ class CompletionResultSetWrapper(private val resultSet: CompletionResultSet,
 
     public fun addElement(element: LookupElement) {
         if (prefixMatcher.prefixMatches(element)) {
-            resultSet.addElement(element)
+            elements.add(element)
-            isSomethingAdded = true
         }
     }
 }

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

@@ -35,13 +35,13 @@ import org.jetbrains.jet.plugin.completion.handlers.JetJavaClassInsertHandler
 import org.jetbrains.jet.plugin.project.ProjectStructureUtil
 import org.jetbrains.jet.plugin.project.ResolveSessionForBodies
 
-class TypesCompletion(val parameters: CompletionParameters, val resolveSession: ResolveSessionForBodies) {
+class TypesCompletion(val parameters: CompletionParameters, val resolveSession: ResolveSessionForBodies, val prefixMatcher: PrefixMatcher) {
-    fun addAllTypes(result: CompletionResultSetWrapper) {
+    fun addAllTypes(result: LookupElementsCollector) {
         result.addDescriptorElements(KotlinBuiltIns.getInstance().getNonPhysicalClasses())
 
         val project = parameters.getOriginalFile().getProject()
         val namesCache = JetShortNamesCache.getKotlinInstance(project)
-        result.addDescriptorElements(namesCache.getJetClassesDescriptors({ result.prefixMatcher.prefixMatches(it!!) }, resolveSession, GlobalSearchScope.allScope(project)))
+        result.addDescriptorElements(namesCache.getJetClassesDescriptors({ prefixMatcher.prefixMatches(it!!) }, resolveSession, GlobalSearchScope.allScope(project)))
 
         if (!ProjectStructureUtil.isJsKotlinModule(parameters.getOriginalFile() as JetFile)) {
             addAdaptedJavaCompletion(result)
@@ -51,8 +51,8 @@ class TypesCompletion(val parameters: CompletionParameters, val resolveSession:
     /**
      * Add java elements with performing conversion to kotlin elements if necessary.
      */
-    private fun addAdaptedJavaCompletion(result: CompletionResultSetWrapper) {
+    private fun addAdaptedJavaCompletion(result: LookupElementsCollector) {
-        JavaClassNameCompletionContributor.addAllClasses(parameters, false, result.prefixMatcher, object : Consumer<LookupElement> {
+        JavaClassNameCompletionContributor.addAllClasses(parameters, false, prefixMatcher, object : Consumer<LookupElement> {
             override fun consume(lookupElement: LookupElement?) {
                 if (lookupElement is JavaPsiClassReferenceElement) {
                     val psiClass = lookupElement.getObject()
@@ -69,7 +69,7 @@ class TypesCompletion(val parameters: CompletionParameters, val resolveSession:
         })
     }
 
-    private fun addJavaClassAsJetLookupElement(aClass: PsiClass, result: CompletionResultSetWrapper): Boolean {
+    private fun addJavaClassAsJetLookupElement(aClass: PsiClass, result: LookupElementsCollector): Boolean {
         if (aClass is KotlinLightClass) {
             // Do nothing. Kotlin not-compiled class should have already been added as kotlin element before.
             return true