FIR IDE: add type arguments and -> symbol to when conditions completion

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/FirCompletionDummyIdentifierProviderService.kt

@@ -5,14 +5,21 @@
 
 package org.jetbrains.kotlin.idea.completion
 
-import org.jetbrains.kotlin.idea.frontend.api.analyse
+import com.intellij.codeInsight.completion.CompletionInitializationContext
-import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassKind
+import com.intellij.codeInsight.completion.CompletionUtilCore
-import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassOrObjectSymbol
+import com.intellij.psi.util.parentOfType
-import org.jetbrains.kotlin.idea.frontend.api.symbols.KtFunctionSymbol
-import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.psi.KtNameReferenceExpression
+import org.jetbrains.kotlin.psi.KtWhenEntry
 
 class FirCompletionDummyIdentifierProviderService : CompletionDummyIdentifierProviderService() {
+    override fun handleDefaultCase(context: CompletionInitializationContext): String? {
+        val elementAtOffset = context.file.findElementAt(context.startOffset) ?: return null
+        return when {
+            elementAtOffset.parentOfType<KtWhenEntry>() != null -> CompletionUtilCore.DUMMY_IDENTIFIER_TRIMMED
+            else -> null
+        }
+    }
+
     override fun allTargetsAreFunctionsOrClasses(nameReferenceExpression: KtNameReferenceExpression): Boolean {
         return true
         // TODO fir cannot handle invalid code and handles listOf< as binary expression

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/contributors/FirWhenWithSubjectConditionContributor.kt

@@ -24,10 +24,15 @@ import org.jetbrains.kotlin.idea.frontend.api.symbols.*
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.idea.completion.contributors.helpers.FirClassifierProvider.getAvailableClassifiersCurrentScope
+import org.jetbrains.kotlin.idea.completion.contributors.helpers.addSymbolToCompletion
+import org.jetbrains.kotlin.idea.completion.contributors.helpers.addTypeArguments
+import org.jetbrains.kotlin.idea.completion.contributors.helpers.createStarTypeArgumentsList
 import org.jetbrains.kotlin.idea.completion.lookups.KotlinLookupObject
 import org.jetbrains.kotlin.idea.completion.lookups.shortenReferencesForFirCompletion
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtNamedSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSymbolWithTypeParameters
 import org.jetbrains.kotlin.idea.frontend.api.types.*
+import org.jetbrains.kotlin.miniStdLib.letIf
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -44,15 +49,22 @@ internal class FirWhenWithSubjectConditionContributor(
         val subjectType = subject.getKtType()
         val classSymbol = getClassSymbol(subjectType)
         val visibilityChecker = CompletionVisibilityChecker.create(basicContext, positionContext)
+        val isSingleCondition = whenCondition.isSingleConditionInEntry()
         when {
             classSymbol?.classKind == KtClassKind.ENUM_CLASS -> {
-                completeEnumEntries(classSymbol, allConditionsExceptCurrent, visibilityChecker)
+                completeEnumEntries(classSymbol, allConditionsExceptCurrent, visibilityChecker, isSingleCondition)
             }
             classSymbol?.modality == Modality.SEALED -> {
-                completeSubClassesOfSealedClass(classSymbol, allConditionsExceptCurrent, whenCondition, visibilityChecker)
+                completeSubClassesOfSealedClass(
+                    classSymbol,
+                    allConditionsExceptCurrent,
+                    whenCondition,
+                    visibilityChecker,
+                    isSingleCondition
+                )
             }
             else -> {
-                completeAllTypes(whenCondition, visibilityChecker)
+                completeAllTypes(whenCondition, visibilityChecker, isSingleCondition)
             }
         }
         addNullIfWhenExpressionCanReturnNull(subjectType)
@@ -75,6 +87,7 @@ internal class FirWhenWithSubjectConditionContributor(
     private fun KtAnalysisSession.completeAllTypes(
         whenCondition: KtWhenCondition,
         visibilityChecker: CompletionVisibilityChecker,
+        isSingleCondition: Boolean,
     ) {
         getAvailableClassifiersCurrentScope(originalKtFile, whenCondition, scopeNameFilter, indexHelper, visibilityChecker)
             .forEach { classifier ->
@@ -84,7 +97,8 @@ internal class FirWhenWithSubjectConditionContributor(
                     classifier.name.asString(),
                     classifier,
                     (classifier as? KtNamedClassOrObjectSymbol)?.classIdIfNonLocal?.asSingleFqName(),
-                    isPrefixNeeded(classifier)
+                    isPrefixNeeded(classifier),
+                    isSingleCondition,
                 )
             }
     }
@@ -104,6 +118,7 @@ internal class FirWhenWithSubjectConditionContributor(
         conditions: List<KtWhenCondition>,
         whenCondition: KtWhenCondition,
         visibilityChecker: CompletionVisibilityChecker,
+        isSingleCondition: Boolean,
     ) {
         require(classSymbol.modality == Modality.SEALED)
         val handledCasesClassIds = getHandledClassIds(conditions)
@@ -115,11 +130,17 @@ internal class FirWhenWithSubjectConditionContributor(
             .filter { with(visibilityChecker) { isVisible(it as KtClassifierSymbol) } }
             .forEach { inheritor ->
                 val classId = inheritor.classIdIfNonLocal ?: return@forEach
-                addLookupElement(classId.relativeClassName.asString(), inheritor, classId.asSingleFqName(), isPrefixNeeded(inheritor))
+                addLookupElement(
+                    classId.relativeClassName.asString(),
+                    inheritor,
+                    classId.asSingleFqName(),
+                    isPrefixNeeded(inheritor),
+                    isSingleCondition
+                )
             }
 
         if (allInheritors.any { it.modality == Modality.ABSTRACT }) {
-            completeAllTypes(whenCondition, visibilityChecker)
+            completeAllTypes(whenCondition, visibilityChecker, isSingleCondition)
         }
     }
 
@@ -162,7 +183,8 @@ internal class FirWhenWithSubjectConditionContributor(
     private fun KtAnalysisSession.completeEnumEntries(
         classSymbol: KtNamedClassOrObjectSymbol,
         conditions: List<KtWhenCondition>,
-        visibilityChecker: CompletionVisibilityChecker
+        visibilityChecker: CompletionVisibilityChecker,
+        isSingleCondition: Boolean,
     ) {
         require(classSymbol.classKind == KtClassKind.ENUM_CLASS)
         val handledCasesNames = conditions.mapNotNullTo(hashSetOf()) { condition ->
@@ -179,18 +201,33 @@ internal class FirWhenWithSubjectConditionContributor(
                     "${classSymbol.name.asString()}.${entry.name.asString()}",
                     entry,
                     entry.callableIdIfNonLocal?.asSingleFqName(),
-                    isPrefixNeeded = false
+                    isPrefixNeeded = false,
+                    isSingleCondition,
                 )
             }
     }
 
-    private fun KtAnalysisSession.addLookupElement(lookupString: String, symbol: KtNamedSymbol, fqName: FqName?, isPrefixNeeded: Boolean) {
+    private fun KtWhenCondition.isSingleConditionInEntry(): Boolean {
-        val lookupObject = WhenConditionLookupObject(symbol.name, fqName, isPrefixNeeded)
+        val entry = parent as KtWhenEntry
+        return entry.conditions.size == 1
+    }
+
+    private fun KtAnalysisSession.addLookupElement(
+        lookupString: String,
+        symbol: KtNamedSymbol,
+        fqName: FqName?,
+        isPrefixNeeded: Boolean,
+        isSingleCondition: Boolean
+    ) {
+        val typeArgumentsCount = (symbol as? KtSymbolWithTypeParameters)?.typeParameters?.size ?: 0
+        val lookupObject = WhenConditionLookupObject(symbol.name, fqName, isPrefixNeeded, isSingleCondition, typeArgumentsCount)
 
         LookupElementBuilder.create(lookupObject, getIsPrefix(isPrefixNeeded) + lookupString)
             .withIcon(getIconFor(symbol))
             .withPsiElement(symbol.psi)
             .withInsertHandler(WhenConditionInsertionHandler)
+            .withTailText(createStarTypeArgumentsList(typeArgumentsCount), /*grayed*/true)
+            .letIf(isSingleCondition) { it.appendTailText(" -> ",  /*grayed*/true) }
             .let(sink::addElement)
     }
 }
@@ -199,20 +236,41 @@ private data class WhenConditionLookupObject(
     override val shortName: Name,
     val fqName: FqName?,
     val needIsPrefix: Boolean,
+    val isSingleCondition: Boolean,
+    val typeArgumentsCount: Int,
 ) : KotlinLookupObject
 
+
 private object WhenConditionInsertionHandler : InsertionHandlerBase<WhenConditionLookupObject>(WhenConditionLookupObject::class) {
     override fun handleInsert(context: InsertionContext, item: LookupElement, ktFile: KtFile, lookupObject: WhenConditionLookupObject) {
+        context.insertName(lookupObject, ktFile)
+        context.addTypeArguments(lookupObject.typeArgumentsCount)
+        context.addArrow(lookupObject)
+    }
+
+    private fun InsertionContext.addArrow(
+        lookupObject: WhenConditionLookupObject
+    ) {
+        if (lookupObject.isSingleCondition && completionChar != ',') {
+            addSymbolToCompletion(" -> ")
+            commitDocument()
+        }
+    }
+
+    private fun InsertionContext.insertName(
+        lookupObject: WhenConditionLookupObject,
+        ktFile: KtFile
+    ) {
         if (lookupObject.fqName != null) {
             val fqName = lookupObject.fqName
-            context.document.replaceString(
+            document.replaceString(
-                context.startOffset,
+                startOffset,
-                context.tailOffset,
+                tailOffset,
                 getIsPrefix(lookupObject.needIsPrefix) + fqName.render()
             )
-            context.commitDocument()
+            commitDocument()
 
-            shortenReferencesForFirCompletion(ktFile, TextRange(context.startOffset, context.tailOffset))
+            shortenReferencesForFirCompletion(ktFile, TextRange(startOffset, tailOffset))
         }
     }
 }

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/contributors/helpers/completionHelpers.kt

@@ -9,9 +9,10 @@ import com.intellij.codeInsight.completion.CodeCompletionHandlerBase
 import com.intellij.codeInsight.completion.CompletionType
 import com.intellij.codeInsight.completion.InsertionContext
 import com.intellij.openapi.application.ApplicationManager
+import com.intellij.psi.PsiDocumentManager
 
 internal fun InsertionContext.addSymbolAndInvokeCompletion(symbol: String) {
-    addSymbolToCompletion(this, symbol)
+    this.addSymbolToCompletion(symbol)
     invokeCompletion(this)
 }
 
@@ -24,8 +25,24 @@ private fun invokeCompletion(context: InsertionContext) {
     }
 }
 
-private fun addSymbolToCompletion(context: InsertionContext, symbol: String) {
+internal fun InsertionContext.addSymbolToCompletion(symbol: String) {
-    context.document.insertString(context.tailOffset, symbol)
+    PsiDocumentManager.getInstance(project).doPostponedOperationsAndUnblockDocument(document)
-    context.commitDocument()
+    document.insertString(tailOffset, symbol)
-    context.editor.caretModel.moveToOffset(context.tailOffset)
+    commitDocument()
+    editor.caretModel.moveToOffset(tailOffset)
+}
+
+internal fun InsertionContext.addTypeArguments(typeArgumentsCount: Int) {
+    when {
+        typeArgumentsCount == 0 -> {
+            return
+        }
+        typeArgumentsCount < 0 -> {
+            error("Count of type arguments should be non-negative, but was $typeArgumentsCount")
+        }
+        else -> {
+            commitDocument()
+            addSymbolToCompletion(createStarTypeArgumentsList(typeArgumentsCount))
+        }
+    }
 }

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/contributors/helpers/utils.kt

@@ -0,0 +1,13 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.completion.contributors.helpers
+
+internal fun createStarTypeArgumentsList(typeArgumentsCount: Int): String =
+    if (typeArgumentsCount > 0) {
+        List(typeArgumentsCount) { "*" }.joinToString(prefix = "<", postfix = ">")
+    } else {
+        ""
+    }

idea/idea-fir/testData/completion/basic/common/when/nonSealedClassWithTypeArgs.kt

@@ -0,0 +1,16 @@
+class NON_SEALED
+class AAAA<E, S>: NON_SEALED()
+object BBBB: NON_SEALED()
+class CCCC<E>: NON_SEALED()
+
+fun foo(e: NON_SEALED) {
+    when (e) {
+        <caret>
+    }
+}
+
+// EXIST: { lookupString: "is AAAA", tailText: "<*, *> -> " }
+// EXIST: BBBB
+// EXIST: { lookupString: "is CCCC", tailText: "<*> -> " }
+// EXIST: { lookupString: "else -> "}
+// FIR_COMPARISON

idea/idea-fir/testData/completion/basic/common/when/sealedClassWithTypeArgs.kt

@@ -0,0 +1,17 @@
+sealed class SEALED
+class AAAA<E, S>: SEALED()
+object BBBB: SEALED()
+class CCCC<E>: SEALED()
+
+fun foo(e: SEALED) {
+    when (e) {
+        <caret>
+    }
+}
+
+// EXIST: { lookupString: "is AAAA", tailText: "<*, *> -> " }
+// EXIST: BBBB
+// EXIST: { lookupString: "is CCCC", tailText: "<*> -> " }
+// EXIST: { lookupString: "else -> "}
+// NOTHING_ELSE
+// FIR_COMPARISON

idea/idea-fir/tests/org/jetbrains/kotlin/idea/fir/completion/HighLevelJvmBasicCompletionTestGenerated.java

@@ -3382,6 +3382,11 @@ public class HighLevelJvmBasicCompletionTestGenerated extends AbstractHighLevelJ
                 runTest("idea/idea-completion/testData/../../idea-fir/testData/completion/basic/common/when/nonSealedClassInWhen.kt");
             }
 
+            @TestMetadata("nonSealedClassWithTypeArgs.kt")
+            public void testNonSealedClassWithTypeArgs() throws Exception {
+                runTest("idea/idea-completion/testData/../../idea-fir/testData/completion/basic/common/when/nonSealedClassWithTypeArgs.kt");
+            }
+
             @TestMetadata("nullableEnumInWhen.kt")
             public void testNullableEnumInWhen() throws Exception {
                 runTest("idea/idea-completion/testData/../../idea-fir/testData/completion/basic/common/when/nullableEnumInWhen.kt");
@@ -3407,6 +3412,11 @@ public class HighLevelJvmBasicCompletionTestGenerated extends AbstractHighLevelJ
                 runTest("idea/idea-completion/testData/../../idea-fir/testData/completion/basic/common/when/sealedClassWithAbstractInheritorsInWhen.kt");
             }
 
+            @TestMetadata("sealedClassWithTypeArgs.kt")
+            public void testSealedClassWithTypeArgs() throws Exception {
+                runTest("idea/idea-completion/testData/../../idea-fir/testData/completion/basic/common/when/sealedClassWithTypeArgs.kt");
+            }
+
             @TestMetadata("sealedInterfaceInWhen.kt")
             public void testSealedInterfaceInWhen() throws Exception {
                 runTest("idea/idea-completion/testData/../../idea-fir/testData/completion/basic/common/when/sealedInterfaceInWhen.kt");

idea/idea-frontend-independent/src/org/jetbrains/kotlin/idea/completion/CompletionDummyIdentifierProviderService.kt

@@ -65,10 +65,13 @@ abstract class CompletionDummyIdentifierProviderService {
                 ?: specialInTypeArgsDummyIdentifier(tokenBefore)
                 ?: specialInArgumentListDummyIdentifier(tokenBefore)
                 ?: isInTypeParametersList(tokenBefore)
+                ?: handleDefaultCase(context)
                 ?: DEFAULT_DUMMY_IDENTIFIER
         }
     }
 
+    protected open fun handleDefaultCase(context: CompletionInitializationContext): String? = null
+
     private fun isInTypeParametersList(tokenBefore: PsiElement?): String? {
         if (tokenBefore == null) return null
         if (tokenBefore.parents.any { it is KtTypeParameterList }) {