FIR IDE: implement expected type weigher for completion

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

@@ -6,9 +6,11 @@
 package org.jetbrains.kotlin.idea.completion
 
 import com.intellij.codeInsight.completion.*
+import com.intellij.codeInsight.lookup.LookupElement
 import com.intellij.patterns.PlatformPatterns
 import com.intellij.patterns.PsiJavaPatterns
 import com.intellij.util.ProcessingContext
+import org.jetbrains.kotlin.idea.completion.weighers.Weighers
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalKtFile
 import org.jetbrains.kotlin.idea.frontend.api.InvalidWayOfUsingAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
@@ -36,9 +38,17 @@ private object KotlinFirCompletionProvider : CompletionProvider<CompletionParame
     override fun addCompletions(parameters: CompletionParameters, context: ProcessingContext, result: CompletionResultSet) {
         if (shouldSuppressCompletion(parameters, result.prefixMatcher)) return
 
-        KotlinAvailableScopesCompletionProvider(result.prefixMatcher).addCompletions(parameters, result)
+        val resultSet = createResultSet(parameters, result)
+        KotlinAvailableScopesCompletionProvider(resultSet.prefixMatcher).addCompletions(parameters, resultSet)
     }
 
+    private fun createResultSet(parameters: CompletionParameters, result: CompletionResultSet): CompletionResultSet =
+        result.withRelevanceSorter(createSorter(parameters, result))
+
+    private fun createSorter(parameters: CompletionParameters, result: CompletionResultSet): CompletionSorter =
+        CompletionSorter.defaultSorter(parameters, result.prefixMatcher)
+            .let(Weighers::addWeighersToCompletionSorter)
+
     private val AFTER_NUMBER_LITERAL = PsiJavaPatterns.psiElement().afterLeafSkipping(
         PsiJavaPatterns.psiElement().withText(""),
         PsiJavaPatterns.psiElement().withElementType(PsiJavaPatterns.elementType().oneOf(KtTokens.FLOAT_LITERAL, KtTokens.INTEGER_LITERAL))
@@ -68,9 +78,21 @@ private class KotlinAvailableScopesCompletionProvider(prefixMatcher: PrefixMatch
     private val scopeNameFilter: KtScopeNameFilter =
         { name -> !name.isSpecial && prefixMatcher.prefixMatches(name.identifier) }
 
-    private fun KtAnalysisSession.addSymbolToCompletion(completionResultSet: CompletionResultSet, symbol: KtSymbol) {
+    private fun KtAnalysisSession.addSymbolToCompletion(completionResultSet: CompletionResultSet, expectedType: KtType?, symbol: KtSymbol) {
         if (symbol !is KtNamedSymbol) return
-        with(lookupElementFactory) { createLookupElement(symbol)?.let(completionResultSet::addElement) }
+        with(lookupElementFactory) {
+            createLookupElement(symbol)
+                ?.let { applyWeighers(it, symbol, expectedType) }
+                ?.let(completionResultSet::addElement)
+        }
+    }
+
+    private fun KtAnalysisSession.applyWeighers(
+        lookupElement: LookupElement,
+        symbol: KtSymbol,
+        expectedType: KtType?
+    ): LookupElement = lookupElement.apply {
+        with(Weighers) { applyWeighsToLookupElement(lookupElement, symbol, expectedType) }
     }
 
     private fun recordOriginalFile(completionParameters: CompletionParameters) {
@@ -90,29 +112,41 @@ private class KotlinAvailableScopesCompletionProvider(prefixMatcher: PrefixMatch
         val explicitReceiver = nameExpression.getReceiverExpression()
 
         with(getAnalysisSessionFor(originalFile).createContextDependentCopy()) {
+            val expectedType = nameExpression.getExpectedType()
+
             val (implicitScopes, _) = originalFile.getScopeContextForPosition(nameExpression)
 
             fun KtCallableSymbol.hasSuitableExtensionReceiver(): Boolean =
                 checkExtensionIsSuitable(originalFile, nameExpression, explicitReceiver)
 
             when {
-                nameExpression.parent is KtUserType -> collectTypesCompletion(result, implicitScopes)
-                explicitReceiver != null ->
-                    collectDotCompletion(result, implicitScopes, explicitReceiver, KtCallableSymbol::hasSuitableExtensionReceiver)
-                else -> collectDefaultCompletion(result, implicitScopes, KtCallableSymbol::hasSuitableExtensionReceiver)
+                nameExpression.parent is KtUserType -> collectTypesCompletion(result, implicitScopes, expectedType)
+                explicitReceiver != null -> collectDotCompletion(
+                    result,
+                    implicitScopes,
+                    explicitReceiver,
+                    expectedType,
+                    KtCallableSymbol::hasSuitableExtensionReceiver
+                )
+                else -> collectDefaultCompletion(result, implicitScopes, expectedType, KtCallableSymbol::hasSuitableExtensionReceiver)
             }
         }
     }
 
-    private fun KtAnalysisSession.collectTypesCompletion(result: CompletionResultSet, implicitScopes: KtScope) {
+    private fun KtAnalysisSession.collectTypesCompletion(
+        result: CompletionResultSet,
+        implicitScopes: KtScope,
+        expectedType: KtType?,
+    ) {
         val availableClasses = implicitScopes.getClassifierSymbols(scopeNameFilter)
-        availableClasses.forEach { addSymbolToCompletion(result, it) }
+        availableClasses.forEach { addSymbolToCompletion(result, expectedType, it) }
     }
 
     private fun KtAnalysisSession.collectDotCompletion(
         result: CompletionResultSet,
         implicitScopes: KtCompositeScope,
         explicitReceiver: KtExpression,
+        expectedType: KtType?,
         hasSuitableExtensionReceiver: KtCallableSymbol.() -> Boolean
     ) {
         val typeOfPossibleReceiver = explicitReceiver.getKtType()
@@ -126,13 +160,14 @@ private class KotlinAvailableScopesCompletionProvider(prefixMatcher: PrefixMatch
             .getCallableSymbols(scopeNameFilter)
             .filter { it.isExtension && it.hasSuitableExtensionReceiver() }
 
-        nonExtensionMembers.forEach { addSymbolToCompletion(result, it) }
-        extensionNonMembers.forEach { addSymbolToCompletion(result, it) }
+        nonExtensionMembers.forEach { addSymbolToCompletion(result, expectedType, it) }
+        extensionNonMembers.forEach { addSymbolToCompletion(result, expectedType, it) }
     }
 
     private fun KtAnalysisSession.collectDefaultCompletion(
         result: CompletionResultSet,
         implicitScopes: KtCompositeScope,
+        expectedType: KtType?,
         hasSuitableExtensionReceiver: KtCallableSymbol.() -> Boolean,
     ) {
         val availableNonExtensions = implicitScopes
@@ -143,19 +178,9 @@ private class KotlinAvailableScopesCompletionProvider(prefixMatcher: PrefixMatch
             .getCallableSymbols(scopeNameFilter)
             .filter { it.isExtension && it.hasSuitableExtensionReceiver() }
 
-        availableNonExtensions.forEach { addSymbolToCompletion(result, it) }
-        extensionsWhichCanBeCalled.forEach { addSymbolToCompletion(result, it) }
+        availableNonExtensions.forEach { addSymbolToCompletion(result, expectedType, it) }
+        extensionsWhichCanBeCalled.forEach { addSymbolToCompletion(result, expectedType, it) }
 
-        collectTypesCompletion(result, implicitScopes)
-    }
-}
-
-private object ExpectedTypeProvider {
-    fun KtAnalysisSession.getExpectedType(nameExpression: KtSimpleNameExpression, parameters: CompletionParameters): KtType? =
-        getExpectedTypeByReturnExpression(nameExpression)
-
-    private fun KtAnalysisSession.getExpectedTypeByReturnExpression(nameExpression: KtSimpleNameExpression): KtType? {
-        val parentReturn = nameExpression.parent as? KtReturnExpression ?: return null
-        TODO()
+        collectTypesCompletion(result, implicitScopes, expectedType)
     }
 }

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

@@ -19,6 +19,7 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.symbols.*
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtNamedSymbol
+import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtTypeArgumentList

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/weighers/ExpectedTypeWeigher.kt

@@ -0,0 +1,51 @@
+/*
+ * Copyright 2010-2020 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.weighers
+
+import com.intellij.codeInsight.lookup.LookupElement
+import com.intellij.codeInsight.lookup.LookupElementWeigher
+import com.intellij.openapi.util.Key
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtTypedSymbol
+import org.jetbrains.kotlin.idea.frontend.api.types.KtType
+import org.jetbrains.kotlin.psi.UserDataProperty
+
+
+internal object ExpectedTypeWeigher {
+    object Weigher : LookupElementWeigher(WEIGHER_ID) {
+        override fun weigh(element: LookupElement): Comparable<*>? =
+            element.matchesExpectedType
+    }
+
+    fun KtAnalysisSession.addWeight(lookupElement: LookupElement, symbol: KtSymbol, expectedType: KtType?) {
+        lookupElement.matchesExpectedType = matchesExpectedType(symbol, expectedType)
+    }
+
+    private fun KtAnalysisSession.matchesExpectedType(
+        symbol: KtSymbol,
+        expectedType: KtType?
+    ) = when {
+        expectedType == null -> MatchesExpectedType.NON_TYPABLE
+        symbol !is KtTypedSymbol -> MatchesExpectedType.NON_TYPABLE
+        else -> MatchesExpectedType.matches(symbol.type isSubTypeOf expectedType)
+    }
+
+
+    private var LookupElement.matchesExpectedType by UserDataProperty(Key<MatchesExpectedType>("MATCHES_EXPECTED_TYPE"))
+
+    enum class MatchesExpectedType {
+        MATCHES, NON_TYPABLE, NOT_MATCHES, ;
+
+        companion object {
+            fun matches(matches: Boolean) = if (matches) MATCHES else NOT_MATCHES
+        }
+    }
+
+    const val WEIGHER_ID = "kotlin.expected.type"
+}
+
+

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/weighers/Weighers.kt

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2010-2020 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.weighers
+
+import com.intellij.codeInsight.completion.CompletionSorter
+import com.intellij.codeInsight.lookup.LookupElement
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
+import org.jetbrains.kotlin.idea.frontend.api.types.KtType
+
+internal object Weighers {
+    fun KtAnalysisSession.applyWeighsToLookupElement(lookupElement: LookupElement, symbol: KtSymbol, expectedType: KtType?) {
+        with(ExpectedTypeWeigher) { addWeight(lookupElement, symbol, expectedType) }
+    }
+
+    fun addWeighersToCompletionSorter(sorter: CompletionSorter): CompletionSorter =
+        sorter
+            .weighBefore(PlatformWeighersIds.STATS, ExpectedTypeWeigher.Weigher)
+
+    private object PlatformWeighersIds {
+        const val PREFIX = "prefix"
+        const val STATS = "stats"
+    }
+}

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/KtAnalysisSession.kt

@@ -62,9 +62,11 @@ abstract class KtAnalysisSession(final override val token: ValidityToken) : Vali
 
     fun KtDeclaration.getReturnKtType(): KtType = typeProvider.getReturnTypeForKtDeclaration(this)
 
-    fun KtType.isEqualTo(other: KtType): Boolean = typeProvider.isEqualTo(this, other)
+    infix fun KtType.isEqualTo(other: KtType): Boolean = typeProvider.isEqualTo(this, other)
 
-    fun KtType.isSubTypeOf(superType: KtType): Boolean = typeProvider.isSubTypeOf(this, superType)
+    infix fun KtType.isSubTypeOf(superType: KtType): Boolean = typeProvider.isSubTypeOf(this, superType)
+
+    fun KtExpression.getExpectedType(): KtType? = typeProvider.getExpectedType(this)
 
     fun KtType.isBuiltInFunctionalType(): Boolean = typeProvider.isBuiltinFunctionalType(this)
 

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtTypeProvider.kt

@@ -18,4 +18,6 @@ abstract class KtTypeProvider : KtAnalysisSessionComponent() {
 
     //TODO get rid of
     abstract fun isBuiltinFunctionalType(type: KtType): Boolean
+
+    abstract fun getExpectedType(expression: KtExpression): KtType?
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirTypeProvider.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.fir.components
 
+import com.intellij.psi.util.parentOfType
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.resolve.inference.isBuiltinFunctionalType
@@ -20,6 +21,7 @@ import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtReturnExpression
 import org.jetbrains.kotlin.types.AbstractTypeChecker
 import org.jetbrains.kotlin.types.AbstractTypeCheckerContext
 
@@ -36,6 +38,15 @@ internal class KtFirTypeProvider(
         expression.getOrBuildFirOfType<FirExpression>(firResolveState).typeRef.coneType.asKtType()
     }
 
+    override fun getExpectedType(expression: KtExpression): KtType? =
+        getExpectedTypeByReturnExpression(expression)
+
+    private fun getExpectedTypeByReturnExpression(expression: KtExpression): KtType? {
+        val returnParent = expression.parentOfType<KtReturnExpression>() ?: return null
+        val targetSymbol = with(analysisSession) { returnParent.getReturnTargetSymbol() } ?: return null
+        return targetSymbol.type
+    }
+
     override fun isEqualTo(first: KtType, second: KtType): Boolean = withValidityAssertion {
         second.assertIsValid()
         check(first is KtFirType)