[Analysis API FIR] Initialize properties of KtCallableSignature lazily

2023-04-13 14:01:57 +02:00
parent 2eab9fc00f
commit 2cc3b542fd
17 changed files with 496 additions and 245 deletions
@@ -62,6 +62,8 @@ import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.types.Variance
 import kotlin.contracts.ExperimentalContracts
 import kotlin.contracts.contract
+import org.jetbrains.kotlin.analysis.api.fir.signatures.KtFirFunctionLikeSubstitutorBasedSignature
+import org.jetbrains.kotlin.analysis.api.fir.signatures.KtFirVariableLikeSubstitutorBasedSignature

 /**
 * Maps FirElement to KtSymbol & ConeType to KtType, thread safe
@@ -219,19 +221,7 @@ internal class KtSymbolByFirBuilder constructor(

        fun buildFunctionSignature(firSymbol: FirNamedFunctionSymbol): KtFunctionLikeSignature<KtFirFunctionSymbol> {
            firSymbol.lazyResolveToPhase(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE)
-            val functionSymbol = buildFunctionSymbol(firSymbol)
-            return KtFunctionLikeSignature(
-                functionSymbol,
-                typeBuilder.buildKtType(firSymbol.resolvedReturnType),
-                firSymbol.resolvedReceiverTypeRef?.let { typeBuilder.buildKtType(it) },
-                functionSymbol.valueParameters.zip(firSymbol.fir.valueParameters).map { (ktSymbol, fir) ->
-                    var type = fir.returnTypeRef.coneType
-                    if (fir.isVararg) {
-                        type = type.arrayElementType() ?: type
-                    }
-                    KtVariableLikeSignature(ktSymbol, typeBuilder.buildKtType(type), null)
-                }
-            )
+            return KtFirFunctionLikeSubstitutorBasedSignature(analysisSession.token, firSymbol, analysisSession.firSymbolBuilder)
        }

        fun buildAnonymousFunctionSymbol(firSymbol: FirAnonymousFunctionSymbol): KtFirAnonymousFunctionSymbol {
@@ -311,11 +301,7 @@ internal class KtSymbolByFirBuilder constructor(

        fun buildPropertySignature(firSymbol: FirPropertySymbol): KtVariableLikeSignature<KtVariableSymbol> {
            firSymbol.lazyResolveToPhase(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE)
-            return KtVariableLikeSignature(
-                buildPropertySymbol(firSymbol),
-                typeBuilder.buildKtType(firSymbol.fir.returnTypeRef),
-                firSymbol.resolvedReceiverTypeRef?.let { typeBuilder.buildKtType(it) }
-            )
+            return KtFirVariableLikeSubstitutorBasedSignature(analysisSession.token, firSymbol, analysisSession.firSymbolBuilder)
        }

        fun buildLocalVariableSymbol(firSymbol: FirPropertySymbol): KtFirLocalVariableSymbol {
@@ -332,6 +318,9 @@ internal class KtSymbolByFirBuilder constructor(
        }

        fun buildValueParameterSymbol(firSymbol: FirValueParameterSymbol): KtValueParameterSymbol {
+            firSymbol.fir.unwrapSubstitutionOverrideIfNeeded()?.let {
+                return buildValueParameterSymbol(it.symbol)
+            }
            return symbolsCache.cache(firSymbol) {
                KtFirValueParameterSymbol(firSymbol, analysisSession)
            }
@@ -529,8 +518,8 @@ internal class KtSymbolByFirBuilder constructor(
    private inline fun <reified T : FirCallableDeclaration> T.unwrapUseSiteSubstitutionOverride(): T? {
        val originalDeclaration = originalForSubstitutionOverride ?: return null

-        val containingClass = getContainingClass(rootSession) ?: return null
-        val originalContainingClass = originalDeclaration.getContainingClass(rootSession) ?: return null
+        val containingClass = getContainingMemberOrSelf().getContainingClass(rootSession) ?: return null
+        val originalContainingClass = originalDeclaration.getContainingMemberOrSelf().getContainingClass(rootSession) ?: return null

        // If substitution override does not change the containing class of the FIR declaration,
        // it is a use-site substitution override
@@ -539,6 +528,11 @@ internal class KtSymbolByFirBuilder constructor(
        return originalDeclaration
    }

+    private fun FirCallableDeclaration.getContainingMemberOrSelf(): FirCallableDeclaration = when (this) {
+        is FirValueParameter -> containingFunctionSymbol.fir
+        else -> this
+    }
+
    /**
     * We want to unwrap a SUBSTITUTION_OVERRIDE wrapper if it doesn't affect the declaration's signature in any way. If the signature
     * is somehow changed, then we want to keep the wrapper.
@@ -6,8 +6,27 @@
 package org.jetbrains.kotlin.analysis.api.fir.components

 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
-import org.jetbrains.kotlin.analysis.api.impl.base.components.AbstractKtSignatureSubstitutorImpl
+import org.jetbrains.kotlin.analysis.api.fir.signatures.KtFirFunctionLikeDummySignature
+import org.jetbrains.kotlin.analysis.api.fir.signatures.KtFirVariableLikeDummySignature
+import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirSymbol
+import org.jetbrains.kotlin.analysis.api.impl.base.components.AbstractKtSignatureSubstitutor
+import org.jetbrains.kotlin.analysis.api.signatures.KtFunctionLikeSignature
+import org.jetbrains.kotlin.analysis.api.signatures.KtVariableLikeSignature
+import org.jetbrains.kotlin.analysis.api.symbols.KtFunctionLikeSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KtVariableLikeSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol

 internal class KtFirSignatureSubstitutor(
    override val analysisSession: KtFirAnalysisSession
-) : AbstractKtSignatureSubstitutorImpl(), KtFirAnalysisSessionComponent
+) : AbstractKtSignatureSubstitutor(), KtFirAnalysisSessionComponent {
+    override fun <S : KtFunctionLikeSymbol> asSignature(symbol: S): KtFunctionLikeSignature<S> {
+        val firSymbol = (symbol as KtFirSymbol<*>).firSymbol as FirFunctionSymbol<*>
+        return KtFirFunctionLikeDummySignature<S>(analysisSession.token, firSymbol, analysisSession.firSymbolBuilder)
+    }
+
+    override fun <S : KtVariableLikeSymbol> asSignature(symbol: S): KtVariableLikeSignature<S> {
+        val firSymbol = (symbol as KtFirSymbol<*>).firSymbol as FirVariableSymbol<*>
+        return KtFirVariableLikeDummySignature<S>(analysisSession.token, firSymbol, analysisSession.firSymbolBuilder)
+    }
+}
@@ -0,0 +1,14 @@
+/*
+ * Copyright 2010-2023 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.signatures
+
+import org.jetbrains.kotlin.analysis.api.fir.KtSymbolByFirBuilder
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+
+internal interface FirSymbolBasedSignature {
+    val firSymbol: FirCallableSymbol<*>
+    val firSymbolBuilder: KtSymbolByFirBuilder
+}
@@ -0,0 +1,106 @@
+/*
+ * Copyright 2010-2023 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.signatures
+
+import org.jetbrains.kotlin.analysis.api.fir.KtSymbolByFirBuilder
+import org.jetbrains.kotlin.analysis.api.fir.buildSymbol
+import org.jetbrains.kotlin.analysis.api.fir.types.AbstractKtFirSubstitutor
+import org.jetbrains.kotlin.analysis.api.fir.utils.cached
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
+import org.jetbrains.kotlin.analysis.api.signatures.KtFunctionLikeSignature
+import org.jetbrains.kotlin.analysis.api.signatures.KtVariableLikeSignature
+import org.jetbrains.kotlin.analysis.api.symbols.KtValueParameterSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KtFunctionLikeSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.receiverType
+import org.jetbrains.kotlin.analysis.api.types.KtSubstitutor
+import org.jetbrains.kotlin.analysis.api.types.KtType
+import org.jetbrains.kotlin.fir.declarations.FirPropertyAccessor
+import org.jetbrains.kotlin.fir.resolve.substitution.ChainedSubstitutor
+import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
+import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirValueParameterSymbol
+
+internal sealed class KtFirFunctionLikeSignature<out S : KtFunctionLikeSymbol> : KtFunctionLikeSignature<S>(), FirSymbolBasedSignature {
+    abstract override fun substitute(substitutor: KtSubstitutor): KtFirFunctionLikeSignature<S>
+
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (javaClass != other?.javaClass) return false
+
+        other as KtFirFunctionLikeSignature<*>
+        return firSymbol == other.firSymbol
+    }
+
+    override fun hashCode(): Int = firSymbol.hashCode()
+}
+
+internal class KtFirFunctionLikeDummySignature<out S : KtFunctionLikeSymbol>(
+    override val token: KtLifetimeToken,
+    override val firSymbol: FirFunctionSymbol<*>,
+    override val firSymbolBuilder: KtSymbolByFirBuilder,
+) : KtFirFunctionLikeSignature<S>() {
+    @Suppress("UNCHECKED_CAST")
+    override val symbol: S
+        get() = withValidityAssertion { firSymbol.buildSymbol(firSymbolBuilder) as S }
+    override val returnType: KtType
+        get() = withValidityAssertion { symbol.returnType }
+    override val receiverType: KtType?
+        get() = withValidityAssertion { symbol.receiverType }
+    override val valueParameters: List<KtVariableLikeSignature<KtValueParameterSymbol>> by cached {
+        firSymbol.valueParameterSymbols.map { KtFirVariableLikeDummySignature(token, it, firSymbolBuilder) }
+    }
+
+    override fun substitute(substitutor: KtSubstitutor): KtFirFunctionLikeSignature<S> = withValidityAssertion {
+        if (substitutor is KtSubstitutor.Empty) return@withValidityAssertion this
+        require(substitutor is AbstractKtFirSubstitutor<*>)
+
+        KtFirFunctionLikeSubstitutorBasedSignature(token, firSymbol, firSymbolBuilder, substitutor.substitutor)
+    }
+}
+
+internal class KtFirFunctionLikeSubstitutorBasedSignature<out S : KtFunctionLikeSymbol>(
+    override val token: KtLifetimeToken,
+    override val firSymbol: FirFunctionSymbol<*>,
+    override val firSymbolBuilder: KtSymbolByFirBuilder,
+    override val coneSubstitutor: ConeSubstitutor = ConeSubstitutor.Empty,
+) : KtFirFunctionLikeSignature<S>(), SubstitutorBasedSignature {
+    @Suppress("UNCHECKED_CAST")
+    override val symbol: S
+        get() = withValidityAssertion { firSymbol.buildSymbol(firSymbolBuilder) as S }
+    override val returnType: KtType by cached {
+        firSymbolBuilder.typeBuilder.buildKtType(coneSubstitutor.substituteOrSelf(firSymbol.resolvedReturnType))
+    }
+    override val receiverType: KtType? by cached {
+        val receiverTypeRef = when (val fir = firSymbol.fir) {
+            is FirPropertyAccessor -> fir.propertySymbol.resolvedReceiverTypeRef
+            else -> firSymbol.resolvedReceiverTypeRef
+        }
+        receiverTypeRef?.let { firSymbolBuilder.typeBuilder.buildKtType(coneSubstitutor.substituteOrSelf(it.type)) }
+    }
+    override val valueParameters: List<KtVariableLikeSignature<KtValueParameterSymbol>> by cached {
+        firSymbol.fir.valueParameters.map { firValueParameter ->
+            KtFirVariableLikeSubstitutorBasedSignature(token, firValueParameter.symbol, firSymbolBuilder, coneSubstitutor)
+        }
+    }
+
+    override fun substitute(substitutor: KtSubstitutor): KtFirFunctionLikeSignature<S> = withValidityAssertion {
+        if (substitutor is KtSubstitutor.Empty) return@withValidityAssertion this
+        require(substitutor is AbstractKtFirSubstitutor<*>)
+        val chainedSubstitutor = ChainedSubstitutor(coneSubstitutor, substitutor.substitutor)
+
+        KtFirFunctionLikeSubstitutorBasedSignature(token, firSymbol, firSymbolBuilder, chainedSubstitutor)
+    }
+
+    override fun equals(other: Any?): Boolean {
+        if (!super.equals(other)) return false
+
+        other as KtFirFunctionLikeSubstitutorBasedSignature<*>
+        return coneSubstitutor == other.coneSubstitutor
+    }
+
+    override fun hashCode(): Int = 31 * super.hashCode() + coneSubstitutor.hashCode()
+}
@@ -0,0 +1,99 @@
+/*
+ * Copyright 2010-2023 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.signatures
+
+import org.jetbrains.kotlin.analysis.api.fir.KtSymbolByFirBuilder
+import org.jetbrains.kotlin.analysis.api.fir.buildSymbol
+import org.jetbrains.kotlin.analysis.api.fir.types.AbstractKtFirSubstitutor
+import org.jetbrains.kotlin.analysis.api.fir.utils.cached
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
+import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
+import org.jetbrains.kotlin.analysis.api.signatures.KtVariableLikeSignature
+import org.jetbrains.kotlin.analysis.api.symbols.KtVariableLikeSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.receiverType
+import org.jetbrains.kotlin.analysis.api.types.KtSubstitutor
+import org.jetbrains.kotlin.analysis.api.types.KtType
+import org.jetbrains.kotlin.fir.resolve.substitution.ChainedSubstitutor
+import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
+import org.jetbrains.kotlin.fir.symbols.impl.FirValueParameterSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
+import org.jetbrains.kotlin.fir.types.arrayElementType
+import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.utils.addToStdlib.applyIf
+
+internal sealed class KtFirVariableLikeSignature<out S : KtVariableLikeSymbol> : KtVariableLikeSignature<S>(), FirSymbolBasedSignature {
+    abstract override fun substitute(substitutor: KtSubstitutor): KtFirVariableLikeSignature<S>
+
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (javaClass != other?.javaClass) return false
+
+        other as KtFirVariableLikeSignature<*>
+        return firSymbol == other.firSymbol
+    }
+
+    override fun hashCode(): Int = firSymbol.hashCode()
+}
+
+internal class KtFirVariableLikeDummySignature<out S : KtVariableLikeSymbol>(
+    override val token: KtLifetimeToken,
+    override val firSymbol: FirVariableSymbol<*>,
+    override val firSymbolBuilder: KtSymbolByFirBuilder,
+) : KtFirVariableLikeSignature<S>() {
+    @Suppress("UNCHECKED_CAST")
+    override val symbol: S
+        get() = withValidityAssertion { firSymbol.buildSymbol(firSymbolBuilder) as S }
+    override val returnType: KtType
+        get() = withValidityAssertion { symbol.returnType }
+    override val receiverType: KtType?
+        get() = withValidityAssertion { symbol.receiverType }
+
+    override fun substitute(substitutor: KtSubstitutor): KtFirVariableLikeSignature<S> = withValidityAssertion {
+        if (substitutor is KtSubstitutor.Empty) return@withValidityAssertion this
+        require(substitutor is AbstractKtFirSubstitutor<*>)
+
+        KtFirVariableLikeSubstitutorBasedSignature(token, firSymbol, firSymbolBuilder, substitutor.substitutor)
+    }
+}
+
+internal class KtFirVariableLikeSubstitutorBasedSignature<out S : KtVariableLikeSymbol>(
+    override val token: KtLifetimeToken,
+    override val firSymbol: FirVariableSymbol<*>,
+    override val firSymbolBuilder: KtSymbolByFirBuilder,
+    override val coneSubstitutor: ConeSubstitutor = ConeSubstitutor.Empty,
+) : KtFirVariableLikeSignature<S>(), SubstitutorBasedSignature {
+    @Suppress("UNCHECKED_CAST")
+    override val symbol: S
+        get() = withValidityAssertion { firSymbol.buildSymbol(firSymbolBuilder) as S }
+    override val returnType: KtType by cached {
+        val isVarargValueParameter = (firSymbol as? FirValueParameterSymbol)?.isVararg == true
+        val coneType = firSymbol.resolvedReturnType.applyIf(isVarargValueParameter) { arrayElementType() ?: this }
+
+        firSymbolBuilder.typeBuilder.buildKtType(coneSubstitutor.substituteOrSelf(coneType))
+    }
+    override val receiverType: KtType? by cached {
+        firSymbol.resolvedReceiverTypeRef?.let { typeRef ->
+            firSymbolBuilder.typeBuilder.buildKtType(coneSubstitutor.substituteOrSelf(typeRef.coneType))
+        }
+    }
+
+    override fun substitute(substitutor: KtSubstitutor): KtFirVariableLikeSignature<S> = withValidityAssertion {
+        if (substitutor is KtSubstitutor.Empty) return@withValidityAssertion this
+        require(substitutor is AbstractKtFirSubstitutor<*>)
+        val chainedSubstitutor = ChainedSubstitutor(coneSubstitutor, substitutor.substitutor)
+
+        KtFirVariableLikeSubstitutorBasedSignature(token, firSymbol, firSymbolBuilder, chainedSubstitutor)
+    }
+
+    override fun equals(other: Any?): Boolean {
+        if (!super.equals(other)) return false
+
+        other as KtFirVariableLikeSubstitutorBasedSignature<*>
+        return coneSubstitutor == other.coneSubstitutor
+    }
+
+    override fun hashCode(): Int = 31 * super.hashCode() + coneSubstitutor.hashCode()
+}
@@ -0,0 +1,12 @@
+/*
+ * Copyright 2010-2023 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.signatures
+
+import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
+
+internal interface SubstitutorBasedSignature {
+    val coneSubstitutor: ConeSubstitutor
+}