FIR IDE: Add extensions function completion from indices

The insertion handling test is not correct, we need to check file imports before adding new ones
2021-03-02 22:43:09 +03:00
parent 8396e17c52
commit 3427e0aebb
18 changed files with 125 additions and 13 deletions
@@ -1,4 +1,3 @@
-// FIR_COMPARISON
 fun String.foo() {
    val v = ::xxx_<caret>
 }
@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 package first

 interface TestedTrait() {
@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 package second

 import first.FirstClass
@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 package first

 class FirstClass() {
@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 package first

 // For KT-3102
@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 package first

 fun firstFun(p: () -> Unit) {
@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 package first

 fun firstFun(p: String.() -> Unit) {
@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 package first

 fun firstFun(x: third.Dependency) {
@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 package first

 fun firstFun() {
@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 package first

 import second.helloFun
@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 package first

 fun foo() {
@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 package first

 import second.extensionProp
@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 package first

 fun foo() {
@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 package first

 import second.`package`
@@ -16,14 +16,17 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.IndexHelper
 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
+import org.jetbrains.kotlin.idea.frontend.api.components.KtScopeContext
 import org.jetbrains.kotlin.idea.frontend.api.getAnalysisSessionFor
 import org.jetbrains.kotlin.idea.frontend.api.scopes.KtCompositeScope
 import org.jetbrains.kotlin.idea.frontend.api.scopes.KtScope
 import org.jetbrains.kotlin.idea.frontend.api.scopes.KtScopeNameFilter
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtCallableSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassOrObjectSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtNamedSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.isExtension
+import org.jetbrains.kotlin.idea.frontend.api.types.KtClassType
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.lexer.KtTokens
@@ -89,6 +92,7 @@ private class KotlinCommonCompletionProvider(
    private val indexHelper: IndexHelper
 ) {
    private val lookupElementFactory = KotlinFirLookupElementFactory()
+    private val typeNamesProvider = TypeNamesProvider(indexHelper)

    private val scopeNameFilter: KtScopeNameFilter =
        { name -> !name.isSpecial && prefixMatcher.prefixMatches(name.identifier) }
@@ -132,21 +136,21 @@ private class KotlinCommonCompletionProvider(
        with(getAnalysisSessionFor(originalFile).createContextDependentCopy(originalFile, nameExpression)) {
            val expectedType = nameExpression.getExpectedType()

-            val (implicitScopes, _) = originalFile.getScopeContextForPosition(nameExpression)
+            val scopesContext = originalFile.getScopeContextForPosition(nameExpression)

            fun KtCallableSymbol.hasSuitableExtensionReceiver(): Boolean =
                checkExtensionIsSuitable(originalFile, nameExpression, explicitReceiver)

            when {
-                nameExpression.parent is KtUserType -> collectTypesCompletion(result, implicitScopes, expectedType)
+                nameExpression.parent is KtUserType -> collectTypesCompletion(result, scopesContext.scopes, expectedType)
                explicitReceiver != null -> collectDotCompletion(
                    result,
-                    implicitScopes,
+                    scopesContext.scopes,
                    explicitReceiver,
                    expectedType,
                    KtCallableSymbol::hasSuitableExtensionReceiver
                )
-                else -> collectDefaultCompletion(result, implicitScopes, expectedType, KtCallableSymbol::hasSuitableExtensionReceiver)
+                else -> collectDefaultCompletion(result, scopesContext, expectedType, KtCallableSymbol::hasSuitableExtensionReceiver)
            }
        }
    }
@@ -183,14 +187,20 @@ private class KotlinCommonCompletionProvider(

        nonExtensionMembers.forEach { addSymbolToCompletion(result, expectedType, it) }
        extensionNonMembers.forEach { addSymbolToCompletion(result, expectedType, it) }
+
+        collectTopLevelExtensionsFromIndices(listOf(typeOfPossibleReceiver), hasSuitableExtensionReceiver)
+            .forEach { addSymbolToCompletion(result, expectedType, it) }
+
    }

    private fun KtAnalysisSession.collectDefaultCompletion(
        result: CompletionResultSet,
-        implicitScopes: KtCompositeScope,
+        implicitScopesContext: KtScopeContext,
        expectedType: KtType?,
        hasSuitableExtensionReceiver: KtCallableSymbol.() -> Boolean,
    ) {
+        val (implicitScopes, implicitReceiversTypes) = implicitScopesContext
+
        val availableNonExtensions = implicitScopes
            .getCallableSymbols(scopeNameFilter)
            .filterNot { it.isExtension }
@@ -209,7 +219,43 @@ private class KotlinCommonCompletionProvider(
                .forEach { addSymbolToCompletion(result, expectedType, it) }
        }

+        collectTopLevelExtensionsFromIndices(implicitReceiversTypes, hasSuitableExtensionReceiver)
+            .forEach { addSymbolToCompletion(result, expectedType, it) }

        collectTypesCompletion(result, implicitScopes, expectedType)
    }
+
+    private fun KtAnalysisSession.collectTopLevelExtensionsFromIndices(
+        receiverTypes: List<KtType>,
+        hasSuitableExtensionReceiver: KtCallableSymbol.() -> Boolean
+    ): Sequence<KtCallableSymbol> {
+        val implicitReceiverNames = findAllNamesOfTypes(receiverTypes)
+        val topLevelExtensions = indexHelper.getTopLevelExtensions(scopeNameFilter, implicitReceiverNames)
+
+        return topLevelExtensions.asSequence()
+            .map { it.getSymbol() as KtCallableSymbol }
+            .filter(hasSuitableExtensionReceiver)
+    }
+
+    private fun KtAnalysisSession.findAllNamesOfTypes(implicitReceiversTypes: List<KtType>) =
+        implicitReceiversTypes.flatMapTo(hashSetOf()) { with(typeNamesProvider) { findAllNames(it) } }
+}
+
+private class TypeNamesProvider(private val indexHelper: IndexHelper) {
+    fun KtAnalysisSession.findAllNames(type: KtType): Set<String> {
+        if (type !is KtClassType) return emptySet()
+
+        val result = hashSetOf<String>()
+        val typeName = type.classId.shortClassName.identifier
+
+        result += typeName
+        result += indexHelper.getPossibleTypeAliasExpansionNames(typeName)
+
+        val superTypes = (type.classSymbol as? KtClassOrObjectSymbol)?.superTypes
+        superTypes?.forEach { superType ->
+            result += findAllNames(superType.type)
+        }
+
+        return result
+    }
 }
@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.idea.core.withRootPrefixIfNeeded
 import org.jetbrains.kotlin.idea.frontend.api.HackToForceAllowRunningAnalyzeOnEDT
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.analyze
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.addImportToFile
 import org.jetbrains.kotlin.idea.frontend.api.hackyAllowRunningOnEdt
 import org.jetbrains.kotlin.idea.frontend.api.symbols.*
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtNamedSymbol
@@ -115,10 +116,13 @@ private class VariableLookupElementFactory {
        if (setterName != null) "$getterName()/$setterName()" else "$getterName()"

    private fun createInsertHandler(symbol: KtVariableLikeSymbol): InsertHandler<LookupElement> {
-        val callableId = symbol.callableIdIfExists
-        if (callableId == null || !symbol.canBeCalledByFqName) return QuotedNamesAwareInsertionHandler(symbol.name)
+        val callableId = symbol.callableIdIfExists ?: return QuotedNamesAwareInsertionHandler(symbol.name)

-        return ShorteningVariableInsertionHandler(callableId)
+        return if (symbol.canBeCalledByFqName) {
+            ShorteningVariableInsertionHandler(callableId)
+        } else {
+            SimpleVariableInsertionHandler(callableId)
+        }
    }

    private val KtVariableLikeSymbol.callableIdIfExists: FqName?
@@ -181,7 +185,6 @@ private class FunctionLookupElementFactory {

    private fun KtAnalysisSession.createInsertHandler(symbol: KtFunctionSymbol): InsertHandler<LookupElement> {
        val functionFqName = symbol.callableIdIfNonLocal
-
        return if (functionFqName != null && canBeCalledByFqName(symbol)) {
            ShorteningFunctionInsertionHandler(
                functionFqName,
@@ -192,11 +195,15 @@ private class FunctionLookupElementFactory {
            SimpleFunctionInsertionHandler(
                symbol.name,
                inputValueArguments = symbol.valueParameters.isNotEmpty(),
-                insertEmptyLambda = insertLambdaBraces(symbol)
+                insertEmptyLambda = insertLambdaBraces(symbol),
+                nameToImport = symbol.importableFqName,
            )
        }
    }

+    private val KtFunctionSymbol.importableFqName: FqName?
+        get() = if (dispatchType == null) callableIdIfNonLocal else null
+
    private fun canBeCalledByFqName(symbol: KtFunctionSymbol): Boolean {
        return !symbol.isExtension && symbol.dispatchType == null
    }
@@ -299,15 +306,22 @@ private abstract class AbstractFunctionInsertionHandler(
 private class SimpleFunctionInsertionHandler(
    name: Name,
    inputValueArguments: Boolean,
-    insertEmptyLambda: Boolean
+    insertEmptyLambda: Boolean,
+    private val nameToImport: FqName?
 ) : AbstractFunctionInsertionHandler(name, inputValueArguments, insertEmptyLambda) {
    override fun handleInsert(context: InsertionContext, item: LookupElement) {
        super.handleInsert(context, item)

+        val targetFile = context.file as? KtFile ?: return
+
        val startOffset = context.startOffset
        val element = context.file.findElementAt(startOffset) ?: return

        addArguments(context, element)
+
+        if (nameToImport != null && targetFile.importDirectives.none { it.importPath?.fqName == nameToImport }) {
+            addImportToFile(context.project, targetFile, nameToImport)
+        }
    }
 }

@@ -352,6 +366,19 @@ private class ShorteningVariableInsertionHandler(private val name: FqName) : Ins
    }
 }

+private class SimpleVariableInsertionHandler(private val nameToImport: FqName) :
+    QuotedNamesAwareInsertionHandler(nameToImport.shortName()) {
+    override fun handleInsert(context: InsertionContext, item: LookupElement) {
+        super.handleInsert(context, item)
+
+        val targetFile = context.file as? KtFile ?: return
+
+        if (targetFile.importDirectives.none { it.importPath?.fqName == nameToImport }) {
+            addImportToFile(context.project, targetFile, nameToImport)
+        }
+    }
+}
+
 private open class QuotedNamesAwareInsertionHandler(private val name: Name) : InsertHandler<LookupElement> {
    override fun handleInsert(context: InsertionContext, item: LookupElement) {
        val startOffset = context.startOffset
@@ -98,6 +98,33 @@ public class IndexHelper(val project: Project, private val scope: GlobalSearchSc
        return (functions + properties).toList()
    }

+    fun getTopLevelExtensions(nameFilter: (Name) -> Boolean, receiverTypeNames: Set<String>): Collection<KtCallableDeclaration> {
+        val index = KotlinTopLevelExtensionsByReceiverTypeIndex.INSTANCE
+
+        return index.getAllKeys(project).asSequence()
+            .onEach { ProgressManager.checkCanceled() }
+            .filter { KotlinTopLevelExtensionsByReceiverTypeIndex.receiverTypeNameFromKey(it) in receiverTypeNames }
+            .filter { nameFilter(Name.identifier(KotlinTopLevelExtensionsByReceiverTypeIndex.callableNameFromKey(it))) }
+            .flatMap { key -> index[key, project, scope] }
+            .toList()
+    }
+
+    fun getPossibleTypeAliasExpansionNames(originalTypeName: String): Set<String> {
+        val index = KotlinTypeAliasByExpansionShortNameIndex.INSTANCE
+        val out = mutableSetOf<String>()
+
+        fun searchRecursively(typeName: String) {
+            ProgressManager.checkCanceled()
+            index[typeName, project, scope].asSequence()
+                .mapNotNull { it.name }
+                .filter { out.add(it) }
+                .forEach(::searchRecursively)
+        }
+
+        searchRecursively(originalTypeName)
+        return out
+    }
+
    companion object {
        private fun CallableId.asStringForIndexes(): String =
            (if (packageName.isRoot) callableName.asString() else toString()).replace('/', '.')
@@ -21,7 +21,7 @@ private val SimpleImportPathComparator: Comparator<ImportPath> = compareBy(Impor
 *
 * We want it as a copy because we do not yet care about imports ordering, so we do not need a fancy comparator.
 */
-internal fun addImportToFile(
+fun addImportToFile(
    project: Project,
    file: KtFile,
    fqName: FqName,