FIR IDE: add completion with package receiver

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

@@ -16,6 +16,7 @@ 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.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassOrObjectSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtPackageSymbol
 import org.jetbrains.kotlin.idea.frontend.api.types.KtClassType
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.psi.KtExpression
@@ -88,6 +89,35 @@ internal class FirCallableCompletionContributor(
         expectedType: KtType?,
         extensionChecker: ExtensionApplicabilityChecker,
         visibilityChecker: CompletionVisibilityChecker,
+    ) {
+        val packageSymbol = explicitReceiver.reference()?.resolveToSymbol() as? KtPackageSymbol
+        if (packageSymbol == null) {
+            collectDotCompletionForCallableReceiver(implicitScopes, explicitReceiver, expectedType, extensionChecker, visibilityChecker)
+        } else {
+            collectDotCompletionForPackageReceiver(packageSymbol, expectedType, visibilityChecker)
+        }
+    }
+
+    private fun KtAnalysisSession.collectDotCompletionForPackageReceiver(
+        packageSymbol: KtPackageSymbol,
+        expectedType: KtType?,
+        visibilityChecker: CompletionVisibilityChecker
+    ) {
+        packageSymbol.getPackageScope()
+            .getCallableSymbols(scopeNameFilter)
+            .filterNot { it.isExtension }
+            .filter { with(visibilityChecker) { isVisible(it) } }
+            .forEach { callable ->
+                addSymbolToCompletion(expectedType, callable)
+            }
+    }
+
+    private fun KtAnalysisSession.collectDotCompletionForCallableReceiver(
+        implicitScopes: KtCompositeScope,
+        explicitReceiver: KtExpression,
+        expectedType: KtType?,
+        extensionChecker: ExtensionApplicabilityChecker,
+        visibilityChecker: CompletionVisibilityChecker,
     ) {
         val typeOfPossibleReceiver = explicitReceiver.getKtType()
         val possibleReceiverScope = typeOfPossibleReceiver.getTypeScope() ?: return

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

@@ -41,9 +41,13 @@ internal open class FirClassifierCompletionContributor(
         receiver: KtExpression,
         visibilityChecker: CompletionVisibilityChecker
     ) {
-        val reference = receiver.mainReference ?: return
-        val symbol = reference.resolveToSymbol() as? KtSymbolWithMembers ?: return
-        symbol.getMemberScope()
+        val reference = receiver.reference() ?: return
+        val scope = when (val symbol = reference.resolveToSymbol()) {
+            is KtSymbolWithMembers -> symbol.getMemberScope()
+            is KtPackageSymbol -> symbol.getPackageScope()
+            else -> return
+        }
+        scope
             .getClassifierSymbols(scopeNameFilter)
             .filter { filterClassifiers(it) }
             .filter { with(visibilityChecker) { isVisible(it) } }

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

@@ -20,8 +20,11 @@ import org.jetbrains.kotlin.idea.frontend.api.scopes.KtScopeNameFilter
 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.idea.references.mainReference
 import org.jetbrains.kotlin.idea.stubindex.KotlinSourceFilterScope
 import org.jetbrains.kotlin.platform.TargetPlatform
+import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
+import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtFile
 
 internal abstract class FirCompletionContributorBase<C : FirRawPositionCompletionContext>(protected val basicContext: FirBasicCompletionContext) {
@@ -61,6 +64,11 @@ internal abstract class FirCompletionContributorBase<C : FirRawPositionCompletio
         result.addElement(lookup)
     }
 
+    protected fun KtExpression.reference() = when (this) {
+        is KtDotQualifiedExpression -> selectorExpression?.mainReference
+        else -> mainReference
+    }
+
     private fun KtAnalysisSession.applyWeighers(
         lookupElement: LookupElement,
         symbol: KtSymbol,

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

@@ -34,9 +34,4 @@ internal class FirPackageCompletionContributor(
                 result.addElement(lookupElementFactory.createPackageLookupElement(packageSymbol.fqName))
             }
     }
-
-    private fun KtExpression.reference() = when (this) {
-        is KtDotQualifiedExpression -> selectorExpression?.mainReference
-        else -> mainReference
-    }
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/scopes/KtFirPackageScope.kt

@@ -48,19 +48,21 @@ internal class KtFirPackageScope(
 
     override fun getPossibleCallableNames() = withValidityAssertion {
         hashSetOf<Name>().apply {
-            KotlinTopLevelPropertyByPackageIndex.getInstance()[fqName.asString(), project, firScope.session.searchScope]
+            KotlinTopLevelPropertyByPackageIndex.getInstance()[fqName.asString(), project, searchScope]
                 .mapNotNullTo(this) { it.nameAsName }
-            KotlinTopLevelFunctionByPackageIndex.getInstance()[fqName.asString(), project, firScope.session.searchScope]
+            KotlinTopLevelFunctionByPackageIndex.getInstance()[fqName.asString(), project, searchScope]
                 .mapNotNullTo(this) { it.nameAsName }
         }
     }
 
     override fun getPossibleClassifierNames(): Set<Name> = withValidityAssertion {
-        KotlinTopLevelClassByPackageIndex.getInstance()[fqName.asString(), project, firScope.session.searchScope]
-            .mapNotNullTo(hashSetOf()) { it.nameAsName }
+        hashSetOf<Name>().apply {
+            KotlinTopLevelClassByPackageIndex.getInstance()[fqName.asString(), project, searchScope]
+                .mapNotNullTo(this) { it.nameAsName }
 
-        KotlinTopLevelTypeAliasByPackageIndex.getInstance()[fqName.asString(), project, firScope.session.searchScope]
-            .mapNotNullTo(hashSetOf()) { it.nameAsName }
+            KotlinTopLevelTypeAliasByPackageIndex.getInstance()[fqName.asString(), project, searchScope]
+                .mapNotNullTo(this) { it.nameAsName }
+        }
     }
 
     override fun getCallableSymbols(nameFilter: KtScopeNameFilter): Sequence<KtCallableSymbol> = withValidityAssertion {