FIR: replace receiver in anon initializer with containing decl

This way the abstraction represents the intended use better and doesn't stay on the way of using FirAnonymousInitializer in scripts
2023-11-30 17:49:51 +01:00
parent 12da3d2b9f
commit ff4f30ddc1
13 changed files with 59 additions and 46 deletions
@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.api
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.nullableJavaSymbolProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirLibraryOrLibrarySourceResolvableModuleSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.FirElementFinder
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.containingClassId
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.getContainingFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isScriptDependentDeclaration
 import org.jetbrains.kotlin.analysis.utils.errors.requireIsInstance
@@ -21,10 +20,11 @@ import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticPropertyAccessor
 import org.jetbrains.kotlin.fir.declarations.utils.isLocal
 import org.jetbrains.kotlin.fir.diagnostics.ConeDestructuringDeclarationsOnTopLevel
-import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.providers.firProvider
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.resolve.toSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirScriptSymbol
 import org.jetbrains.kotlin.fir.utils.exceptions.withFirEntry
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.name.ClassId
@@ -106,8 +106,8 @@ private fun collectDesignationPath(target: FirElementWithResolveState): List<Fir
        }

        is FirAnonymousInitializer -> {
-            val containingClassId = target.containingClassId()
-            if (containingClassId.isLocal) return null
+            val containingClassId = (target.containingDeclarationSymbol as? FirClassSymbol<*>)?.classId
+            if (containingClassId == null || containingClassId.isLocal) return null
            return collectDesignationPathWithContainingClass(target, containingClassId)
        }

@@ -293,8 +293,17 @@ fun FirElementWithResolveState.tryCollectDesignationWithFile(): FirDesignationWi
 }

 private fun FirDeclaration.scriptDesignation(): FirDesignationWithFile? {
-    if (this !is FirStatement || !isScriptDependentDeclaration) return null
-    val firFile = getContainingFile() ?: return null
-    val firScript = firFile.declarations.singleOrNull() as? FirScript ?: return null
-    return FirDesignationWithFile(path = emptyList(), firScript, firFile)
+    return when {
+        this is FirAnonymousInitializer -> {
+            val firScriptSymbol = (containingDeclarationSymbol as? FirScriptSymbol) ?: return null
+            val firFile = firScriptSymbol.fir.getContainingFile() ?: return null
+            FirDesignationWithFile(path = emptyList(), firScriptSymbol.fir, firFile)
+        }
+        isScriptDependentDeclaration -> {
+            val firFile = getContainingFile() ?: return null
+            val firScript = firFile.declarations.singleOrNull() as? FirScript ?: return null
+            FirDesignationWithFile(path = emptyList(), firScript, firFile)
+        }
+        else -> null
+    }
 }
@@ -8,7 +8,7 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.fir
 import org.jetbrains.kotlin.analysis.low.level.api.fir.ContextByDesignationCollector
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignationWithFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.collectDesignation
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.containingClassId
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.containingClassIdOrNull
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContextForProvider
@@ -64,7 +64,7 @@ internal object PersistenceContextCollector {
            is FirClassLikeDeclaration -> declaration.symbol.classId.isLocal
            is FirCallableDeclaration -> declaration.symbol.callableId.isLocal
            is FirDanglingModifierList -> declaration.containingClass()?.classId?.isLocal == true
-            is FirAnonymousInitializer -> declaration.containingClassId().isLocal
+            is FirAnonymousInitializer -> declaration.containingClassIdOrNull()?.isLocal == true
            is FirScript, is FirCodeFragment -> false
            else -> errorWithAttachment("Unsupported declaration ${declaration::class}") {
                withFirEntry("declaration", declaration)
@@ -25,7 +25,7 @@ fun FirElementWithResolveState.getContainingFile(): FirFile? {
        is FirBackingField -> propertySymbol.fir.getContainingFile()
        is FirCallableDeclaration -> provider.getFirCallableContainerFile(symbol)
        is FirClassLikeDeclaration -> provider.getFirClassifierContainerFileIfAny(symbol)
-        is FirAnonymousInitializer -> provider.getFirClassifierContainerFileIfAny(containingClassId())
+        is FirAnonymousInitializer -> containingDeclarationSymbol?.fir?.getContainingFile()
        is FirDanglingModifierList, is FirCodeFragment -> {
            val ktFile = psi?.containingFile as? KtFile
                ?: error("File for dangling modifier list cannot be null")
@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.fir.utils.exceptions.withFirEntry
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.psi.*
@@ -171,17 +172,8 @@ private fun KtDeclaration.findSourceNonLocalFirDeclarationByProvider(
    return candidate?.takeIf { it.psi == this }
 }

-fun FirAnonymousInitializer.containingClassId(): ClassId {
-    val dispatchReceiverType = this.dispatchReceiverType
-    checkWithAttachment(
-        condition = dispatchReceiverType != null,
-        message = { "dispatchReceiverType for ${FirAnonymousInitializer::class.simpleName} modifier cannot be null" },
-    ) {
-        withFirEntry("initializer", this@containingClassId)
-    }
-
-    return dispatchReceiverType.lookupTag.classId
-}
+fun FirAnonymousInitializer.containingClassIdOrNull(): ClassId? =
+    (containingDeclarationSymbol as? FirClassSymbol<*>)?.classId

 val ORIGINAL_DECLARATION_KEY = com.intellij.openapi.util.Key<KtDeclaration>("ORIGINAL_DECLARATION_KEY")
 var KtDeclaration.originalDeclaration by UserDataProperty(ORIGINAL_DECLARATION_KEY)
@@ -158,7 +158,7 @@ fun FirTypeRef.toRegularClassSymbol(session: FirSession): FirRegularClassSymbol?
 fun FirBasedSymbol<*>.getContainingClassSymbol(session: FirSession): FirClassLikeSymbol<*>? = when (this) {
    is FirCallableSymbol<*> -> containingClassLookupTag()?.toSymbol(session)
    is FirClassLikeSymbol<*> -> getContainingClassLookupTag()?.toSymbol(session)
-    is FirAnonymousInitializerSymbol -> dispatchReceiverType?.toSymbol(session) as? FirClassLikeSymbol<*>
+    is FirAnonymousInitializerSymbol -> containingDeclarationSymbol as? FirClassLikeSymbol<*>
    else -> null
 }

@@ -10,8 +10,6 @@ import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.diagnostics.reportOn
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.analysis.checkers.classKind
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.checkers.getContainingClassSymbol
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
@@ -25,7 +23,6 @@ import org.jetbrains.kotlin.fir.references.toResolvedNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.SymbolInternals
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
-import org.jetbrains.kotlin.fir.types.toSymbol
 import org.jetbrains.kotlin.utils.addToStdlib.flatAssociateBy

 object FirUninitializedEnumChecker : FirQualifiedAccessExpressionChecker() {
@@ -150,7 +147,7 @@ object FirUninitializedEnumChecker : FirQualifiedAccessExpressionChecker() {
        //       fun foo() = ...
        //     }
        //   }
-        if (accessedContext in enumEntries && containingDeclarationForAccess?.isEnumEntryInitializer(context.session) != true) {
+        if (accessedContext in enumEntries && containingDeclarationForAccess?.isEnumEntryInitializer() != true) {
            return
        }

@@ -242,14 +239,14 @@ object FirUninitializedEnumChecker : FirQualifiedAccessExpressionChecker() {
            return (lazyCallArgument.expression as? FirAnonymousFunctionExpression)?.anonymousFunction
        }

-    private fun FirDeclaration.isEnumEntryInitializer(session: FirSession): Boolean {
+    private fun FirDeclaration.isEnumEntryInitializer(): Boolean {
        val containingClassSymbol = when (this) {
            is FirConstructor -> {
                if (!isPrimary) return false
                (containingClassForStaticMemberAttr as? ConeClassLookupTagWithFixedSymbol)?.symbol
            }
            is FirAnonymousInitializer -> {
-                dispatchReceiverType?.toSymbol(session)
+                containingDeclarationSymbol as? FirClassSymbol
            }
            else -> null
        } ?: return false
@@ -54,6 +54,7 @@ import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
 import org.jetbrains.kotlin.util.getChildren
 import org.jetbrains.kotlin.utils.addToStdlib.runIf
 import org.jetbrains.kotlin.utils.addToStdlib.runUnless
+import org.jetbrains.kotlin.utils.addToStdlib.shouldNotBeCalled

 class LightTreeRawFirDeclarationBuilder(
    session: FirSession,
@@ -136,7 +137,7 @@ class LightTreeRawFirDeclarationBuilder(
                KtNodeTypes.PROPERTY -> container += convertPropertyDeclaration(node) as FirStatement
                DESTRUCTURING_DECLARATION -> container += convertDestructingDeclaration(node).toFirDestructingDeclaration(baseModuleData)
                TYPEALIAS -> container += convertTypeAlias(node) as FirStatement
-                CLASS_INITIALIZER -> container += convertAnonymousInitializer(node) as FirStatement
+                CLASS_INITIALIZER -> shouldNotBeCalled("CLASS_INITIALIZER expected to be processed during class body conversion")
                else -> if (node.isExpression()) container += expressionConverter.getAsFirStatement(node)
            }
        }
@@ -891,7 +892,7 @@ class LightTreeRawFirDeclarationBuilder(
                KtNodeTypes.PROPERTY -> container += convertPropertyDeclaration(node, classWrapper)
                TYPEALIAS -> container += convertTypeAlias(node)
                OBJECT_DECLARATION -> container += convertClass(node)
-                CLASS_INITIALIZER -> container += convertAnonymousInitializer(node) //anonymousInitializer
+                CLASS_INITIALIZER -> container += convertAnonymousInitializer(node, classWrapper) //anonymousInitializer
                SECONDARY_CONSTRUCTOR -> container += convertSecondaryConstructor(node, classWrapper)
                MODIFIER_LIST -> modifierLists += node
                DESTRUCTURING_DECLARATION -> container += buildErrorTopLevelDestructuringDeclaration(node.toFirSourceElement())
@@ -1053,7 +1054,10 @@ class LightTreeRawFirDeclarationBuilder(
     * @see org.jetbrains.kotlin.parsing.KotlinParsing.parseMemberDeclarationRest
     * at INIT keyword
     */
-    private fun convertAnonymousInitializer(anonymousInitializer: LighterASTNode): FirDeclaration {
+    private fun convertAnonymousInitializer(
+        anonymousInitializer: LighterASTNode,
+        classWrapper: ClassWrapper
+    ): FirDeclaration {
        val initializerSymbol = FirAnonymousInitializerSymbol()
        withContainerSymbol(initializerSymbol) {
            var firBlock: FirBlock? = null
@@ -1073,7 +1077,7 @@ class LightTreeRawFirDeclarationBuilder(
                moduleData = baseModuleData
                origin = FirDeclarationOrigin.Source
                body = firBlock ?: buildEmptyExpressionBlock()
-                dispatchReceiverType = context.dispatchReceiverTypesStack.lastOrNull()
+                containingDeclarationSymbol = classWrapper.classBuilder.ownerRegularOrAnonymousObjectSymbol
                annotations += initializerAnnotations
            }
        }
@@ -400,6 +400,9 @@ open class PsiRawFirBuilder(
                is KtDestructuringDeclaration -> {
                    buildErrorTopLevelDestructuringDeclaration(toFirSourceElement())
                }
+                is KtClassInitializer -> {
+                    buildAnonymousInitializer(this, ownerClassBuilder.ownerRegularOrAnonymousObjectSymbol)
+                }
                else -> convert()
            }
        }
@@ -1310,6 +1313,7 @@ open class PsiRawFirBuilder(
                                        source = declarationSource
                                        body = firBlock
                                        declaration.extractAnnotationsTo(this)
+                                        containingDeclarationSymbol = scriptSymbol
                                    }
                                )
                            }
@@ -2209,7 +2213,11 @@ open class PsiRawFirBuilder(
        }

        override fun visitAnonymousInitializer(initializer: KtAnonymousInitializer, data: FirElement?): FirElement {
-            return buildAnonymousInitializer {
+            return buildAnonymousInitializer(initializer, null)
+        }
+
+        private fun buildAnonymousInitializer(initializer: KtAnonymousInitializer, containingDeclarationSymbol: FirBasedSymbol<*>?) =
+            buildAnonymousInitializer {
                withContainerSymbol(symbol) {
                    source = initializer.toFirSourceElement()
                    moduleData = baseModuleData
@@ -2219,11 +2227,10 @@ open class PsiRawFirBuilder(
                            initializer.body.toFirBlock()
                        }
                    }
-                    dispatchReceiverType = context.dispatchReceiverTypesStack.lastOrNull()
+                    this.containingDeclarationSymbol = containingDeclarationSymbol
                    initializer.extractAnnotationsTo(this)
                }
            }
-        }

        override fun visitProperty(property: KtProperty, data: FirElement?): FirElement {
            return property.toFirProperty(
@@ -14,8 +14,8 @@ import org.jetbrains.kotlin.fir.FirModuleData
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirBlock
 import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousInitializerSymbol
-import org.jetbrains.kotlin.fir.types.ConeClassLikeType
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor

@@ -31,7 +31,7 @@ abstract class FirAnonymousInitializer : FirDeclaration(), FirControlFlowGraphOw
    abstract override val controlFlowGraphReference: FirControlFlowGraphReference?
    abstract val body: FirBlock?
    abstract override val symbol: FirAnonymousInitializerSymbol
-    abstract val dispatchReceiverType: ConeClassLikeType?
+    abstract val containingDeclarationSymbol: FirBasedSymbol<*>?

    override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R =
        visitor.visitAnonymousInitializer(this, data)
@@ -23,8 +23,8 @@ import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.impl.FirAnonymousInitializerImpl
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousInitializerSymbol
-import org.jetbrains.kotlin.fir.types.ConeClassLikeType

@FirBuilderDsl
 class FirAnonymousInitializerBuilder : FirDeclarationBuilder, FirAnnotationContainerBuilder {
@@ -36,7 +36,7 @@ class FirAnonymousInitializerBuilder : FirDeclarationBuilder, FirAnnotationConta
    override var attributes: FirDeclarationAttributes = FirDeclarationAttributes()
    var body: FirBlock? = null
    var symbol: FirAnonymousInitializerSymbol = FirAnonymousInitializerSymbol()
-    var dispatchReceiverType: ConeClassLikeType? = null
+    var containingDeclarationSymbol: FirBasedSymbol<*>? = null

    override fun build(): FirAnonymousInitializer {
        return FirAnonymousInitializerImpl(
@@ -48,7 +48,7 @@ class FirAnonymousInitializerBuilder : FirDeclarationBuilder, FirAnnotationConta
            attributes,
            body,
            symbol,
-            dispatchReceiverType,
+            containingDeclarationSymbol,
        )
    }

@@ -18,8 +18,8 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirBlock
 import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousInitializerSymbol
-import org.jetbrains.kotlin.fir.types.ConeClassLikeType
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 import org.jetbrains.kotlin.fir.visitors.transformInplace
@@ -34,7 +34,7 @@ internal class FirAnonymousInitializerImpl(
    override val attributes: FirDeclarationAttributes,
    override var body: FirBlock?,
    override val symbol: FirAnonymousInitializerSymbol,
-    override val dispatchReceiverType: ConeClassLikeType?,
+    override val containingDeclarationSymbol: FirBasedSymbol<*>?,
 ) : FirAnonymousInitializer() {
    override var controlFlowGraphReference: FirControlFlowGraphReference? = null

@@ -12,6 +12,6 @@ import org.jetbrains.kotlin.fir.types.ConeSimpleKotlinType
 class FirAnonymousInitializerSymbol : FirBasedSymbol<FirAnonymousInitializer>() {
    override fun toString(): String = "${this::class.simpleName} <init>"

-    val dispatchReceiverType: ConeSimpleKotlinType?
-        get() = fir.dispatchReceiverType
+    val containingDeclarationSymbol: FirBasedSymbol<*>?
+        get() = fir.containingDeclarationSymbol
 }
@@ -474,7 +474,11 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
        anonymousInitializer.configure {
            +body(nullable = true, withReplace = true)
            +symbol("FirAnonymousInitializerSymbol")
-            +field("dispatchReceiverType", coneClassLikeTypeType, nullable = true)
+            // the containing declaration is nullable, because it is not immediately clear how to obtain it in all places in the fir builder
+            // TODO: review and consider making not-nullable (KT-64195)
+            +field("containingDeclarationSymbol", firBasedSymbolType.withArgs(TypeRef.Star), nullable = true).apply {
+                withBindThis = false
+            }
        }

        danglingModifierList.configure {