FIR IDE: add this labels to keyword completion

idea/ide-common/build.gradle.kts

@@ -8,6 +8,7 @@ dependencies {
     compileOnly(project(":compiler:util"))
     compileOnly(project(":compiler:frontend"))
     compileOnly(project(":compiler:frontend.java"))
+    compileOnly(project(":idea:idea-frontend-independent"))
     compileOnly(project(":js:js.frontend"))
     compileOnly(project(":js:js.serializer"))
     compileOnly(intellijCoreDep()) { includeJars("intellij-core", "guava", rootProject = rootProject) }

idea/ide-common/src/org/jetbrains/kotlin/idea/util/Utils.kt

@@ -18,31 +18,6 @@ import org.jetbrains.kotlin.types.FlexibleType
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.checker.KotlinTypeChecker
 
-fun KtFunctionLiteral.findLabelAndCall(): Pair<Name?, KtCallExpression?> {
-    val literalParent = (this.parent as KtLambdaExpression).parent
-
-    fun KtValueArgument.callExpression(): KtCallExpression? {
-        val parent = parent
-        return (if (parent is KtValueArgumentList) parent else this).parent as? KtCallExpression
-    }
-
-    when (literalParent) {
-        is KtLabeledExpression -> {
-            val callExpression = (literalParent.parent as? KtValueArgument)?.callExpression()
-            return Pair(literalParent.getLabelNameAsName(), callExpression)
-        }
-
-        is KtValueArgument -> {
-            val callExpression = literalParent.callExpression()
-            val label = (callExpression?.calleeExpression as? KtSimpleNameExpression)?.getReferencedNameAsName()
-            return Pair(label, callExpression)
-        }
-
-        else -> {
-            return Pair(null, null)
-        }
-    }
-}
 
 fun SmartCastManager.getSmartCastVariantsWithLessSpecificExcluded(
     receiverToCast: ReceiverValue,

idea/idea-completion/testData/keywords/LabeledLambdaThis.kt

@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 fun foo(): String.() -> Unit {
     return (label@ {
         f {

idea/idea-completion/testData/keywords/This.kt

@@ -1,3 +1,4 @@
+// FIR_COMPARISON
 class Outer {
     class Nested {
         inner class Inner {

idea/idea-completion/testData/keywords/ThisInCompanion.kt

@@ -0,0 +1,11 @@
+// FIR_COMPARISON
+class C {
+    companion object {
+        fun String.foo() {
+            <caret>
+        }
+    }
+}
+
+// INVOCATION_COUNT: 1
+// EXIST: "this@Companion"

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

@@ -579,6 +579,11 @@ public class KeywordCompletionTestGenerated extends AbstractKeywordCompletionTes
         runTest("idea/idea-completion/testData/keywords/This.kt");
     }
 
+    @TestMetadata("ThisInCompanion.kt")
+    public void testThisInCompanion() throws Exception {
+        runTest("idea/idea-completion/testData/keywords/ThisInCompanion.kt");
+    }
+
     @TestMetadata("ThisPrefixMatching.kt")
     public void testThisPrefixMatching() throws Exception {
         runTest("idea/idea-completion/testData/keywords/ThisPrefixMatching.kt");

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

@@ -226,7 +226,8 @@ private class KotlinWithNameReferenceCompletionProvider(
         extensionChecker: ExtensionApplicabilityChecker,
         visibilityChecker: CompletionVisibilityChecker,
     ) {
-        val (implicitScopes, implicitReceiversTypes) = implicitScopesContext
+        val (implicitScopes, implicitReceivers) = implicitScopesContext
+        val implicitReceiversTypes = implicitReceivers.map { it.type }
 
         val availableNonExtensions = implicitScopes.collectNonExtensions(visibilityChecker)
         val extensionsWhichCanBeCalled = implicitScopes.collectSuitableExtensions(extensionChecker, visibilityChecker)

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

@@ -0,0 +1,83 @@
+/*
+ * 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.keywords
+
+import com.intellij.codeInsight.completion.CompletionParameters
+import com.intellij.codeInsight.lookup.LookupElement
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
+import com.intellij.psi.util.isAncestor
+import org.jetbrains.kotlin.idea.completion.KeywordLookupObject
+import org.jetbrains.kotlin.idea.completion.context.FirBasicCompletionContext
+import org.jetbrains.kotlin.idea.completion.createKeywordElement
+import org.jetbrains.kotlin.idea.completion.keywords.CompletionKeywordHandler
+import org.jetbrains.kotlin.idea.completion.labelNameToTail
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.components.KtImplicitReceiver
+import org.jetbrains.kotlin.idea.frontend.api.components.KtTypeRendererOptions
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtAnonymousFunctionSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassKind
+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.util.findLabelAndCall
+import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtClassOrObject
+import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtFunctionLiteral
+import org.jetbrains.kotlin.psi.KtPsiUtil
+
+internal class ThisKeywordHandler(
+    private val basicContext: FirBasicCompletionContext
+) : CompletionKeywordHandler<KtAnalysisSession>(KtTokens.THIS_KEYWORD) {
+    @OptIn(ExperimentalStdlibApi::class)
+    override fun KtAnalysisSession.createLookups(
+        parameters: CompletionParameters,
+        expression: KtExpression?,
+        lookup: LookupElement,
+        project: Project
+    ): Collection<LookupElement> {
+        if (expression == null) {
+            // for completion in secondary constructor delegation call
+            return listOf(lookup)
+        }
+
+        val result = mutableListOf<LookupElement>()
+        val receivers = basicContext.originalKtFile.getScopeContextForPosition(expression).implicitReceivers
+
+        receivers.forEachIndexed { index, receiver ->
+            if (!canReferenceSymbolByThis(parameters, receiver.ownerSymbol)) {
+                return@forEachIndexed
+            }
+            val labelName = if (index != 0) getThisLabelBySymbol(receiver.ownerSymbol) else null
+            result += createThisLookupElement(receiver, labelName)
+        }
+
+        return result
+    }
+
+    private fun KtAnalysisSession.canReferenceSymbolByThis(parameters: CompletionParameters, symbol: KtSymbol): Boolean {
+        if (symbol !is KtClassOrObjectSymbol) return true
+        if (symbol.classKind != KtClassKind.COMPANION_OBJECT) return true
+        val companionPsi = symbol.psi as KtClassOrObject
+        return parameters.offset in companionPsi.textRange
+    }
+
+    private fun KtAnalysisSession.createThisLookupElement(receiver: KtImplicitReceiver, labelName: Name?): LookupElement {
+        return createKeywordElement("this", labelName.labelNameToTail(), lookupObject = KeywordLookupObject())
+            .withTypeText(receiver.type.render(KtTypeRendererOptions.SHORT_NAMES))
+    }
+
+    private fun KtAnalysisSession.getThisLabelBySymbol(symbol: KtSymbol): Name? = when {
+        symbol is KtNamedSymbol && !symbol.name.isSpecial -> symbol.name
+        symbol is KtAnonymousFunctionSymbol -> {
+            val psi = symbol.psi as KtFunctionLiteral
+            psi.findLabelAndCall().first
+        }
+        else -> null
+    }
+}

idea/idea-fir/tests/org/jetbrains/kotlin/idea/completion/FirKeywordCompletionTestGenerated.java

@@ -579,6 +579,11 @@ public class FirKeywordCompletionTestGenerated extends AbstractFirKeywordComplet
         runTest("idea/idea-completion/testData/keywords/This.kt");
     }
 
+    @TestMetadata("ThisInCompanion.kt")
+    public void testThisInCompanion() throws Exception {
+        runTest("idea/idea-completion/testData/keywords/ThisInCompanion.kt");
+    }
+
     @TestMetadata("ThisPrefixMatching.kt")
     public void testThisPrefixMatching() throws Exception {
         runTest("idea/idea-completion/testData/keywords/ThisPrefixMatching.kt");

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

@@ -5,11 +5,14 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.components
 
+import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.scopes.*
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtFileSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtPackageSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSymbolWithDeclarations
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSymbolWithMembers
+import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
@@ -55,4 +58,10 @@ interface KtScopeProviderMixIn : KtAnalysisSessionMixIn {
         analysisSession.scopeProvider.getScopeContextForPosition(this, this)
 }
 
-data class KtScopeContext(val scopes: KtCompositeScope, val implicitReceiversTypes: List<KtType>)
+data class KtScopeContext(val scopes: KtCompositeScope, val implicitReceivers: List<KtImplicitReceiver>)
+
+class KtImplicitReceiver(
+    override val token: ValidityToken,
+    val type: KtType,
+    val ownerSymbol: KtSymbol
+): ValidityTokenOwner

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

@@ -9,6 +9,7 @@ import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousObject
 import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.calls.FirSyntheticPropertiesScope
@@ -18,6 +19,7 @@ import org.jetbrains.kotlin.fir.scopes.impl.*
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getTowerDataContextUnsafe
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
+import org.jetbrains.kotlin.idea.frontend.api.components.KtImplicitReceiver
 import org.jetbrains.kotlin.idea.frontend.api.components.KtScopeContext
 import org.jetbrains.kotlin.idea.frontend.api.components.KtScopeProvider
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
@@ -143,7 +145,13 @@ internal class KtFirScopeProvider(
         val towerDataContext = analysisSession.firResolveState.getTowerDataContextUnsafe(positionInFakeFile)
 
         val implicitReceivers = towerDataContext.nonLocalTowerDataElements.mapNotNull { it.implicitReceiver }.distinct()
-        val implicitReceiversTypes = implicitReceivers.map { builder.typeBuilder.buildKtType(it.type) }
+        val implicitKtReceivers = implicitReceivers.map { receiver ->
+            KtImplicitReceiver(
+                token,
+                builder.typeBuilder.buildKtType(receiver.type),
+                builder.buildSymbol(receiver.boundSymbol.fir as FirDeclaration),
+            )
+        }
 
         val implicitReceiverScopes = implicitReceivers.mapNotNull { it.implicitScope }
         val nonLocalScopes = towerDataContext.nonLocalTowerDataElements.mapNotNull { it.scope }.distinct()
@@ -158,7 +166,7 @@ internal class KtFirScopeProvider(
 
         KtScopeContext(
             getCompositeScope(allKtScopes.asReversed()),
-            implicitReceiversTypes.asReversed()
+            implicitKtReceivers.asReversed()
         )
     }
 

idea/idea-frontend-independent/src/org/jetbrains/kotlin/idea/util/findLabelAndCall.kt

@@ -0,0 +1,35 @@
+/*
+ * 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.util
+
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.*
+
+fun KtFunctionLiteral.findLabelAndCall(): Pair<Name?, KtCallExpression?> {
+    val literalParent = (this.parent as KtLambdaExpression).parent
+
+    fun KtValueArgument.callExpression(): KtCallExpression? {
+        val parent = parent
+        return (if (parent is KtValueArgumentList) parent else this).parent as? KtCallExpression
+    }
+
+    when (literalParent) {
+        is KtLabeledExpression -> {
+            val callExpression = (literalParent.parent as? KtValueArgument)?.callExpression()
+            return Pair(literalParent.getLabelNameAsName(), callExpression)
+        }
+
+        is KtValueArgument -> {
+            val callExpression = literalParent.callExpression()
+            val label = (callExpression?.calleeExpression as? KtSimpleNameExpression)?.getReferencedNameAsName()
+            return Pair(label, callExpression)
+        }
+
+        else -> {
+            return Pair(null, null)
+        }
+    }
+}

plugins/scripting/scripting-ide-services-test/build.gradle.kts

@@ -28,6 +28,7 @@ dependencies {
     testRuntimeOnly(project(":kotlin-compiler"))
     testRuntimeOnly(commonDep("org.jetbrains.intellij.deps", "trove4j"))
     testRuntimeOnly(project(":idea:ide-common")) { isTransitive = false }
+    testRuntimeOnly(project(":idea:idea-frontend-independent")) { isTransitive = false }
 
     embeddableTestRuntime(project(":kotlin-scripting-ide-services", configuration="runtimeElements"))
     embeddableTestRuntime(project(":kotlin-scripting-compiler-impl-embeddable", configuration="runtimeElements"))