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

idea/idea-completion/testData/basic/multifile/CallableReferenceNotImported/CallableReferenceNotImported.kt

@@ -1,4 +1,3 @@
-// FIR_COMPARISON
 fun String.foo() {
     val v = ::xxx_<caret>
 }

idea/idea-completion/testData/basic/multifile/DoNotCompleteWithConstraints/DoNotCompleteWithConstraints.kt

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

idea/idea-completion/testData/basic/multifile/ExtensionFunction/ExtensionFunction.dependency.kt

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

idea/idea-completion/testData/basic/multifile/ExtensionFunction/ExtensionFunction.kt

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

idea/idea-completion/testData/basic/multifile/ExtensionFunctionOnImportedFunction/ExtensionFunctionOnImportedFunction.kt

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

idea/idea-completion/testData/basic/multifile/NotImportedExtensionFunction2/NotImportedExtensionFunction2.kt

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

idea/idea-completion/testData/basic/multifile/NotImportedExtensionFunction3/NotImportedExtensionFunction3.kt

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

idea/idea-completion/testData/basic/multifile/NotImportedExtensionFunctionAndAlias/NotImportedExtensionFunctionAndAlias.kt

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

idea/idea-completion/testData/handlers/multifile/ExtensionFunctionImport-1.kt

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

idea/idea-completion/testData/handlers/multifile/ExtensionFunctionImport.kt.after

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

idea/idea-completion/testData/handlers/multifile/ExtensionPropertyImport-1.kt

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

idea/idea-completion/testData/handlers/multifile/ExtensionPropertyImport.kt.after

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

idea/idea-completion/testData/handlers/multifile/KeywordExtensionFunctionName-1.kt

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

idea/idea-completion/testData/handlers/multifile/KeywordExtensionFunctionName.kt.after

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

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

@@ -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
+    }
 }

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

@@ -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
+        val callableId = symbol.callableIdIfExists ?: return QuotedNamesAwareInsertionHandler(symbol.name)
-        if (callableId == null || !symbol.canBeCalledByFqName) 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

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/IndexHelper.kt

@@ -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('/', '.')

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/shortenReferencesUtils.kt

@@ -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,