FIR IDE: introduce import statement completion

idea/idea-analysis/src/org/jetbrains/kotlin/idea/stubindex/PackageIndexUtil.kt

@@ -8,14 +8,40 @@ package org.jetbrains.kotlin.idea.stubindex
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.newvfs.persistent.PersistentFS
 import com.intellij.openapi.vfs.newvfs.persistent.PersistentFSImpl
+import com.intellij.psi.JavaPsiFacade
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.stubs.StubIndex
+import org.jetbrains.kotlin.idea.search.getKotlinFqName
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.platform.TargetPlatform
+import org.jetbrains.kotlin.platform.jvm.isJvm
 import org.jetbrains.kotlin.psi.KtFile
 
 
 object PackageIndexUtil {
+    fun getJavaAndKotlinSubPackageFqNames(
+        packageFqName: FqName,
+        scope: GlobalSearchScope,
+        project: Project,
+        targetPlatform: TargetPlatform,
+        nameFilter: (Name) -> Boolean
+    ) = sequence {
+        yieldAll(getSubPackageFqNames(packageFqName, scope, project, nameFilter))
+
+        if (targetPlatform.isJvm()) {
+            val javaPackage = JavaPsiFacade.getInstance(project).findPackage(packageFqName.asString())
+            if (javaPackage != null) {
+                for (psiPackage in javaPackage.getSubPackages(scope)) {
+                    val fqName = psiPackage.getKotlinFqName() ?: continue
+                    if (nameFilter(fqName.shortName())) {
+                        yield(fqName)
+                    }
+                }
+            }
+        }
+    }
+
     @JvmStatic
     fun getSubPackageFqNames(
         packageFqName: FqName,

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

@@ -10,8 +10,6 @@ import com.intellij.openapi.components.service
 import com.intellij.patterns.PlatformPatterns
 import com.intellij.patterns.PsiJavaPatterns
 import com.intellij.util.ProcessingContext
-import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
-import org.jetbrains.kotlin.idea.completion.checkers.CompletionVisibilityChecker
 import org.jetbrains.kotlin.idea.completion.context.*
 import org.jetbrains.kotlin.idea.completion.context.FirBasicCompletionContext
 import org.jetbrains.kotlin.idea.completion.context.FirExpressionNameReferencePositionContext
@@ -25,7 +23,6 @@ import org.jetbrains.kotlin.idea.completion.contributors.FirClassifierCompletion
 import org.jetbrains.kotlin.idea.completion.contributors.FirKeywordCompletionContributor
 import org.jetbrains.kotlin.idea.completion.contributors.complete
 import org.jetbrains.kotlin.idea.completion.weighers.Weighers
-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.KtAnalysisSession
 import org.jetbrains.kotlin.lexer.KtTokens
@@ -70,6 +67,7 @@ private object KotlinFirCompletionProvider : CompletionProvider<CompletionParame
         val classifierContributor = FirClassifierCompletionContributor(basicContext)
         val annotationsContributor = FirAnnotationCompletionContributor(basicContext)
         val packageCompletionContributor = FirPackageCompletionContributor(basicContext)
+        val importDirectivePackageMembersCompletionContributor = FirImportDirectivePackageMembersCompletionContributor(basicContext)
 
         when (positionContext) {
             is FirExpressionNameReferencePositionContext -> {
@@ -95,6 +93,11 @@ private object KotlinFirCompletionProvider : CompletionProvider<CompletionParame
                 complete(FirSuperEntryContributor(basicContext), positionContext)
             }
 
+            is FirImportDirectivePositionContext -> {
+                complete(packageCompletionContributor, positionContext)
+                complete(importDirectivePackageMembersCompletionContributor, positionContext)
+            }
+
             is FirUnknownPositionContext -> {
                 complete(keywordContributor, positionContext)
             }

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

@@ -13,7 +13,7 @@ import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.*
 import javax.swing.Icon
 
 internal object KotlinFirIconProvider {
-    fun getIconFor(symbol: KtNamedSymbol): Icon? {
+    fun getIconFor(symbol: KtSymbol): Icon? {
         if (symbol is KtFunctionSymbol) {
             val isAbstract = symbol.modality == Modality.ABSTRACT
 

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

@@ -11,10 +11,12 @@ import com.intellij.codeInsight.completion.InsertionContext
 import com.intellij.codeInsight.lookup.Lookup
 import com.intellij.codeInsight.lookup.LookupElement
 import com.intellij.codeInsight.lookup.LookupElementBuilder
+import com.intellij.icons.AllIcons
 import com.intellij.openapi.diagnostic.ControlFlowException
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.idea.completion.contributors.helpers.addDotAndInvokeCompletion
 import org.jetbrains.kotlin.idea.completion.handlers.isTextAt
 import org.jetbrains.kotlin.idea.core.asFqNameWithRootPrefixIfNeeded
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
@@ -27,6 +29,7 @@ import org.jetbrains.kotlin.idea.frontend.api.tokens.HackToForceAllowRunningAnal
 import org.jetbrains.kotlin.idea.frontend.api.tokens.hackyAllowRunningOnEdt
 import org.jetbrains.kotlin.idea.frontend.api.types.KtFunctionalType
 import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.miniStdLib.letIf
 import org.jetbrains.kotlin.name.*
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtTypeArgumentList
@@ -40,38 +43,50 @@ internal class KotlinFirLookupElementFactory {
     private val typeParameterLookupElementFactory = TypeParameterLookupElementFactory()
 
     fun KtAnalysisSession.createLookupElement(symbol: KtNamedSymbol): LookupElement? {
-        val elementBuilder = when (symbol) {
-            is KtFunctionSymbol -> with(functionLookupElementFactory) { createLookup(symbol) }
-            is KtVariableLikeSymbol -> with(variableLookupElementFactory) { createLookup(symbol) }
-            is KtClassLikeSymbol -> classLookupElementFactory.createLookup(symbol)
-            is KtTypeParameterSymbol -> typeParameterLookupElementFactory.createLookup(symbol)
+        return when (symbol) {
+            is KtCallableSymbol -> createCallableLookupElement(symbol, importingStrategy = detectImportStrategy(symbol), CallableInsertionStrategy.AS_CALL)
+            is KtClassLikeSymbol -> with(classLookupElementFactory) { createLookup(symbol) }
+            is KtTypeParameterSymbol -> with(typeParameterLookupElementFactory) { createLookup(symbol) }
             else -> throw IllegalArgumentException("Cannot create a lookup element for $symbol")
-        } ?: return null
-
-        return elementBuilder
-            .withPsiElement(symbol.psi) // TODO check if it is a heavy operation and should be postponed
-            .withIcon(KotlinFirIconProvider.getIconFor(symbol))
+        }
     }
 
-    fun createPackageLookupElement(packageFqName: FqName): LookupElement {
-        return LookupElementBuilder.create(PackageLookupObject(packageFqName), packageFqName.shortName().asString())
-            .withInsertHandler(QuotedNamesAwareInsertionHandler())
-            .let {
-                if (!packageFqName.parent().isRoot) {
-                    it.appendTailText(" (${packageFqName.asString()})", true)
-                } else it
+    fun KtAnalysisSession.createCallableLookupElement(
+        symbol: KtCallableSymbol,
+        importingStrategy: CallableImportStrategy,
+        insertionStrategy: CallableInsertionStrategy,
+    ): LookupElementBuilder? {
+        return when (symbol) {
+            is KtFunctionSymbol -> with(functionLookupElementFactory) { createLookup(symbol, importingStrategy, insertionStrategy) }
+            is KtVariableLikeSymbol -> with(variableLookupElementFactory) { createLookup(symbol, importingStrategy) }
+            else -> throw IllegalArgumentException("Cannot create a lookup element for $symbol")
+        }
+    }
+
+    fun createPackagePartLookupElement(packagePartFqName: FqName): LookupElement {
+        val shortName = packagePartFqName.shortName()
+        return LookupElementBuilder.create(PackagePartLookupObject(shortName), "${shortName.render()}.")
+            .withInsertHandler(PackagePartInsertionHandler)
+            .withIcon(AllIcons.Nodes.Package)
+            .letIf(!packagePartFqName.parent().isRoot) {
+                it.appendTailText("(${packagePartFqName.asString()})", true)
             }
     }
 
     fun KtAnalysisSession.createLookupElementForClassLikeSymbol(symbol: KtClassLikeSymbol, insertFqName: Boolean = true): LookupElement? {
         if (symbol !is KtNamedSymbol) return null
-        return classLookupElementFactory.createLookup(symbol, insertFqName)
-            .withPsiElement(symbol.psi) // TODO check if it is a heavy operation and should be postponed
-            .withIcon(KotlinFirIconProvider.getIconFor(symbol))
+        return with(classLookupElementFactory) { createLookup(symbol, insertFqName) }
     }
 }
 
-private sealed class CallableImportStrategy {
+private fun KtAnalysisSession.withSymbolInfo(
+    symbol: KtSymbol,
+    elementBuilder: LookupElementBuilder
+): LookupElementBuilder = elementBuilder
+    .withPsiElement(symbol.psi) // TODO check if it is a heavy operation and should be postponed
+    .withIcon(KotlinFirIconProvider.getIconFor(symbol))
+
+internal sealed class CallableImportStrategy {
     object DoNothing : CallableImportStrategy()
     data class AddImport(val nameToImport: CallableId) : CallableImportStrategy()
     data class InsertFqNameAndShorten(val callableId: CallableId) : CallableImportStrategy()
@@ -86,14 +101,12 @@ private interface KotlinLookupObject {
     val shortName: Name
 }
 
-private data class PackageLookupObject(
-    val packageFqName: FqName,
-) : KotlinLookupObject {
-    override val shortName: Name = packageFqName.shortName()
-}
+private data class PackagePartLookupObject(
+    override val shortName: Name,
+) : KotlinLookupObject
 
-
-private data class ClassifierLookupObject(override val shortName: Name, val classId: ClassId?, val insertFqName: Boolean) : KotlinLookupObject
+private data class ClassifierLookupObject(override val shortName: Name, val classId: ClassId?, val insertFqName: Boolean) :
+    KotlinLookupObject
 
 /**
  * Simplest lookup object so two lookup elements for the same function will clash.
@@ -116,30 +129,37 @@ private data class VariableLookupObject(
 ) : KotlinLookupObject
 
 class ClassLookupElementFactory {
-    fun createLookup(symbol: KtClassLikeSymbol, insertFqName: Boolean = true): LookupElementBuilder {
+    fun KtAnalysisSession.createLookup(symbol: KtClassLikeSymbol, insertFqName: Boolean = true): LookupElementBuilder {
         val name = symbol.nameOrAnonymous
         return LookupElementBuilder.create(ClassifierLookupObject(name, symbol.classIdIfNonLocal, insertFqName), name.asString())
             .withInsertHandler(ClassifierInsertionHandler)
+            .let { withSymbolInfo(symbol, it) }
     }
 }
 
 private class TypeParameterLookupElementFactory {
-    fun createLookup(symbol: KtTypeParameterSymbol): LookupElementBuilder {
-        return LookupElementBuilder.create(UniqueLookupObject(), symbol.name.asString())
+    fun KtAnalysisSession.createLookup(symbol: KtTypeParameterSymbol): LookupElementBuilder {
+        return LookupElementBuilder
+            .create(UniqueLookupObject(), symbol.name.asString())
+            .let { withSymbolInfo(symbol, it) }
     }
 }
 
 private class VariableLookupElementFactory {
-    fun KtAnalysisSession.createLookup(symbol: KtVariableLikeSymbol): LookupElementBuilder {
+    fun KtAnalysisSession.createLookup(
+        symbol: KtVariableLikeSymbol,
+        importStrategy: CallableImportStrategy = detectImportStrategy(symbol)
+    ): LookupElementBuilder {
         val lookupObject = VariableLookupObject(
             symbol.name,
-            importStrategy = detectImportStrategy(symbol)
+            importStrategy = importStrategy
         )
 
         return LookupElementBuilder.create(lookupObject, symbol.name.asString())
             .withTypeText(symbol.annotatedType.type.render(WITH_TYPE_RENDERING_OPTIONS))
             .markIfSyntheticJavaProperty(symbol)
             .withInsertHandler(VariableInsertionHandler)
+            .let { withSymbolInfo(symbol, it) }
     }
 
     private fun LookupElementBuilder.markIfSyntheticJavaProperty(symbol: KtVariableLikeSymbol): LookupElementBuilder = when (symbol) {
@@ -154,35 +174,50 @@ private class VariableLookupElementFactory {
 
     private fun buildSyntheticPropertyTailText(getterName: String, setterName: String?): String =
         if (setterName != null) "$getterName()/$setterName()" else "$getterName()"
+}
 
-    private fun detectImportStrategy(symbol: KtVariableLikeSymbol): CallableImportStrategy {
-        if (symbol !is KtKotlinPropertySymbol || symbol.dispatchType != null) return CallableImportStrategy.DoNothing
+private fun detectImportStrategy(symbol: KtCallableSymbol): CallableImportStrategy {
+    if (symbol !is KtKotlinPropertySymbol || symbol.dispatchType != null) return CallableImportStrategy.DoNothing
 
-        val propertyId = symbol.callableIdIfNonLocal ?: return CallableImportStrategy.DoNothing
+    val propertyId = symbol.callableIdIfNonLocal ?: return CallableImportStrategy.DoNothing
 
-        return if (symbol.isExtension) {
-            CallableImportStrategy.AddImport(propertyId)
-        } else {
-            CallableImportStrategy.InsertFqNameAndShorten(propertyId)
-        }
+    return if (symbol.isExtension) {
+        CallableImportStrategy.AddImport(propertyId)
+    } else {
+        CallableImportStrategy.InsertFqNameAndShorten(propertyId)
     }
 }
 
+internal enum class CallableInsertionStrategy {
+    AS_CALL,
+    AS_IDENTIFIER
+}
+
 private class FunctionLookupElementFactory {
-    fun KtAnalysisSession.createLookup(symbol: KtFunctionSymbol): LookupElementBuilder? {
+    fun KtAnalysisSession.createLookup(
+        symbol: KtFunctionSymbol,
+        importStrategy: CallableImportStrategy,
+        insertionStrategy: CallableInsertionStrategy
+    ): LookupElementBuilder? {
         val lookupObject = FunctionLookupObject(
             symbol.name,
-            importStrategy = detectImportStrategy(symbol),
+            importStrategy = importStrategy,
             inputValueArguments = symbol.valueParameters.isNotEmpty(),
             insertEmptyLambda = insertLambdaBraces(symbol),
             renderedFunctionParameters = with(ShortNamesRenderer) { renderFunctionParameters(symbol) }
         )
 
+        val insertionHandler = when (insertionStrategy) {
+            CallableInsertionStrategy.AS_CALL -> FunctionInsertionHandler
+            CallableInsertionStrategy.AS_IDENTIFIER -> QuotedNamesAwareInsertionHandler()
+        }
+
         return try {
             LookupElementBuilder.create(lookupObject, symbol.name.asString())
                 .withTailText(getTailText(symbol), true)
                 .withTypeText(symbol.annotatedType.type.render(WITH_TYPE_RENDERING_OPTIONS))
-                .withInsertHandler(FunctionInsertionHandler)
+                .withInsertHandler(insertionHandler)
+                .let { withSymbolInfo(symbol, it) }
         } catch (e: Throwable) {
             if (e is ControlFlowException) throw e
             LOG.error(e)
@@ -190,18 +225,6 @@ private class FunctionLookupElementFactory {
         }
     }
 
-    private fun detectImportStrategy(symbol: KtFunctionSymbol): CallableImportStrategy {
-        if (symbol.dispatchType != null) return CallableImportStrategy.DoNothing
-
-        val functionFqName = symbol.callableIdIfNonLocal ?: return CallableImportStrategy.DoNothing
-
-        return if (symbol.isExtension) {
-            CallableImportStrategy.AddImport(functionFqName)
-        } else {
-            CallableImportStrategy.InsertFqNameAndShorten(functionFqName)
-        }
-    }
-
     private fun KtAnalysisSession.getTailText(symbol: KtFunctionSymbol): String {
         return if (insertLambdaBraces(symbol)) " {...}" else with(ShortNamesRenderer) { renderFunctionParameters(symbol) }
     }
@@ -371,11 +394,22 @@ private object VariableInsertionHandler : InsertHandler<LookupElement> {
                 shortenReferencesForFirCompletion(targetFile, TextRange(context.startOffset, context.tailOffset))
             }
 
-            is CallableImportStrategy.DoNothing -> {}
+            is CallableImportStrategy.DoNothing -> {
+            }
         }
     }
 }
 
+private object PackagePartInsertionHandler : InsertHandler<LookupElement> {
+    override fun handleInsert(context: InsertionContext, item: LookupElement) {
+        val lookupElement = item.`object` as PackagePartLookupObject
+        val name = lookupElement.shortName.render()
+        context.document.replaceString(context.startOffset, context.tailOffset, name)
+        context.commitDocument()
+        context.addDotAndInvokeCompletion()
+    }
+}
+
 private open class QuotedNamesAwareInsertionHandler : InsertHandler<LookupElement> {
     override fun handleInsert(context: InsertionContext, item: LookupElement) {
         val lookupElement = item.`object` as KotlinLookupObject

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/context/FirBasicCompletionContext.kt

@@ -9,6 +9,7 @@ import com.intellij.codeInsight.completion.CompletionParameters
 import com.intellij.codeInsight.completion.CompletionResultSet
 import com.intellij.codeInsight.completion.PrefixMatcher
 import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.completion.KotlinFirLookupElementFactory
 import org.jetbrains.kotlin.idea.fir.low.level.api.IndexHelper
@@ -29,6 +30,7 @@ internal class FirBasicCompletionContext(
     val lookupElementFactory: KotlinFirLookupElementFactory = KotlinFirLookupElementFactory(),
 ) {
     val visibleScope = KotlinSourceFilterScope.projectSourceAndClassFiles(originalKtFile.resolveScope, project)
+    val moduleInfo: IdeaModuleInfo = originalKtFile.getModuleInfo()
 
     companion object {
         fun createFromParameters(parameters: CompletionParameters, result: CompletionResultSet): FirBasicCompletionContext? {

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/context/FirPositionCompletionContext.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.idea.completion.context
 
 import com.intellij.psi.PsiElement
+import com.intellij.psi.util.parentOfType
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.analyse
 import org.jetbrains.kotlin.idea.frontend.api.analyseInDependedAnalysisSession
@@ -24,6 +25,15 @@ internal sealed class FirNameReferencePositionContext : FirRawPositionCompletion
     abstract val explicitReceiver: KtExpression?
 }
 
+internal class FirImportDirectivePositionContext(
+    override val position: PsiElement,
+    override val reference: KtSimpleNameReference,
+    override val nameExpression: KtSimpleNameExpression,
+    override val explicitReceiver: KtExpression?,
+    val importDirective: KtImportDirective,
+) : FirNameReferencePositionContext()
+
+
 internal class FirTypeNameReferencePositionContext(
     override val position: PsiElement,
     override val reference: KtSimpleNameReference,
@@ -68,17 +78,36 @@ internal object FirPositionCompletionContextDetector {
         val nameExpression = reference.expression.takeIf { it !is KtLabelReferenceExpression }
             ?: return FirUnknownPositionContext(position)
         val explicitReceiver = nameExpression.getReceiverExpression()
+        val parent = nameExpression.parent
 
-        return when (val parent = nameExpression.parent) {
-            is KtUserType -> {
+        return when {
+            parent is KtUserType -> {
                 detectForTypeContext(parent, position, reference, nameExpression, explicitReceiver)
             }
+            nameExpression.isReferenceExpressionInImportDirective() -> {
+                FirImportDirectivePositionContext(
+                    position,
+                    reference,
+                    nameExpression,
+                    explicitReceiver,
+                    parent.parentOfType(withSelf = true)!!
+                )
+            }
             else -> {
                 FirExpressionNameReferencePositionContext(position, reference, nameExpression, explicitReceiver)
             }
         }
     }
 
+    private fun KtExpression.isReferenceExpressionInImportDirective() = when (val parent = parent) {
+        is KtImportDirective -> parent.importedReference == this
+        is KtDotQualifiedExpression -> {
+            val importDirective = parent.parent as? KtImportDirective
+            importDirective?.importedReference == parent
+        }
+        else -> false
+    }
+
     private fun detectForTypeContext(
         userType: KtUserType,
         position: PsiElement,
@@ -116,7 +145,7 @@ internal object FirPositionCompletionContextDetector {
                 positionContext.nameExpression,
                 action
             )
-            is FirUnknownPositionContext -> {
+            is FirUnknownPositionContext, is FirImportDirectivePositionContext -> {
                 analyse(basicContext.originalKtFile, action)
             }
         }

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

@@ -8,13 +8,10 @@ package org.jetbrains.kotlin.idea.completion.contributors
 import org.jetbrains.kotlin.idea.completion.checkers.CompletionVisibilityChecker
 import org.jetbrains.kotlin.idea.completion.context.FirBasicCompletionContext
 import org.jetbrains.kotlin.idea.completion.context.FirNameReferenceRawPositionContext
+import org.jetbrains.kotlin.idea.completion.contributors.helpers.getStaticScope
 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.symbols.markers.KtSymbolWithMembers
 import org.jetbrains.kotlin.idea.frontend.api.types.KtClassType
-import org.jetbrains.kotlin.idea.references.mainReference
-import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtEnumEntry
 import org.jetbrains.kotlin.psi.KtExpression
 
@@ -42,11 +39,7 @@ internal open class FirClassifierCompletionContributor(
         visibilityChecker: CompletionVisibilityChecker
     ) {
         val reference = receiver.reference() ?: return
-        val scope = when (val symbol = reference.resolveToSymbol()) {
-            is KtSymbolWithMembers -> symbol.getMemberScope()
-            is KtPackageSymbol -> symbol.getPackageScope()
-            else -> return
-        }
+        val scope = getStaticScope(reference) ?: return
         scope
             .getClassifierSymbols(scopeNameFilter)
             .filter { filterClassifiers(it) }

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

@@ -10,6 +10,8 @@ import com.intellij.codeInsight.completion.CompletionResultSet
 import com.intellij.codeInsight.completion.PrefixMatcher
 import com.intellij.codeInsight.lookup.LookupElement
 import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.idea.completion.CallableImportStrategy
+import org.jetbrains.kotlin.idea.completion.CallableInsertionStrategy
 import org.jetbrains.kotlin.idea.completion.KotlinFirLookupElementFactory
 import org.jetbrains.kotlin.idea.completion.context.FirBasicCompletionContext
 import org.jetbrains.kotlin.idea.completion.context.FirRawPositionCompletionContext
@@ -21,7 +23,6 @@ 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
@@ -64,6 +65,18 @@ internal abstract class FirCompletionContributorBase<C : FirRawPositionCompletio
         result.addElement(lookup)
     }
 
+    protected fun KtAnalysisSession.addCallableSymbolToCompletion(
+        symbol: KtCallableSymbol,
+        importingStrategy: CallableImportStrategy,
+        insertionStrategy: CallableInsertionStrategy,
+    ) {
+        if (symbol !is KtNamedSymbol) return
+        val lookup = with(lookupElementFactory) {
+            createCallableLookupElement(symbol, importingStrategy, insertionStrategy)
+        } ?: return
+        result.addElement(lookup)
+    }
+
     protected fun KtExpression.reference() = when (this) {
         is KtDotQualifiedExpression -> selectorExpression?.mainReference
         else -> mainReference

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

@@ -0,0 +1,38 @@
+/*
+ * Copyright 2010-2021 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.contributors
+
+import org.jetbrains.kotlin.idea.completion.CallableImportStrategy
+import org.jetbrains.kotlin.idea.completion.CallableInsertionStrategy
+import org.jetbrains.kotlin.idea.completion.checkers.CompletionVisibilityChecker
+import org.jetbrains.kotlin.idea.completion.context.FirBasicCompletionContext
+import org.jetbrains.kotlin.idea.completion.context.FirImportDirectivePositionContext
+import org.jetbrains.kotlin.idea.completion.contributors.helpers.getStaticScope
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+
+internal class FirImportDirectivePackageMembersCompletionContributor(
+    basicContext: FirBasicCompletionContext
+) : FirCompletionContributorBase<FirImportDirectivePositionContext>(basicContext) {
+    override fun KtAnalysisSession.complete(positionContext: FirImportDirectivePositionContext) {
+        val reference = positionContext.explicitReceiver?.reference() ?: return
+        val scope = getStaticScope(reference) ?: return
+        val visibilityChecker = CompletionVisibilityChecker.create(basicContext, positionContext)
+
+        scope.getClassifierSymbols(scopeNameFilter)
+            .filter { with(visibilityChecker) { isVisible(it) } }
+            .forEach { addClassifierSymbolToCompletion(it, insertFqName = false) }
+
+        scope.getCallableSymbols(scopeNameFilter)
+            .filter { with(visibilityChecker) { isVisible(it) } }
+            .forEach {
+                addCallableSymbolToCompletion(
+                    it,
+                    importingStrategy = CallableImportStrategy.DoNothing,
+                    insertionStrategy = CallableInsertionStrategy.AS_IDENTIFIER
+                )
+            }
+    }
+}

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

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.idea.completion.contributors
 
-import com.intellij.psi.JavaPsiFacade
 import org.jetbrains.kotlin.idea.completion.context.FirBasicCompletionContext
 import org.jetbrains.kotlin.idea.completion.context.FirNameReferenceRawPositionContext
 import org.jetbrains.kotlin.idea.completion.context.FirRawPositionCompletionContext
@@ -13,12 +12,6 @@ import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtPackageSymbol
 import org.jetbrains.kotlin.idea.isExcludedFromAutoImport
 import org.jetbrains.kotlin.idea.project.languageVersionSettings
-import org.jetbrains.kotlin.idea.references.mainReference
-import org.jetbrains.kotlin.idea.stubindex.PackageIndexUtil
-import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.platform.jvm.isJvm
-import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
-import org.jetbrains.kotlin.psi.KtExpression
 
 internal class FirPackageCompletionContributor(
     basicContext: FirBasicCompletionContext,
@@ -36,7 +29,7 @@ internal class FirPackageCompletionContributor(
                 packageName.fqName.isExcludedFromAutoImport(project, originalKtFile, originalKtFile.languageVersionSettings)
             }
             .forEach { packageSymbol ->
-                result.addElement(lookupElementFactory.createPackageLookupElement(packageSymbol.fqName))
+                result.addElement(lookupElementFactory.createPackagePartLookupElement(packageSymbol.fqName))
             }
     }
 }

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

@@ -37,15 +37,14 @@ internal object SuperCallInsertionHandler : InsertHandler<LookupElement> {
         val lookupObject = item.`object` as SuperCallLookupObject
 
         replaceWithClassIdAndShorten(lookupObject, context)
-        addDot(context)
-        invokeCompletion(context)
+        context.addDotAndInvokeCompletion()
     }
 
     private fun replaceWithClassIdAndShorten(
         lookupObject: SuperCallLookupObject,
         context: InsertionContext
     ) {
-        val replaceTo = lookupObject.replaceTo  ?: return
+        val replaceTo = lookupObject.replaceTo ?: return
         context.document.replaceString(context.startOffset, context.tailOffset, replaceTo)
         context.commitDocument()
 
@@ -54,19 +53,6 @@ internal object SuperCallInsertionHandler : InsertHandler<LookupElement> {
             shortenReferencesForFirCompletion(targetFile, TextRange(context.startOffset, context.tailOffset))
         }
     }
-
-    private fun invokeCompletion(context: InsertionContext) {
-        ApplicationManager.getApplication().invokeLater {
-            CodeCompletionHandlerBase(CompletionType.BASIC, true, false, true)
-                .invokeCompletion(context.project, context.editor)
-        }
-    }
-
-    private fun addDot(context: InsertionContext) {
-        context.document.insertString(context.tailOffset, ".")
-        context.commitDocument()
-        context.editor.caretModel.moveToOffset(context.tailOffset)
-    }
 }
 
 internal interface SuperCallLookupObject {

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

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2010-2021 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.contributors.helpers
+
+import com.intellij.codeInsight.completion.CodeCompletionHandlerBase
+import com.intellij.codeInsight.completion.CompletionType
+import com.intellij.codeInsight.completion.InsertionContext
+import com.intellij.openapi.application.ApplicationManager
+
+internal fun InsertionContext.addDotAndInvokeCompletion() {
+    addDotToCompletion(this)
+    invokeCompletion(this)
+}
+
+private fun invokeCompletion(context: InsertionContext) {
+    ApplicationManager.getApplication().invokeLater {
+        CodeCompletionHandlerBase(CompletionType.BASIC, true, false, true)
+            .invokeCompletion(context.project, context.editor)
+    }
+}
+
+private fun addDotToCompletion(context: InsertionContext) {
+    context.document.insertString(context.tailOffset, ".")
+    context.commitDocument()
+    context.editor.caretModel.moveToOffset(context.tailOffset)
+}

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

@@ -0,0 +1,19 @@
+/*
+ * Copyright 2010-2021 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.contributors.helpers
+
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.scopes.KtScope
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtPackageSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSymbolWithMembers
+import org.jetbrains.kotlin.idea.references.KtReference
+
+internal fun KtAnalysisSession.getStaticScope(reference: KtReference): KtScope? =
+    when (val symbol = reference.resolveToSymbol()) {
+        is KtSymbolWithMembers -> symbol.getStaticMemberScope()
+        is KtPackageSymbol -> symbol.getPackageScope()
+        else -> null
+    }

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

@@ -19,6 +19,8 @@ import org.jetbrains.kotlin.psi.KtFile
 
 abstract class KtScopeProvider : KtAnalysisSessionComponent() {
     abstract fun getMemberScope(classSymbol: KtSymbolWithMembers): KtMemberScope
+    abstract fun getStaticMemberScope(symbol: KtSymbolWithMembers): KtScope
+
     abstract fun getDeclaredMemberScope(classSymbol: KtSymbolWithMembers): KtDeclaredMemberScope
     abstract fun getFileScope(fileSymbol: KtFileSymbol): KtDeclarationScope<KtSymbolWithDeclarations>
     abstract fun getPackageScope(packageSymbol: KtPackageSymbol): KtPackageScope
@@ -39,6 +41,9 @@ interface KtScopeProviderMixIn : KtAnalysisSessionMixIn {
     fun KtSymbolWithMembers.getDeclaredMemberScope(): KtDeclaredMemberScope =
         analysisSession.scopeProvider.getDeclaredMemberScope(this)
 
+    fun KtSymbolWithMembers.getStaticMemberScope(): KtScope =
+        analysisSession.scopeProvider.getStaticMemberScope(this)
+
     fun KtFileSymbol.getFileScope(): KtDeclarationScope<KtSymbolWithDeclarations> =
         analysisSession.scopeProvider.getFileScope(this)
 

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

@@ -86,6 +86,15 @@ internal class KtFirScopeProvider(
         }
     }
 
+    override fun getStaticMemberScope(symbol: KtSymbolWithMembers): KtScope {
+        val firScope = symbol.withFirForScope { fir ->
+            fir.scopeProvider.getStaticScope(fir, analysisSession.rootModuleSession, ScopeSession())
+        } ?: return KtFirEmptyMemberScope(symbol)
+        firScopeStorage.register(firScope)
+        check(firScope is FirContainingNamesAwareScope)
+        return KtFirDelegatingScopeImpl(firScope, builder, token)
+    }
+
     override fun getDeclaredMemberScope(classSymbol: KtSymbolWithMembers): KtDeclaredMemberScope = withValidityAssertion {
         declaredMemberScopeCache.getOrPut(classSymbol) {
             val firScope = classSymbol.withFirForScope {

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

@@ -79,6 +79,13 @@ internal fun FirScope.getCallableSymbols(callableNames: Collection<Name>, builde
     }
 }
 
+internal class KtFirDelegatingScopeImpl<S>(
+    override val firScope: S,
+    builder: KtSymbolByFirBuilder,
+    token: ValidityToken
+) : KtFirDelegatingScope<S>(builder, token) where S : FirContainingNamesAwareScope, S : FirScope
+
+
 internal fun FirScope.getClassifierSymbols(classLikeNames: Collection<Name>, builder: KtSymbolByFirBuilder): Sequence<KtClassifierSymbol> =
     sequence {
         classLikeNames.forEach { name ->

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

@@ -62,6 +62,11 @@ internal class KtFirPackageScope(
 
             KotlinTopLevelTypeAliasByPackageIndex.getInstance()[fqName.asString(), project, searchScope]
                 .mapNotNullTo(this) { it.nameAsName }
+
+            JavaPsiFacade.getInstance(project)
+                .findPackage(fqName.asString())
+                ?.getClasses(searchScope)
+                ?.mapNotNullTo(this) { it.name?.let(Name::identifier) }
         }
     }
 
@@ -74,21 +79,7 @@ internal class KtFirPackageScope(
     }
 
     override fun getPackageSymbols(nameFilter: KtScopeNameFilter): Sequence<KtPackageSymbol> = withValidityAssertion {
-        sequence {
-            PackageIndexUtil
-                .getSubPackageFqNames(fqName, searchScope, project, nameFilter)
-                .forEach {
-                    yield(builder.createPackageSymbol(it))
-                }
-
-            if (targetPlatform.isJvm()) {
-                JavaPsiFacade.getInstance(project).findPackage(fqName.asString())?.getSubPackages(searchScope)?.forEach { psiPackage ->
-                    val fqName = psiPackage.getKotlinFqName() ?: return@forEach
-                    if (nameFilter(fqName.shortName())) {
-                        yield(builder.createPackageSymbol(fqName))
-                    }
-                }
-            }
-        }
+        PackageIndexUtil.getJavaAndKotlinSubPackageFqNames(fqName, searchScope, project, targetPlatform, nameFilter)
+            .map { builder.createPackageSymbol(it) }
     }
 }

idea/idea-frontend-independent/src/org/jetbrains/kotlin/miniStdLib/chainUtils.kt

@@ -0,0 +1,18 @@
+/*
+ * Copyright 2010-2021 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.miniStdLib
+
+import kotlin.contracts.ExperimentalContracts
+import kotlin.contracts.InvocationKind
+import kotlin.contracts.contract
+
+@OptIn(ExperimentalContracts::class)
+inline fun <T : R, R> T.letIf(condition: Boolean, block: (T) -> R): R {
+    contract {
+        callsInPlace(block, InvocationKind.AT_MOST_ONCE)
+    }
+    return if (condition) block(this) else this
+}