Use index to get symbols from importing scopes

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirAbstractSimpleImportingScope.kt

@@ -21,7 +21,8 @@ abstract class FirAbstractSimpleImportingScope(
     scopeSession: ScopeSession
 ) : FirAbstractImportingScope(session, scopeSession, lookupInFir = true) {
 
-    protected abstract val simpleImports: Map<Name, List<FirResolvedImport>>
+    // TODO try to hide this
+    abstract val simpleImports: Map<Name, List<FirResolvedImport>>
 
     override fun processClassifiersByNameWithSubstitution(name: Name, processor: (FirClassifierSymbol<*>, ConeSubstitutor) -> Unit) {
         val imports = simpleImports[name] ?: return

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirAbstractStarImportingScope.kt

@@ -21,7 +21,8 @@ abstract class FirAbstractStarImportingScope(
     lookupInFir: Boolean = true
 ) : FirAbstractImportingScope(session, scopeSession, lookupInFir) {
 
-    protected abstract val starImports: List<FirResolvedImport>
+    // TODO try to hide this
+    abstract val starImports: List<FirResolvedImport>
 
     private val absentClassifierNames = mutableSetOf<Name>()
 

idea/idea-completion/testData/basic/common/primitiveCompletion/defaultKotlinImports.kt

@@ -0,0 +1,6 @@
+fun usage() {
+    val a = <caret>
+}
+
+// EXIST: listOf
+// EXIST: setOf

idea/idea-completion/testData/basic/common/primitiveCompletion/insideFunctionWithParameters.kt

@@ -0,0 +1,6 @@
+fun test(a: Int, b: Int) {
+    <caret>
+}
+
+// EXIST: a
+// EXIST: b

idea/idea-completion/testData/basic/common/primitiveCompletion/nullableCompletion.kt

@@ -0,0 +1,7 @@
+fun usage() {
+    val a = "10".toIntOrNull()
+
+    a?.<caret>
+}
+
+// EXIST: toString

idea/idea-completion/testData/basic/common/primitiveCompletion/stringCompletion.kt

@@ -0,0 +1,6 @@
+fun usage() {
+    "".<caret>
+}
+
+// EXIST: length
+// EXIST: get

idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/JSBasicCompletionTestGenerated.java

@@ -2252,6 +2252,11 @@ public class JSBasicCompletionTestGenerated extends AbstractJSBasicCompletionTes
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/classFieldsAndFunctionsFromInnerClassExplicitReceiver.kt");
             }
 
+            @TestMetadata("defaultKotlinImports.kt")
+            public void testDefaultKotlinImports() throws Exception {
+                runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/defaultKotlinImports.kt");
+            }
+
             @TestMetadata("explicitReceiverCompletion.kt")
             public void testExplicitReceiverCompletion() throws Exception {
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/explicitReceiverCompletion.kt");
@@ -2282,6 +2287,11 @@ public class JSBasicCompletionTestGenerated extends AbstractJSBasicCompletionTes
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/implicitReceiverFromExtensionMethod.kt");
             }
 
+            @TestMetadata("insideFunctionWithParameters.kt")
+            public void testInsideFunctionWithParameters() throws Exception {
+                runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/insideFunctionWithParameters.kt");
+            }
+
             @TestMetadata("localVariablesAndFunctions.kt")
             public void testLocalVariablesAndFunctions() throws Exception {
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/localVariablesAndFunctions.kt");
@@ -2292,11 +2302,21 @@ public class JSBasicCompletionTestGenerated extends AbstractJSBasicCompletionTes
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/localVariablesAndFunctionsFromNestedScope.kt");
             }
 
+            @TestMetadata("nullableCompletion.kt")
+            public void testNullableCompletion() throws Exception {
+                runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/nullableCompletion.kt");
+            }
+
             @TestMetadata("smartCast.kt")
             public void testSmartCast() throws Exception {
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/smartCast.kt");
             }
 
+            @TestMetadata("stringCompletion.kt")
+            public void testStringCompletion() throws Exception {
+                runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/stringCompletion.kt");
+            }
+
             @TestMetadata("topLevelVariablesAndFunctions.kt")
             public void testTopLevelVariablesAndFunctions() throws Exception {
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/topLevelVariablesAndFunctions.kt");

idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/JvmBasicCompletionTestGenerated.java

@@ -2252,6 +2252,11 @@ public class JvmBasicCompletionTestGenerated extends AbstractJvmBasicCompletionT
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/classFieldsAndFunctionsFromInnerClassExplicitReceiver.kt");
             }
 
+            @TestMetadata("defaultKotlinImports.kt")
+            public void testDefaultKotlinImports() throws Exception {
+                runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/defaultKotlinImports.kt");
+            }
+
             @TestMetadata("explicitReceiverCompletion.kt")
             public void testExplicitReceiverCompletion() throws Exception {
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/explicitReceiverCompletion.kt");
@@ -2282,6 +2287,11 @@ public class JvmBasicCompletionTestGenerated extends AbstractJvmBasicCompletionT
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/implicitReceiverFromExtensionMethod.kt");
             }
 
+            @TestMetadata("insideFunctionWithParameters.kt")
+            public void testInsideFunctionWithParameters() throws Exception {
+                runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/insideFunctionWithParameters.kt");
+            }
+
             @TestMetadata("localVariablesAndFunctions.kt")
             public void testLocalVariablesAndFunctions() throws Exception {
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/localVariablesAndFunctions.kt");
@@ -2292,11 +2302,21 @@ public class JvmBasicCompletionTestGenerated extends AbstractJvmBasicCompletionT
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/localVariablesAndFunctionsFromNestedScope.kt");
             }
 
+            @TestMetadata("nullableCompletion.kt")
+            public void testNullableCompletion() throws Exception {
+                runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/nullableCompletion.kt");
+            }
+
             @TestMetadata("smartCast.kt")
             public void testSmartCast() throws Exception {
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/smartCast.kt");
             }
 
+            @TestMetadata("stringCompletion.kt")
+            public void testStringCompletion() throws Exception {
+                runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/stringCompletion.kt");
+            }
+
             @TestMetadata("topLevelVariablesAndFunctions.kt")
             public void testTopLevelVariablesAndFunctions() throws Exception {
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/topLevelVariablesAndFunctions.kt");

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

@@ -7,14 +7,19 @@ package org.jetbrains.kotlin.idea.completion
 
 import com.intellij.codeInsight.completion.*
 import com.intellij.codeInsight.lookup.LookupElementBuilder
+import com.intellij.openapi.project.Project
 import com.intellij.patterns.PlatformPatterns
+import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.util.ProcessingContext
 import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirTypeParametersOwner
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
+import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.scope
 import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.impl.FirAbstractSimpleImportingScope
+import org.jetbrains.kotlin.fir.scopes.impl.FirAbstractStarImportingScope
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirConstructorSymbol
 import org.jetbrains.kotlin.fir.typeContext
@@ -24,6 +29,9 @@ import org.jetbrains.kotlin.idea.fir.getFirOfClosestParent
 import org.jetbrains.kotlin.idea.fir.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.fir.low.level.api.LowLevelFirApiFacade
 import org.jetbrains.kotlin.idea.references.mainReference
+import org.jetbrains.kotlin.idea.stubindex.KotlinTopLevelFunctionByPackageIndex
+import org.jetbrains.kotlin.idea.stubindex.KotlinTopLevelPropertyByPackageIndex
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
@@ -36,9 +44,6 @@ class KotlinFirCompletionContributor : CompletionContributor() {
     }
 }
 
-/**
- * 1. Collect all receivers for the expression (implicit and explicit)
- */
 private object KotlinFirCompletionProvider : CompletionProvider<CompletionParameters>() {
     override fun addCompletions(parameters: CompletionParameters, context: ProcessingContext, result: CompletionResultSet) {
         val originalFile = parameters.originalFile as? KtFile ?: return
@@ -56,6 +61,8 @@ private object KotlinFirCompletionProvider : CompletionProvider<CompletionParame
         val towerDataContext = completionContext.getTowerDataContext(nameExpression)
 
         val symbols: Sequence<FirCallableSymbol<*>> = sequence {
+            val packageIndexHelper = PackageIndexHelper(parameters.position.project)
+
             val explicitReceiver = element.explicitReceiver
             val explicitReceiverType = explicitReceiver?.typeRef?.coneTypeUnsafe<ConeKotlinType>()
             val explicitReceiverScope = explicitReceiverType?.scope(completionContext.session, ScopeSession())
@@ -72,16 +79,16 @@ private object KotlinFirCompletionProvider : CompletionProvider<CompletionParame
             val typeContext = completionContext.session.typeContext
 
             if (explicitReceiverScope != null) {
-                yieldAll(explicitReceiverScope.collectCallableSymbols())
+                yieldAll(explicitReceiverScope.collectCallableSymbols(packageIndexHelper))
 
                 val allApplicableExtensions = allNonExplicitScopes
-                    .flatMap { it.collectCallableSymbols() }
+                    .flatMap { it.collectCallableSymbols(packageIndexHelper) }
                     .filter { it.isExtension && it.isExtensionThatCanBeCalledOnTypes(listOf(explicitReceiverType), typeContext) }
 
                 yieldAll(allApplicableExtensions)
             } else {
                 val allAvailableSymbols = allNonExplicitScopes
-                    .flatMap { it.collectCallableSymbols() }
+                    .flatMap { it.collectCallableSymbols(packageIndexHelper) }
                     .filter { !it.isExtension || it.isExtensionThatCanBeCalledOnTypes(implicitReceiversTypes, typeContext) }
 
                 yieldAll(allAvailableSymbols)
@@ -100,13 +107,25 @@ private object KotlinFirCompletionProvider : CompletionProvider<CompletionParame
 
 private val FirCallableSymbol<*>.isExtension get() = fir.receiverTypeRef != null
 
-private fun FirScope.collectCallableSymbols(): MutableList<FirCallableSymbol<*>> {
+
+private fun FirScope.getScopeNames(indexHelper: PackageIndexHelper): Set<Name> = when (this) {
+    is FirAbstractStarImportingScope -> {
+        this.starImports.flatMapTo(mutableSetOf()) { indexHelper.getPackageTopLevelNames(it.packageFqName) }
+    }
+
+    is FirAbstractSimpleImportingScope -> simpleImports.keys
+
+    else -> getCallableNames()
+}
+
+private fun FirScope.collectCallableSymbols(packageIndexHelper: PackageIndexHelper): MutableList<FirCallableSymbol<*>> {
     val allCallableSymbols = mutableListOf<FirCallableSymbol<*>>()
     fun symbolsCollector(symbol: FirCallableSymbol<*>) {
         allCallableSymbols.add(symbol)
     }
 
-    for (name in getCallableNames()) {
+    val scopeNames = getScopeNames(packageIndexHelper)
+    for (name in scopeNames) {
         processFunctionsByName(name, ::symbolsCollector)
         processPropertiesByName(name, ::symbolsCollector)
     }
@@ -114,6 +133,22 @@ private fun FirScope.collectCallableSymbols(): MutableList<FirCallableSymbol<*>>
     return allCallableSymbols
 }
 
+private class PackageIndexHelper(private val project: Project) {
+    private val searchScope = GlobalSearchScope.allScope(project)
+
+    private val functionByPackageIndex = KotlinTopLevelFunctionByPackageIndex.getInstance()
+    private val propertyByPackageIndex = KotlinTopLevelPropertyByPackageIndex.getInstance()
+
+    fun getPackageTopLevelNames(packageFqName: FqName): Set<Name> {
+        return getTopLevelCallables(packageFqName).mapTo(mutableSetOf()) { it.nameAsSafeName }
+    }
+
+    private fun getTopLevelCallables(packageFqName: FqName): Sequence<KtCallableDeclaration> = sequence {
+        yieldAll(functionByPackageIndex.get(packageFqName.asString(), project, searchScope))
+        yieldAll(propertyByPackageIndex.get(packageFqName.asString(), project, searchScope))
+    }
+}
+
 private fun FirCallableSymbol<*>.isExtensionThatCanBeCalledOnTypes(
     receivers: List<ConeKotlinType>,
     session: TypeCheckerProviderContext