FIR IDE: add KtExtensionReceiverSymbol

This symbol represents the extension receiver of a callable and currently is used as the resolution target of `this` in a callable with extension.
2021-09-22 15:38:13 -07:00
parent d3c91ee9be
commit e195f0426e
17 changed files with 161 additions and 13 deletions
@@ -56,6 +56,7 @@ internal class KtSymbolByFirBuilder private constructor(
    override val token: ValidityToken,
    val withReadOnlyCaching: Boolean,
    private val symbolsCache: BuilderCache<FirDeclaration, KtSymbol>,
+    private val extensionReceiverSymbolsCache: BuilderCache<FirCallableDeclaration, KtSymbol>,
    private val filesCache: BuilderCache<FirFile, KtFileSymbol>,
    private val backingFieldCache: BuilderCache<FirBackingField, KtBackingFieldSymbol>,
    private val typesCache: BuilderCache<ConeKotlinType, KtType>,
@@ -82,6 +83,7 @@ internal class KtSymbolByFirBuilder private constructor(
        resolveState = resolveState,
        withReadOnlyCaching = false,
        symbolsCache = BuilderCache(),
+        extensionReceiverSymbolsCache = BuilderCache(),
        typesCache = BuilderCache(),
        backingFieldCache = BuilderCache(),
        filesCache = BuilderCache(),
@@ -95,6 +97,7 @@ internal class KtSymbolByFirBuilder private constructor(
            resolveState = newResolveState,
            withReadOnlyCaching = true,
            symbolsCache = symbolsCache.createReadOnlyCopy(),
+            extensionReceiverSymbolsCache = extensionReceiverSymbolsCache.createReadOnlyCopy(),
            typesCache = typesCache.createReadOnlyCopy(),
            filesCache = filesCache.createReadOnlyCopy(),
            backingFieldCache = backingFieldCache.createReadOnlyCopy(),
@@ -337,6 +340,13 @@ internal class KtSymbolByFirBuilder private constructor(
            checkRequirementForBuildingSymbol<KtFirPropertySetterSymbol>(fir, fir.isSetter)
            return symbolsCache.cache(fir) { KtFirPropertySetterSymbol(fir, resolveState, token, this@KtSymbolByFirBuilder) }
        }
+
+        fun buildExtensionReceiverSymbol(fir: FirCallableDeclaration): KtReceiverParameterSymbol? {
+            if (fir.receiverTypeRef == null) return null
+            return extensionReceiverSymbolsCache.cache(fir) {
+                KtFirReceiverParameterSymbol(fir, resolveState, token, this@KtSymbolByFirBuilder)
+            }
+        }
    }

    inner class AnonymousInitializerBuilder {
@@ -6,26 +6,35 @@
 package org.jetbrains.kotlin.analysis.api.fir.components

 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
-import org.jetbrains.kotlin.fir.FirRealSourceElementKind
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
-import org.jetbrains.kotlin.fir.psi
-import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
 import org.jetbrains.kotlin.analysis.api.components.KtSymbolContainingDeclarationProvider
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirReceiverParameterSymbol
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirSymbol
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirTypeParameterSymbol
 import org.jetbrains.kotlin.analysis.api.impl.barebone.parentOfType
 import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolKind
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithKind
-import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
+import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
+import org.jetbrains.kotlin.fir.FirRealSourceElementKind
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtDestructuringDeclaration
+import org.jetbrains.kotlin.psi.KtPrimaryConstructor
+import org.jetbrains.kotlin.psi.KtProperty

 internal class KtFirSymbolContainingDeclarationProvider(
    override val analysisSession: KtFirAnalysisSession,
    override val token: ValidityToken,
 ) : KtSymbolContainingDeclarationProvider(), KtFirAnalysisSessionComponent {
    override fun getContainingDeclaration(symbol: KtSymbol): KtSymbolWithKind? {
+        if (symbol is KtReceiverParameterSymbol) {
+            return (symbol as KtFirReceiverParameterSymbol).firRef.withFir {
+                firSymbolBuilder.buildSymbol(it) as KtSymbolWithKind
+            }
+        }
        if (symbol is KtPackageSymbol) return null
        if (symbol is KtSymbolWithKind && symbol.symbolKind == KtSymbolKind.TOP_LEVEL) return null
        if (symbol is KtCallableSymbol) {
@@ -5,19 +5,21 @@

 package org.jetbrains.kotlin.analysis.api.fir.components

-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.analysis.api.components.KtDeclarationRendererOptions
 import org.jetbrains.kotlin.analysis.api.components.KtSymbolDeclarationRendererProvider
 import org.jetbrains.kotlin.analysis.api.components.KtTypeRendererOptions
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.analysis.api.fir.renderer.ConeTypeIdeRenderer
 import org.jetbrains.kotlin.analysis.api.fir.renderer.FirIdeRenderer
+import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirReceiverParameterSymbol
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirSymbol
 import org.jetbrains.kotlin.analysis.api.fir.types.KtFirType
-import org.jetbrains.kotlin.analysis.api.symbols.*
-import org.jetbrains.kotlin.analysis.api.symbols.markers.*
+import org.jetbrains.kotlin.analysis.api.symbols.KtPackageSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithKind
 import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
 import org.jetbrains.kotlin.analysis.api.types.KtType
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase

 internal class KtFirSymbolDeclarationRendererProvider(
    override val analysisSession: KtFirAnalysisSession,
@@ -45,6 +47,13 @@ internal class KtFirSymbolDeclarationRendererProvider(
                    FirIdeRenderer.render(fir, containingFir, options, fir.moduleData.session)
                }
            }
+            is KtFirReceiverParameterSymbol -> {
+                val containingDeclaration = with(analysisSession) {
+                    symbol.getContainingSymbol() ?: error("expects receiver parameter symbol to return the containing declaration")
+                }
+
+                "extension receiver of " + render(containingDeclaration, options)
+            }
            else -> {
                error("Unexpected Fir Symbol ${symbol::class.simpleName}")
            }
@@ -29,6 +29,7 @@ import org.jetbrains.kotlin.fir.scopes.impl.FirExplicitSimpleImportingScope
 import org.jetbrains.kotlin.fir.scopes.processClassifiersByName
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.analysis.api.fir.getCandidateSymbols
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
@@ -80,7 +81,11 @@ internal object FirReferenceResolveHelper {
        return classLikeDeclaration?.buildSymbol(symbolBuilder)
    }

-    fun FirReference.toTargetSymbol(session: FirSession, symbolBuilder: KtSymbolByFirBuilder): Collection<KtSymbol> {
+    fun FirReference.toTargetSymbol(
+        session: FirSession,
+        symbolBuilder: KtSymbolByFirBuilder,
+        isInLabelReference: Boolean = false
+    ): Collection<KtSymbol> {
        return when (this) {
            is FirBackingFieldReference -> {
                listOfNotNull(resolvedSymbol.fir.buildSymbol(symbolBuilder))
@@ -103,7 +108,12 @@ internal object FirReferenceResolveHelper {
                listOfNotNull(fir.buildSymbol(symbolBuilder))
            }
            is FirThisReference -> {
-                listOfNotNull(boundSymbol?.fir?.buildSymbol(symbolBuilder))
+                val boundSymbol = boundSymbol
+                when {
+                    !isInLabelReference && boundSymbol is FirCallableSymbol<*> ->
+                        symbolBuilder.callableBuilder.buildExtensionReceiverSymbol(boundSymbol.fir)
+                    else -> boundSymbol?.fir?.buildSymbol(symbolBuilder)
+                }.let { listOfNotNull(it) }
            }
            is FirSuperReference -> {
                listOfNotNull((superTypeRef as? FirResolvedTypeRef)?.toTargetSymbol(session, symbolBuilder))
@@ -308,7 +318,7 @@ internal object FirReferenceResolveHelper {
                    fir.dispatchReceiver as? FirQualifiedAccessExpression ?: fir.extensionReceiver as FirQualifiedAccessExpression
                receiver.calleeReference
            } else fir.calleeReference
-        return calleeReference.toTargetSymbol(session, symbolBuilder)
+        return calleeReference.toTargetSymbol(session, symbolBuilder, isInLabelReference = expression is KtLabelReferenceExpression)
    }


@@ -0,0 +1,46 @@
+/*
+ * 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.analysis.api.fir.symbols
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.analysis.api.ValidityTokenOwner
+import org.jetbrains.kotlin.analysis.api.fir.KtSymbolByFirBuilder
+import org.jetbrains.kotlin.analysis.api.fir.findPsi
+import org.jetbrains.kotlin.analysis.api.fir.utils.cached
+import org.jetbrains.kotlin.analysis.api.fir.utils.firRef
+import org.jetbrains.kotlin.analysis.api.fir.utils.weakRef
+import org.jetbrains.kotlin.analysis.api.symbols.KtReceiverParameterSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KtSymbolOrigin
+import org.jetbrains.kotlin.analysis.api.symbols.markers.KtTypeAndAnnotations
+import org.jetbrains.kotlin.analysis.api.symbols.pointers.KtSymbolPointer
+import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
+import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
+
+internal class KtFirReceiverParameterSymbol(
+    fir: FirCallableDeclaration,
+    resolveState: FirModuleResolveState,
+    override val token: ValidityToken,
+    _builder: KtSymbolByFirBuilder
+) : KtReceiverParameterSymbol(), ValidityTokenOwner {
+    init {
+        require(fir.receiverTypeRef != null) { "$fir doesn't have an extension receiver." }
+    }
+
+    private val builder by weakRef(_builder)
+    internal val firRef = firRef(fir, resolveState)
+
+    override val type: KtTypeAndAnnotations by cached {
+        firRef.receiverTypeAndAnnotations(builder) ?: throw IllegalStateException("$fir doesn't have an extension receiver.")
+    }
+    override val origin: KtSymbolOrigin get() = firRef.withFir { it.ktSymbolOrigin() }
+
+    override val psi: PsiElement? by firRef.withFirAndCache { fir -> fir.receiverTypeRef?.findPsi(fir.moduleData.session) }
+
+    override fun createPointer(): KtSymbolPointer<KtFirReceiverParameterSymbol> {
+        TODO("Probably just create a pointer based on the underlying declaration")
+    }
+}
@@ -32,7 +32,7 @@ internal fun KtFirSymbol<*>.symbolEquals(other: Any?): Boolean {

 internal fun KtFirSymbol<*>.symbolHashCode(): Int = firRef.hashCode() * 31 + token.hashCode()

-private tailrec fun FirDeclaration.ktSymbolOrigin(): KtSymbolOrigin = when (origin) {
+internal tailrec fun FirDeclaration.ktSymbolOrigin(): KtSymbolOrigin = when (origin) {
    FirDeclarationOrigin.Source -> {
        when (source?.kind) {
            FirFakeSourceElementKind.ImplicitConstructor,
@@ -492,6 +492,30 @@ public class FirReferenceResolveTestGenerated extends AbstractFirReferenceResolv
        runTest("analysis/analysis-api/testData/referenceResolve/SyntheticProperty.kt");
    }

+    @Test
+    @TestMetadata("thisWithLabelToFunction_caretAtLabel.kt")
+    public void testThisWithLabelToFunction_caretAtLabel() throws Exception {
+        runTest("analysis/analysis-api/testData/referenceResolve/thisWithLabelToFunction_caretAtLabel.kt");
+    }
+
+    @Test
+    @TestMetadata("thisWithLabelToFunction_caretAtThis.kt")
+    public void testThisWithLabelToFunction_caretAtThis() throws Exception {
+        runTest("analysis/analysis-api/testData/referenceResolve/thisWithLabelToFunction_caretAtThis.kt");
+    }
+
+    @Test
+    @TestMetadata("thisWithLabelToProperty_caretAtLabel.kt")
+    public void testThisWithLabelToProperty_caretAtLabel() throws Exception {
+        runTest("analysis/analysis-api/testData/referenceResolve/thisWithLabelToProperty_caretAtLabel.kt");
+    }
+
+    @Test
+    @TestMetadata("thisWithLabelToProperty_caretAtThis.kt")
+    public void testThisWithLabelToProperty_caretAtThis() throws Exception {
+        runTest("analysis/analysis-api/testData/referenceResolve/thisWithLabelToProperty_caretAtThis.kt");
+    }
+
    @Test
    @TestMetadata("TopLevelClassVsLocalClassConstructor.kt")
    public void testTopLevelClassVsLocalClassConstructor() throws Exception {
@@ -18,4 +18,12 @@ public abstract class KtCallableSymbol : KtSymbol, KtSymbolWithKind {
    public abstract val isExtension: Boolean

    abstract override fun createPointer(): KtSymbolPointer<KtCallableSymbol>
+}
+
+/**
+ * Symbol for a receiver parameter of a function or property. For example, consider code `fun String.foo() {...}`, the declaration of
+ * `String` receiver parameter is such a symbol.
+ */
+public abstract class KtReceiverParameterSymbol : KtSymbol {
+    public abstract val type: KtTypeAndAnnotations
 }
@@ -0,0 +1,2 @@
+Resolved to:
+0: extension receiver of fun kotlin.String.foo()
@@ -0,0 +1,5 @@
+class Foo {
+    fun String.foo() {
+        this@f<caret>oo
+    }
+}
@@ -0,0 +1,2 @@
+Resolved to:
+0: (in Foo) fun kotlin.String.foo()
@@ -0,0 +1,5 @@
+class Foo {
+    fun String.foo() {
+        t<caret>his@foo
+    }
+}
@@ -0,0 +1,2 @@
+Resolved to:
+0: extension receiver of fun kotlin.String.foo()
@@ -0,0 +1,6 @@
+class Foo {
+    val String.foo : Any
+    get() {
+        this@f<caret>oo
+    }
+}
@@ -0,0 +1,2 @@
+Resolved to:
+0: (in Foo) val kotlin.String.foo: kotlin.Any
@@ -0,0 +1,6 @@
+class Foo {
+    val String.foo : Any
+    get() {
+        t<caret>his@foo
+    }
+}
@@ -0,0 +1,2 @@
+Resolved to:
+0: extension receiver of val kotlin.String.foo: kotlin.Any