Smart completion: code improvements and refactorings after code review

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

@@ -16,6 +16,8 @@
 
 package org.jetbrains.jet.plugin.completion;
 
+import com.google.common.base.Predicate;
+import com.google.common.collect.Collections2;
 import com.intellij.codeInsight.completion.CompletionParameters;
 import com.intellij.codeInsight.completion.CompletionResultSet;
 import com.intellij.codeInsight.completion.CompletionType;
@@ -43,7 +45,6 @@ import org.jetbrains.jet.plugin.project.ResolveSessionForBodies;
 import org.jetbrains.jet.plugin.references.JetSimpleNameReference;
 
 import java.util.Collection;
-import java.util.Collections;
 
 class CompletionSession {
     @Nullable
@@ -52,13 +53,6 @@ class CompletionSession {
     private final JetCompletionResultSet jetResult;
     private final JetSimpleNameReference jetReference;
 
-    private final Condition<DeclarationDescriptor> descriptorFilter = new Condition<DeclarationDescriptor>() {
-        @Override
-        public boolean value(DeclarationDescriptor descriptor) {
-            return isVisibleDescriptor(descriptor);
-        }
-    };
-
     public CompletionSession(
             @NotNull CompletionParameters parameters,
             @NotNull CompletionResultSet result,
@@ -75,10 +69,17 @@ class CompletionSession {
 
         inDescriptor = scope != null ? scope.getContainingDeclaration() : null;
 
+        Condition<DeclarationDescriptor> descriptorFilter = new Condition<DeclarationDescriptor>() {
+            @Override
+            public boolean value(DeclarationDescriptor descriptor) {
+                return isVisibleDescriptor(descriptor);
+            }
+        };
         this.jetResult = new JetCompletionResultSet(
                 WeigherPackage.addJetSorting(result, parameters),
                 resolveSession,
-                expressionBindingContext, descriptorFilter);
+                expressionBindingContext,
+                descriptorFilter);
     }
 
     public void completeForReference() {
@@ -134,20 +135,18 @@ class CompletionSession {
     public void completeSmart() {
         assert parameters.getCompletionType() == CompletionType.SMART;
 
-        final SmartCompletionData data = CompletionPackage.buildSmartCompletionData(jetReference.getExpression(), getResolveSession(), new Function1<DeclarationDescriptor, Boolean>() {
+        Collection<DeclarationDescriptor> descriptors = TipsManager.getReferenceVariants(
+                jetReference.getExpression(), getExpressionBindingContext());
+        Function1<DeclarationDescriptor, Boolean> visibilityFilter = new Function1<DeclarationDescriptor, Boolean>() {
             @Override
             public Boolean invoke(DeclarationDescriptor descriptor) {
                 return isVisibleDescriptor(descriptor);
             }
-        });
+        };
-        if (data != null) {
+        Collection<LookupElement> elements = CompletionPackage.buildSmartCompletionData(
-            addReferenceVariants(new Function1<DeclarationDescriptor, Iterable<LookupElement>>(){
+                jetReference.getExpression(), getResolveSession(), descriptors, visibilityFilter);
-                @Override
+        if (elements != null) {
-                public Iterable<LookupElement> invoke(DeclarationDescriptor descriptor) {
+            for (LookupElement element : elements) {
-                    return data.toElements(descriptor);
-                }
-            });
-            for (LookupElement element : data.getAdditionalElements()) {
                 jetResult.addElement(element);
             }
         }
@@ -244,33 +243,21 @@ class CompletionSession {
     }
 
     private boolean shouldRunExtensionsCompletion() {
-        return !(parameters.getInvocationCount() <= 1 && jetResult.getResult().getPrefixMatcher().getPrefix().length() < 3);
+        return parameters.getInvocationCount() > 1 || jetResult.getResult().getPrefixMatcher().getPrefix().length() >= 3;
     }
 
     private void addReferenceVariants(@NotNull final Condition<DeclarationDescriptor> filterCondition) {
-        addReferenceVariants(new Function1<DeclarationDescriptor, Iterable<LookupElement>>(){
-            @Override
-            public Iterable<LookupElement> invoke(DeclarationDescriptor descriptor) {
-                return filterCondition.value(descriptor)
-                       ? Collections.singletonList(
-                        DescriptorLookupConverter.createLookupElement(getResolveSession(), getExpressionBindingContext(), descriptor))
-                       : Collections.<LookupElement>emptyList();
-            }
-        });
-    }
-
-    private void addReferenceVariants(@NotNull Function1<DeclarationDescriptor, Iterable<LookupElement>> filter) {
         Collection<DeclarationDescriptor> descriptors = TipsManager.getReferenceVariants(
                 jetReference.getExpression(), getExpressionBindingContext());
 
-        for (DeclarationDescriptor descriptor : descriptors) {
+        Collection<DeclarationDescriptor> filterDescriptors = Collections2.filter(descriptors, new Predicate<DeclarationDescriptor>() {
-            if (descriptor != null && descriptorFilter.value(descriptor)) {
+            @Override
-                Iterable<LookupElement> elements = filter.invoke(descriptor);
+            public boolean apply(@Nullable DeclarationDescriptor descriptor) {
-                for (LookupElement element : elements) {
+                return descriptor != null && filterCondition.value(descriptor);
-                    jetResult.addElement(element);
-                }
             }
-        }
+        });
+
+        jetResult.addAllElements(filterDescriptors);
     }
 
     private boolean isVisibleDescriptor(DeclarationDescriptor descriptor) {

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

@@ -34,11 +34,6 @@ import org.jetbrains.jet.lang.resolve.scopes.receivers.ReceiverValue
 import com.intellij.lang.ASTNode
 import org.jetbrains.jet.lang.resolve.scopes.JetScope
 
-trait SmartCompletionData{
-    fun toElements(descriptor: DeclarationDescriptor): Iterable<LookupElement>
-    val additionalElements: Iterable<LookupElement>
-}
-
 enum class Tail {
   COMMA
   PARENTHESIS
@@ -48,7 +43,8 @@ data class ExpectedTypeInfo(val `type`: JetType, val tail: Tail?)
 
 fun buildSmartCompletionData(expression: JetSimpleNameExpression,
                              resolveSession: ResolveSessionForBodies,
-                             visibilityFilter: (DeclarationDescriptor) -> Boolean): SmartCompletionData? {
+                             referenceVariants: Iterable<DeclarationDescriptor>,
+                             visibilityFilter: (DeclarationDescriptor) -> Boolean): Collection<LookupElement>? {
     val parent = expression.getParent()
     val expressionWithType: JetExpression;
     val receiver: JetExpression?
@@ -67,104 +63,72 @@ fun buildSmartCompletionData(expression: JetSimpleNameExpression,
     val expectedTypes = allExpectedTypes.filter { !it.`type`.isError() }
     if (expectedTypes.isEmpty()) return null
 
+    val result = ArrayList<LookupElement>()
+
+    val typesOf: (DeclarationDescriptor) -> Iterable<JetType> = dataFlowToDescriptorTypes(expressionWithType, receiver, bindingContext)
+
     val itemsToSkip = calcItemsToSkip(expressionWithType, resolveSession)
 
-    val additionalElements = ArrayList<LookupElement>()
+    for (descriptor in referenceVariants) {
+        if (itemsToSkip.contains(descriptor)) continue
+
+        run {
+            val matchedExpectedTypes = expectedTypes.filter { expectedType ->
+                typesOf(descriptor).any { descriptorType -> isSubtypeOf(descriptorType, expectedType.`type`) }
+            }
+            if (matchedExpectedTypes.isNotEmpty()) {
+                val lookupElement = DescriptorLookupConverter.createLookupElement(resolveSession, bindingContext, descriptor)
+                result.add(addTailToLookupElement(lookupElement, matchedExpectedTypes))
+            }
+        }
+
+        val functionExpectedTypes = expectedTypes.filter { KotlinBuiltIns.getInstance().isExactFunctionOrExtensionFunctionType(it.`type`) }
+        if (functionExpectedTypes.isNotEmpty()) {
+            fun functionReferenceLookupElement(descriptor: FunctionDescriptor): LookupElement? {
+                val functionType = functionType(descriptor)
+                if (functionType == null) return null
+
+                val matchedExpectedTypes = functionExpectedTypes.filter { isSubtypeOf(functionType, it.`type`) }
+                if (matchedExpectedTypes.isEmpty()) return null
+                val lookupElement = DescriptorLookupConverter.createLookupElement(resolveSession, bindingContext, descriptor)
+                val text = "::" + (if (descriptor is ConstructorDescriptor) descriptor.getContainingDeclaration().getName() else descriptor.getName())
+                val lookupElementDecorated = object: LookupElementDecorator<LookupElement>(lookupElement) {
+                    override fun getLookupString() = text
+
+                    override fun renderElement(presentation: LookupElementPresentation) {
+                        super.renderElement(presentation)
+                        presentation.setItemText(text)
+                        presentation.setTypeText("")
+                    }
+
+                    override fun handleInsert(context: InsertionContext) {
+                    }
+                }
+
+                return addTailToLookupElement(lookupElementDecorated, matchedExpectedTypes)
+            }
+
+            if (descriptor is SimpleFunctionDescriptor) {
+                functionReferenceLookupElement(descriptor)?.let { result.add(it) }
+            }
+            else if (descriptor is ClassDescriptor && descriptor.getModality() != Modality.ABSTRACT) {
+                val constructors = descriptor.getConstructors().filter(visibilityFilter)
+                if (constructors.size == 1) { //TODO: this code is to be changed if overloads to start work after ::
+                    functionReferenceLookupElement(constructors.single())?.let { result.add(it) }
+                }
+            }
+        }
+    }
 
     if (receiver == null) {
-        additionalElements.addTypeInstantiationItems(expectedTypes, resolveSession, bindingContext)
+        result.addTypeInstantiationItems(expectedTypes, resolveSession, bindingContext)
 
-        additionalElements.addStaticMembers(expressionWithType, expectedTypes, resolveSession, bindingContext)
+        result.addStaticMembers(expressionWithType, expectedTypes, resolveSession, bindingContext)
 
-        additionalElements.addThisItems(expressionWithType, expectedTypes, bindingContext)
+        result.addThisItems(expressionWithType, expectedTypes, bindingContext)
     }
 
-    val dataFlowInfo = bindingContext[BindingContext.EXPRESSION_DATA_FLOW_INFO, expressionWithType]
+    return result
-    val (variableToTypes: Map<VariableDescriptor, Collection<JetType>>, notNullVariables: Set<VariableDescriptor>) = processDataFlowInfo(dataFlowInfo, receiver, bindingContext)
-
-    fun typesOf(descriptor: DeclarationDescriptor): Iterable<JetType> {
-        if (descriptor is CallableDescriptor) {
-            var returnType = descriptor.getReturnType()
-            if (returnType != null && KotlinBuiltIns.getInstance().isNothing(returnType!!)) { //TODO: maybe we should include them on the second press?
-                return listOf()
-            }
-            if (descriptor is VariableDescriptor) {
-                if (notNullVariables.contains(descriptor) && returnType != null) {
-                    returnType = TypeUtils.makeNotNullable(returnType!!)
-                }
-
-                val autoCastTypes = variableToTypes[descriptor]
-                if (autoCastTypes != null && !autoCastTypes.isEmpty()) {
-                    return autoCastTypes + returnType.toList()
-                }
-            }
-            return returnType.toList()
-        }
-        else if (descriptor is ClassDescriptor && descriptor.getKind() == ClassKind.ENUM_ENTRY) {
-            return listOf(descriptor.getDefaultType())
-        }
-        else {
-            return listOf()
-        }
-    }
-
-    return object: SmartCompletionData {
-        override fun toElements(descriptor: DeclarationDescriptor): Iterable<LookupElement> {
-            if (itemsToSkip.contains(descriptor)) return listOf()
-
-            val result = ArrayList<LookupElement>()
-
-            run {
-                val matchedExpectedTypes = expectedTypes.filter { expectedType -> typesOf(descriptor).any { descriptorType -> isSubtypeOf(descriptorType, expectedType.`type`) } }
-                if (matchedExpectedTypes.isNotEmpty()) {
-                    val tail = mergeTails(matchedExpectedTypes.map { it.tail })
-                    result.add(addTailToLookupElement(DescriptorLookupConverter.createLookupElement(resolveSession, bindingContext, descriptor), tail))
-                }
-            }
-
-            val functionExpectedTypes = expectedTypes.filter { KotlinBuiltIns.getInstance().isExactFunctionOrExtensionFunctionType(it.`type`) }
-            if (functionExpectedTypes.isNotEmpty()) {
-                fun functionReferenceLookupElement(descriptor: FunctionDescriptor): LookupElement? {
-                    val functionType = functionType(descriptor)
-                    if (functionType == null) return null
-
-                    val matchedExpectedTypes = functionExpectedTypes.filter { isSubtypeOf(functionType, it.`type`) }
-                    if (matchedExpectedTypes.isEmpty()) return null
-                    val lookupElement = DescriptorLookupConverter.createLookupElement(resolveSession, bindingContext, descriptor)
-                    val text = "::" + (if (descriptor is ConstructorDescriptor) descriptor.getContainingDeclaration().getName() else descriptor.getName())
-                    val lookupElementDecorated = object: LookupElementDecorator<LookupElement>(lookupElement) {
-                        override fun getLookupString() = text
-
-                        override fun renderElement(presentation: LookupElementPresentation) {
-                            super.renderElement(presentation)
-                            presentation.setItemText(text)
-                            presentation.setTypeText("")
-                        }
-
-                        override fun handleInsert(context: InsertionContext) {
-                        }
-                    }
-
-                    val tail = mergeTails(matchedExpectedTypes.map { it.tail })
-                    return addTailToLookupElement(lookupElementDecorated, tail)
-                }
-
-                if (descriptor is SimpleFunctionDescriptor) {
-                    functionReferenceLookupElement(descriptor)?.let { result.add(it) }
-                }
-                else if (descriptor is ClassDescriptor && descriptor.getModality() != Modality.ABSTRACT) {
-                    val constructors = descriptor.getConstructors().filter(visibilityFilter)
-                    if (constructors.size == 1) { //TODO: this code is to be changed if overloads to start work after ::
-                        functionReferenceLookupElement(constructors.single())?.let { result.add(it) }
-                    }
-                }
-            }
-
-            return result
-        }
-
-        override val additionalElements = additionalElements
-    }
 }
 
 private fun calcExpectedTypes(expressionWithType: JetExpression, bindingContext: BindingContext, moduleDescriptor: ModuleDescriptor): Collection<ExpectedTypeInfo>? {
@@ -364,8 +328,7 @@ private fun MutableCollection<LookupElement>.addThisItems(context: JetExpression
             val qualifier = if (i == 0) null else thisQualifierName(receiver, bindingContext) ?: continue
             val expressionText = if (qualifier == null) "this" else "this@" + qualifier
             val lookupElement = LookupElementBuilder.create(expressionText).withTypeText(DescriptorRenderer.TEXT.renderType(thisType))
-            val tailType = mergeTails(matchedExpectedTypes.map { it.tail })
+            add(addTailToLookupElement(lookupElement, matchedExpectedTypes))
-            add(addTailToLookupElement(lookupElement, tailType))
         }
     }
 }
@@ -388,6 +351,40 @@ private fun thisQualifierName(receiver: ReceiverParameterDescriptor, bindingCont
                             ?.getReferencedName()
 }
 
+private fun dataFlowToDescriptorTypes(expression: JetExpression, receiver: JetExpression?, bindingContext: BindingContext): (DeclarationDescriptor) -> Iterable<JetType> {
+    val dataFlowInfo = bindingContext[BindingContext.EXPRESSION_DATA_FLOW_INFO, expression]
+    val (variableToTypes: Map<VariableDescriptor, Collection<JetType>>, notNullVariables: Set<VariableDescriptor>)
+        = processDataFlowInfo(dataFlowInfo, receiver, bindingContext)
+
+    fun typesOf(descriptor: DeclarationDescriptor): Iterable<JetType> {
+        if (descriptor is CallableDescriptor) {
+            var returnType = descriptor.getReturnType()
+            if (returnType != null && KotlinBuiltIns.getInstance().isNothing(returnType!!)) { //TODO: maybe we should include them on the second press?
+                return listOf()
+            }
+            if (descriptor is VariableDescriptor) {
+                if (notNullVariables.contains(descriptor) && returnType != null) {
+                    returnType = TypeUtils.makeNotNullable(returnType!!)
+                }
+
+                val autoCastTypes = variableToTypes[descriptor]
+                if (autoCastTypes != null && !autoCastTypes.isEmpty()) {
+                    return autoCastTypes + returnType.toList()
+                }
+            }
+            return returnType.toList()
+        }
+        else if (descriptor is ClassDescriptor && descriptor.getKind() == ClassKind.ENUM_ENTRY) {
+            return listOf(descriptor.getDefaultType())
+        }
+        else {
+            return listOf()
+        }
+    }
+
+    return ::typesOf
+}
+
 private data class ProcessDataFlowInfoResult(
         val variableToTypes: Map<VariableDescriptor, Collection<JetType>> = Collections.emptyMap(),
         val notNullVariables: Set<VariableDescriptor> = Collections.emptySet()
@@ -550,8 +547,7 @@ private fun MutableCollection<LookupElement>.addStaticMembers(classDescriptor: C
             }
         }
 
-        val tail = mergeTails(descriptorExpectedTypes.map { it.tail })
+        add(addTailToLookupElement(lookupElementDecorated, descriptorExpectedTypes))
-        add(addTailToLookupElement(lookupElementDecorated, tail))
     }
 }
 
@@ -578,6 +574,9 @@ private fun addTailToLookupElement(lookupElement: LookupElement, tail: Tail?): L
     }
 }
 
+private fun addTailToLookupElement(lookupElement: LookupElement, expectedTypes: Collection<ExpectedTypeInfo>): LookupElement
+    = addTailToLookupElement(lookupElement, mergeTails(expectedTypes.map { it.tail }))
+
 private fun functionType(function: FunctionDescriptor): JetType? {
     return KotlinBuiltIns.getInstance().getKFunctionType(function.getAnnotations(),
                                                          null,

idea/src/org/jetbrains/jet/plugin/completion/handlers/WithTailInsertHandler.kt

@@ -1,9 +1,8 @@
-package org.jetbrains.jet.plugin.completion
+package org.jetbrains.jet.plugin.completion.handlers
 
-import com.intellij.codeInsight.lookup.*
 import com.intellij.codeInsight.completion.*
-import com.intellij.openapi.editor.event.DocumentListener
+import com.intellij.codeInsight.lookup.LookupElement
-import com.intellij.openapi.editor.event.DocumentEvent
+import com.intellij.openapi.editor.event.*
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiDocumentManager
 
@@ -29,18 +28,22 @@ class WithTailInsertHandler(val tailChar: Char, val spaceAfter: Boolean) : Inser
         }
 
         document.addDocumentListener(documentListener)
-
+        try{
-        item.handleInsert(context)
+            item.handleInsert(context)
-        PsiDocumentManager.getInstance(context.getProject()).doPostponedOperationsAndUnblockDocument(document)
+            PsiDocumentManager.getInstance(context.getProject()).doPostponedOperationsAndUnblockDocument(document)
-
+        }
-        document.removeDocumentListener(documentListener)
+        finally {
+            document.removeDocumentListener(documentListener)
+        }
 
         val moveCaret = caretModel.getOffset() == maxChangeOffset
 
-        if (maxChangeOffset < document.getTextLength() && document.getText(TextRange(maxChangeOffset, maxChangeOffset + 1))[0] == tailChar) {
+        fun isCharAt(offset: Int, c: Char) = offset < document.getTextLength() && document.getText(TextRange(offset, offset + 1))[0] == c
+
+        if (isCharAt(maxChangeOffset, tailChar)) {
             document.deleteString(maxChangeOffset, maxChangeOffset + 1)
 
-            if (spaceAfter && maxChangeOffset < document.getTextLength() && document.getText(TextRange(maxChangeOffset, maxChangeOffset + 1)) == " ") {
+            if (spaceAfter && isCharAt(maxChangeOffset, ' ')) {
                 document.deleteString(maxChangeOffset, maxChangeOffset + 1)
             }
         }