[kotlin] KTIJ-33671, KTIJ-23721 chained classifiers completion

GitOrigin-RevId: 15f16f46848b2840414db55f79c9f95c35863b2a

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/FirCallableCompletionContributor.kt

@@ -3,7 +3,6 @@ package org.jetbrains.kotlin.idea.completion.impl.k2.contributors
 
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.util.NlsSafe
-import com.intellij.openapi.util.registry.RegistryManager
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiEnumConstant
 import com.intellij.psi.PsiField
@@ -12,7 +11,6 @@ import com.intellij.psi.util.parents
 import com.intellij.util.containers.addIfNotNull
 import org.jetbrains.kotlin.analysis.api.KaExperimentalApi
 import org.jetbrains.kotlin.analysis.api.KaSession
-import org.jetbrains.kotlin.analysis.api.analyze
 import org.jetbrains.kotlin.analysis.api.components.*
 import org.jetbrains.kotlin.analysis.api.lifetime.KaLifetimeOwner
 import org.jetbrains.kotlin.analysis.api.lifetime.KaLifetimeToken
@@ -40,13 +38,11 @@ import org.jetbrains.kotlin.idea.completion.impl.k2.context.getOriginalDeclarati
 import org.jetbrains.kotlin.idea.completion.lookups.CallableInsertionOptions
 import org.jetbrains.kotlin.idea.completion.lookups.CallableInsertionStrategy
 import org.jetbrains.kotlin.idea.completion.lookups.ImportStrategy
-import org.jetbrains.kotlin.idea.completion.lookups.factories.ClassifierLookupObject
 import org.jetbrains.kotlin.idea.completion.lookups.factories.FunctionInsertionHelper
 import org.jetbrains.kotlin.idea.completion.reference
 import org.jetbrains.kotlin.idea.completion.weighers.CallableWeigher.callableWeight
 import org.jetbrains.kotlin.idea.completion.weighers.WeighingContext
 import org.jetbrains.kotlin.idea.references.mainReference
-import org.jetbrains.kotlin.idea.util.positionContext.KotlinExpressionNameReferencePositionContext
 import org.jetbrains.kotlin.idea.util.positionContext.KotlinNameReferencePositionContext
 import org.jetbrains.kotlin.idea.util.positionContext.KotlinSimpleNameReferencePositionContext
 import org.jetbrains.kotlin.kdoc.psi.impl.KDocName
@@ -54,7 +50,6 @@ import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.platform.isMultiPlatform
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.nextSiblingOfSameType
-import org.jetbrains.kotlin.renderer.render
 import org.jetbrains.kotlin.resolve.ArrayFqNames
 import org.jetbrains.kotlin.types.Variance
 
@@ -340,47 +335,28 @@ internal open class FirCallableCompletionContributor(
                 }
 
                 if (showReceiver) return@sequence
-                if (!RegistryManager.getInstance().`is`("kotlin.k2.chain.completion.enabled")) return@sequence
-                sink.runRemainingContributors(parameters.delegate) { completionResult ->
-                    val lookupElement = completionResult.lookupElement
-                    val (_, importStrategy) = lookupElement.`object` as? ClassifierLookupObject
-                        ?: return@runRemainingContributors
+                runChainCompletion(positionContext, explicitReceiver) { receiverExpression,
+                                                                        positionContext,
+                                                                        importingStrategy ->
+                    val weighingContext = WeighingContext.create(parameters, positionContext)
 
-                    val nameToImport = when (importStrategy) {
-                        is ImportStrategy.AddImport -> importStrategy.nameToImport
-                        is ImportStrategy.InsertFqNameAndShorten -> importStrategy.fqName
-                        ImportStrategy.DoNothing -> null
-                    } ?: return@runRemainingContributors
+                    collectDotCompletion(
+                        positionContext = positionContext,
+                        scopeContext = weighingContext.scopeContext!!,
+                        explicitReceiver = receiverExpression,
+                        extensionChecker = null,
+                        showReceiver = true,
+                    ).flatMap { callableWithMetadata ->
+                        val signature = callableWithMetadata.signature
 
-                    val expression = KtPsiFactory.contextual(explicitReceiver)
-                        .createExpression(nameToImport.render() + "." + positionContext.nameExpression.text) as KtDotQualifiedExpression
-
-                    val nameExpression = expression.selectorExpression as? KtNameReferenceExpression ?: return@runRemainingContributors
-
-                    analyze(nameExpression) {
-                        val receiverExpression = expression.receiverExpression
-
-                        val positionContext = KotlinExpressionNameReferencePositionContext(nameExpression)
-                        val weighingContext = WeighingContext.create(parameters, positionContext)
-
-                        collectDotCompletion(
-                            positionContext = positionContext,
-                            scopeContext = weighingContext.scopeContext!!,
-                            explicitReceiver = receiverExpression,
-                            extensionChecker = null,
-                            showReceiver = true,
-                        ).flatMap { callableWithMetadata ->
-                            val signature = callableWithMetadata.signature
-
-                            createCallableLookupElements(
-                                context = weighingContext,
-                                signature = signature,
-                                options = callableWithMetadata.options.copy(importingStrategy = ImportStrategy.AddImport(nameToImport)),
-                                symbolOrigin = callableWithMetadata.symbolOrigin,
-                                presentableText = callableWithMetadata.itemText,
-                                withTrailingLambda = true,
-                            )
-                        }.forEach(sink::addElement)
+                        createCallableLookupElements(
+                            context = weighingContext,
+                            signature = signature,
+                            options = callableWithMetadata.options.copy(importingStrategy),
+                            symbolOrigin = callableWithMetadata.symbolOrigin,
+                            presentableText = callableWithMetadata.itemText,
+                            withTrailingLambda = true,
+                        )
                     }
                 }
             }

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/FirClassifierCompletionContributor.kt

@@ -14,6 +14,7 @@ import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.util.concurrency.annotations.RequiresReadLock
 import org.jetbrains.kotlin.analysis.api.KaSession
 import org.jetbrains.kotlin.analysis.api.analyze
+import org.jetbrains.kotlin.analysis.api.components.KaScopeWithKind
 import org.jetbrains.kotlin.analysis.api.components.ShortenCommand
 import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.analysis.api.types.KaClassType
@@ -31,7 +32,6 @@ import org.jetbrains.kotlin.idea.completion.lookups.factories.shortenCommand
 import org.jetbrains.kotlin.idea.completion.reference
 import org.jetbrains.kotlin.idea.completion.weighers.Weighers.applyWeighs
 import org.jetbrains.kotlin.idea.completion.weighers.WeighingContext
-import org.jetbrains.kotlin.idea.references.KtReference
 import org.jetbrains.kotlin.idea.util.positionContext.KotlinNameReferencePositionContext
 import org.jetbrains.kotlin.idea.util.positionContext.KotlinTypeNameReferencePositionContext
 import org.jetbrains.kotlin.name.FqName
@@ -61,44 +61,71 @@ internal open class FirClassifierCompletionContributor(
         positionContext: KotlinNameReferencePositionContext,
         weighingContext: WeighingContext,
     ) {
+        val defaultFactory = (psiFactory::createExpression).asFactory()
         val factory = when (positionContext) {
             is KotlinTypeNameReferencePositionContext -> ({ type: String -> psiFactory.createType(type) }).asFactory()
-            else -> (psiFactory::createExpression).asFactory()
+            else -> defaultFactory
         }
 
-        when (val receiver = positionContext.explicitReceiver) {
+        when (val explicitReceiver = positionContext.explicitReceiver) {
             null -> completeWithoutReceiver(positionContext, factory, weighingContext)
+                .forEach(sink::addElement)
 
             else -> {
-                receiver.reference()?.let {
-                    completeWithReceiver(positionContext, factory, weighingContext, it)
-                } ?: emptySequence()
+                val reference = explicitReceiver.reference()
+                    ?: return
+
+                val symbols = reference.resolveToSymbols()
+                if (symbols.isNotEmpty()) {
+                    symbols.asSequence()
+                        .mapNotNull { it.staticScope }
+                        .flatMap { scopeWithKind ->
+                            scopeWithKind.completeClassifiers(positionContext)
+                                .mapNotNull { classifierSymbol ->
+                                    createClassifierLookupElement(classifierSymbol, factory)?.applyWeighs(
+                                        context = weighingContext,
+                                        symbolWithOrigin = KtSymbolWithOrigin(
+                                            classifierSymbol,
+                                            CompletionSymbolOrigin.Scope(scopeWithKind.kind)
+                                        )
+                                    )
+                                }
+                        }.forEach(sink::addElement)
+                } else {
+                    runChainCompletion(positionContext, explicitReceiver) { receiverExpression,
+                                                                            positionContext,
+                                                                            importingStrategy ->
+                        val selectorExpression = receiverExpression.selectorExpression
+                            ?: return@runChainCompletion emptySequence()
+
+                        val reference = receiverExpression.reference()
+                            ?: return@runChainCompletion emptySequence()
+
+                        // TODO val weighingContext = WeighingContext.create(parameters, positionContext)
+                        reference.resolveToSymbols()
+                            .asSequence()
+                            .mapNotNull { it.staticScope }
+                            .flatMap { it.completeClassifiers(positionContext) }
+                            .mapNotNull { it ->
+                                createClassifierLookupElement(
+                                    classifierSymbol = it,
+                                    factory = defaultFactory,
+                                    importingStrategy = importingStrategy,
+                                )
+                            }.map { it.withPresentableText(selectorExpression.text + "." + it.lookupString) }
+                    }
+                }
             }
-        }.forEach(sink::addElement)
+        }
     }
 
     context(KaSession)
-    private fun completeWithReceiver(
+    private fun KaScopeWithKind.completeClassifiers(
         positionContext: KotlinNameReferencePositionContext,
-        factory: MethodBasedElementFactory<*>,
-        context: WeighingContext,
-        reference: KtReference,
-    ): Sequence<LookupElementBuilder> = reference
-        .resolveToSymbols()
-        .asSequence()
-        .mapNotNull { it.staticScope }
-        .flatMap { scopeWithKind ->
-            scopeWithKind.scope
-                .classifiers(scopeNameFilter)
-                .filter { filterClassifiers(it) }
-                .filter { visibilityChecker.isVisible(it, positionContext) }
-                .mapNotNull { classifierSymbol ->
-                    createClassifierLookupElement(classifierSymbol, factory)?.applyWeighs(
-                        context = context,
-                        symbolWithOrigin = KtSymbolWithOrigin(classifierSymbol, CompletionSymbolOrigin.Scope(scopeWithKind.kind))
-                    )
-                }
-        }
+    ): Sequence<KaClassifierSymbol> = scope
+        .classifiers(scopeNameFilter)
+        .filter { filterClassifiers(it) }
+        .filter { visibilityChecker.isVisible(it, positionContext) }
 
     context(KaSession)
     private fun completeWithoutReceiver(

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/FirCompletionContributorBase.kt

@@ -7,7 +7,9 @@ import com.intellij.codeInsight.lookup.LookupElementBuilder
 import com.intellij.codeInsight.lookup.LookupElementPresentation
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.NlsSafe
+import com.intellij.openapi.util.registry.RegistryManager
 import org.jetbrains.kotlin.analysis.api.KaSession
+import org.jetbrains.kotlin.analysis.api.analyze
 import org.jetbrains.kotlin.analysis.api.signatures.KaCallableSignature
 import org.jetbrains.kotlin.analysis.api.symbols.KaConstructorSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KaNamedClassSymbol
@@ -26,15 +28,19 @@ import org.jetbrains.kotlin.idea.completion.doPostponedOperationsAndUnblockDocum
 import org.jetbrains.kotlin.idea.completion.impl.k2.ImportStrategyDetector
 import org.jetbrains.kotlin.idea.completion.impl.k2.LookupElementSink
 import org.jetbrains.kotlin.idea.completion.lookups.CallableInsertionOptions
+import org.jetbrains.kotlin.idea.completion.lookups.ImportStrategy
+import org.jetbrains.kotlin.idea.completion.lookups.factories.ClassifierLookupObject
 import org.jetbrains.kotlin.idea.completion.lookups.factories.FunctionCallLookupObject
 import org.jetbrains.kotlin.idea.completion.lookups.factories.KotlinFirLookupElementFactory
 import org.jetbrains.kotlin.idea.completion.weighers.CallableWeigher.callableWeight
 import org.jetbrains.kotlin.idea.completion.weighers.Weighers.applyWeighs
 import org.jetbrains.kotlin.idea.completion.weighers.WeighingContext
+import org.jetbrains.kotlin.idea.util.positionContext.KotlinExpressionNameReferencePositionContext
+import org.jetbrains.kotlin.idea.util.positionContext.KotlinNameReferencePositionContext
 import org.jetbrains.kotlin.idea.util.positionContext.KotlinRawPositionContext
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.renderer.render
 
 internal abstract class FirCompletionContributorBase<C : KotlinRawPositionContext>(
     protected val parameters: KotlinFirCompletionParameters,
@@ -115,6 +121,47 @@ internal abstract class FirCompletionContributorBase<C : KotlinRawPositionContex
         }
     }
 
+    protected fun runChainCompletion(
+        positionContext: KotlinNameReferencePositionContext,
+        explicitReceiver: KtElement,
+        createLookupElements: KaSession.(
+            receiverExpression: KtDotQualifiedExpression,
+            positionContext: KotlinExpressionNameReferencePositionContext,
+            importingStrategy: ImportStrategy.AddImport,
+        ) -> Sequence<LookupElement>,
+    ) {
+        if (!RegistryManager.getInstance().`is`("kotlin.k2.chain.completion.enabled")) return
+
+        sink.runRemainingContributors(parameters.delegate) { completionResult ->
+            val lookupElement = completionResult.lookupElement
+            val (_, importStrategy) = lookupElement.`object` as? ClassifierLookupObject
+                ?: return@runRemainingContributors
+
+            val nameToImport = when (importStrategy) {
+                is ImportStrategy.AddImport -> importStrategy.nameToImport
+                is ImportStrategy.InsertFqNameAndShorten -> importStrategy.fqName
+                ImportStrategy.DoNothing -> null
+            } ?: return@runRemainingContributors
+
+            val expression = KtPsiFactory.contextual(explicitReceiver)
+                .createExpression(nameToImport.render() + "." + positionContext.nameExpression.text) as KtDotQualifiedExpression
+
+            val receiverExpression = expression.receiverExpression as? KtDotQualifiedExpression
+                ?: return@runRemainingContributors
+
+            val nameExpression = expression.selectorExpression as? KtNameReferenceExpression
+                ?: return@runRemainingContributors
+
+            analyze(nameExpression) {
+                createLookupElements(
+                    /* receiverExpression = */ receiverExpression,
+                    /* positionContext = */ KotlinExpressionNameReferencePositionContext(nameExpression),
+                    /* importingStrategy = */ ImportStrategy.AddImport(nameToImport),
+                ).forEach(sink::addElement)
+            }
+        }
+    }
+
     // todo move out
     // todo move to the corresponding assignment
     protected fun LookupElementBuilder.adaptToExplicitReceiver(