FIR: Unify all references to FIR nodes from non-parents

compiler/fir/dump/src/org/jetbrains/kotlin/fir/dump/HtmlFirDump.kt

@@ -1587,7 +1587,7 @@ class HtmlFirDump internal constructor(private var linkResolver: FirLinkResolver
         +"?."
 
         val nestedQualifier = safeCallExpression.regularQualifiedAccess
-        if (nestedQualifier.explicitReceiver == safeCallExpression.checkedSubject.value) {
+        if (nestedQualifier.explicitReceiver == safeCallExpression.checkedSubjectRef.value) {
             when (nestedQualifier) {
                 is FirFunctionCall -> {
                     return generate(nestedQualifier, skipReceiver = true)

compiler/fir/raw-fir/fir-common/src/org/jetbrains/kotlin/fir/builder/BaseFirBuilder.kt

@@ -562,7 +562,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
             source = baseSource
             rValue = value
             calleeReference = nestedAccess.calleeReference
-            explicitReceiver = safeCallNonAssignment.checkedSubject.value
+            explicitReceiver = safeCallNonAssignment.checkedSubjectRef.value
         }
 
         safeCallNonAssignment.replaceRegularQualifiedAccess(

compiler/fir/raw-fir/fir-common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt

@@ -147,10 +147,12 @@ fun FirExpression.generateNotNullOrOther(
 ): FirWhenExpression {
     val subjectName = Name.special("<$caseId>")
     val subjectVariable = generateTemporaryVariable(session, baseSource, subjectName, this)
-    val subject = FirWhenSubject()
+
+    @OptIn(FirContractViolation::class)
+    val ref = FirExpressionRef<FirWhenExpression>()
     val subjectExpression = buildWhenSubjectExpression {
         source = baseSource
-        whenSubject = subject
+        whenRef = ref
     }
 
     return buildWhenExpression {
@@ -176,7 +178,7 @@ fun FirExpression.generateNotNullOrOther(
             result = buildSingleExpressionBlock(generateResolvedAccessExpression(baseSource, subjectVariable))
         }
     }.also {
-        subject.bind(it)
+        ref.bind(it)
     }
 }
 
@@ -487,6 +489,7 @@ private fun FirExpression.checkReceiver(name: String?): Boolean {
     return receiverName == name
 }
 
+
 fun FirModifiableQualifiedAccess.wrapWithSafeCall(receiver: FirExpression): FirSafeCallExpression {
     // TODO: Refactor tree to make FirModifiableQualifiedAccess inherit FirQualifiedAccess
     require(this is FirQualifiedAccess) {
@@ -494,13 +497,19 @@ fun FirModifiableQualifiedAccess.wrapWithSafeCall(receiver: FirExpression): FirS
     }
 
     val checkedSafeCallSubject = buildCheckedSafeCallSubject {
-        this.originalReceiverReference = FirSafeCallOriginalReceiverReference(receiver)
+        @OptIn(FirContractViolation::class)
+        this.originalReceiverRef = FirExpressionRef<FirExpression>().apply {
+            bind(receiver)
+        }
     }
 
     explicitReceiver = checkedSafeCallSubject
     return buildSafeCallExpression {
         this.receiver = receiver
-        this.checkedSubject = FirSafeCallCheckedSubjectReference(checkedSafeCallSubject)
+        @OptIn(FirContractViolation::class)
+        this.checkedSubjectRef = FirExpressionRef<FirCheckedSafeCallSubject>().apply {
+            bind(checkedSafeCallSubject)
+        }
         this.regularQualifiedAccess = this@wrapWithSafeCall
         this.source = this@wrapWithSafeCall.source
     }

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt

@@ -636,7 +636,9 @@ class ExpressionsConverter(
         }
         subjectExpression = subjectVariable?.initializer ?: subjectExpression
         val hasSubject = subjectExpression != null
-        val subject = FirWhenSubject()
+
+        @OptIn(FirContractViolation::class)
+        val subject = FirExpressionRef<FirWhenExpression>()
         whenEntryNodes.mapTo(whenEntries) { convertWhenEntry(it, subject.takeIf { hasSubject }) }
         return buildWhenExpression {
             source = whenExpression.toFirSourceElement()
@@ -676,15 +678,15 @@ class ExpressionsConverter(
      * @see org.jetbrains.kotlin.parsing.KotlinExpressionParsing.parseWhenEntry
      * @see org.jetbrains.kotlin.parsing.KotlinExpressionParsing.parseWhenEntryNotElse
      */
-    private fun convertWhenEntry(whenEntry: LighterASTNode, subject: FirWhenSubject?): WhenEntry {
+    private fun convertWhenEntry(whenEntry: LighterASTNode, whenRefWithSubject: FirExpressionRef<FirWhenExpression>?): WhenEntry {
         var isElse = false
         var firBlock: FirBlock = buildEmptyExpressionBlock()
         val conditions = mutableListOf<FirExpression>()
         whenEntry.forEachChildren {
             when (it.tokenType) {
-                WHEN_CONDITION_EXPRESSION -> conditions += convertWhenConditionExpression(it, subject)
+                WHEN_CONDITION_EXPRESSION -> conditions += convertWhenConditionExpression(it, whenRefWithSubject)
-                WHEN_CONDITION_IN_RANGE -> conditions += convertWhenConditionInRange(it, subject)
+                WHEN_CONDITION_IN_RANGE -> conditions += convertWhenConditionInRange(it, whenRefWithSubject)
-                WHEN_CONDITION_IS_PATTERN -> conditions += convertWhenConditionIsPattern(it, subject)
+                WHEN_CONDITION_IS_PATTERN -> conditions += convertWhenConditionIsPattern(it, whenRefWithSubject)
                 ELSE_KEYWORD -> isElse = true
                 BLOCK -> firBlock = declarationsConverter.convertBlock(it)
                 else -> if (it.isExpression()) firBlock = declarationsConverter.convertBlock(it)
@@ -694,20 +696,20 @@ class ExpressionsConverter(
         return WhenEntry(conditions, firBlock, isElse)
     }
 
-    private fun convertWhenConditionExpression(whenCondition: LighterASTNode, subject: FirWhenSubject?): FirExpression {
+    private fun convertWhenConditionExpression(whenCondition: LighterASTNode, whenRefWithSubject: FirExpressionRef<FirWhenExpression>?): FirExpression {
         var firExpression: FirExpression = buildErrorExpression(null, ConeSimpleDiagnostic("No expression in condition with expression", DiagnosticKind.Syntax))
         whenCondition.forEachChildren {
             when (it.tokenType) {
                 else -> if (it.isExpression()) firExpression = getAsFirExpression(it, "No expression in condition with expression")
             }
         }
-        return if (subject != null) {
+        return if (whenRefWithSubject != null) {
             buildOperatorCall {
                 source = whenCondition.toFirSourceElement()
                 operation = FirOperation.EQ
                 argumentList = buildBinaryArgumentList(
                     buildWhenSubjectExpression {
-                        whenSubject = subject
+                        whenRef = whenRefWithSubject
                     }, firExpression
                 )
             }
@@ -717,7 +719,7 @@ class ExpressionsConverter(
         }
     }
 
-    private fun convertWhenConditionInRange(whenCondition: LighterASTNode, subject: FirWhenSubject?): FirExpression {
+    private fun convertWhenConditionInRange(whenCondition: LighterASTNode, whenRefWithSubject: FirExpressionRef<FirWhenExpression>?): FirExpression {
         var isNegate = false
         var firExpression: FirExpression = buildErrorExpression(null, ConeSimpleDiagnostic("No range in condition with range", DiagnosticKind.Syntax))
         var conditionSource: FirLightSourceElement? = null
@@ -731,9 +733,9 @@ class ExpressionsConverter(
             }
         }
 
-        val subjectExpression = if (subject != null) {
+        val subjectExpression = if (whenRefWithSubject != null) {
             buildWhenSubjectExpression {
-                whenSubject = subject
+                whenRef = whenRefWithSubject
             }
         } else {
             return buildErrorExpression {
@@ -750,7 +752,7 @@ class ExpressionsConverter(
         )
     }
 
-    private fun convertWhenConditionIsPattern(whenCondition: LighterASTNode, subject: FirWhenSubject?): FirExpression {
+    private fun convertWhenConditionIsPattern(whenCondition: LighterASTNode, whenRefWithSubject: FirExpressionRef<FirWhenExpression>?): FirExpression {
         lateinit var firOperation: FirOperation
         lateinit var firType: FirTypeRef
         whenCondition.forEachChildren {
@@ -761,9 +763,9 @@ class ExpressionsConverter(
             }
         }
 
-        val subjectExpression = if (subject != null) {
+        val subjectExpression = if (whenRefWithSubject != null) {
             buildWhenSubjectExpression {
-                whenSubject = subject
+                whenRef = whenRefWithSubject
             }
         } else {
             return buildErrorExpression {

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/PsiConversionUtils.kt

@@ -7,9 +7,9 @@ package org.jetbrains.kotlin.fir.builder
 
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.fir.FirExpressionRef
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSourceElement
-import org.jetbrains.kotlin.fir.FirWhenSubject
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
 import org.jetbrains.kotlin.fir.declarations.FirVariable
 import org.jetbrains.kotlin.fir.declarations.builder.buildProperty
@@ -24,14 +24,14 @@ import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.psi.*
 
 internal fun KtWhenCondition.toFirWhenCondition(
-    subject: FirWhenSubject,
+    whenRefWithSibject: FirExpressionRef<FirWhenExpression>,
     convert: KtExpression?.(String) -> FirExpression,
     toFirOrErrorTypeRef: KtTypeReference?.() -> FirTypeRef,
 ): FirExpression {
     val baseSource = this.toFirPsiSourceElement()
     val firSubjectExpression = buildWhenSubjectExpression {
         source = baseSource
-        whenSubject = subject
+        whenRef = whenRefWithSibject
     }
     return when (this) {
         is KtWhenConditionWithExpression -> {
@@ -68,7 +68,7 @@ internal fun KtWhenCondition.toFirWhenCondition(
 
 internal fun Array<KtWhenCondition>.toFirWhenCondition(
     baseSource: FirSourceElement?,
-    subject: FirWhenSubject,
+    subject: FirExpressionRef<FirWhenExpression>,
     convert: KtExpression?.(String) -> FirExpression,
     toFirOrErrorTypeRef: KtTypeReference?.() -> FirTypeRef,
 ): FirExpression {

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -1343,7 +1343,8 @@ class RawFirBuilder(
                 else -> null
             }
             val hasSubject = subjectExpression != null
-            val subject = FirWhenSubject()
+            @OptIn(FirContractViolation::class)
+            val ref = FirExpressionRef<FirWhenExpression>()
             return buildWhenExpression {
                 source = expression.toFirSourceElement()
                 this.subject = subjectExpression
@@ -1358,7 +1359,7 @@ class RawFirBuilder(
                                 source = entrySource
                                 condition = entry.conditions.toFirWhenCondition(
                                     entrySource,
-                                    subject,
+                                    ref,
                                     { toFirExpression(it) },
                                     { toFirOrErrorType() },
                                 )
@@ -1382,7 +1383,7 @@ class RawFirBuilder(
                 }
             }.also {
                 if (hasSubject) {
-                    subject.bind(it)
+                    ref.bind(it)
                 }
             }
         }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/VariableStorage.kt

@@ -43,10 +43,10 @@ class VariableStorage(private val session: FirSession) {
     }
 
     private fun FirElement.unwrapElement(): FirElement = when (this) {
-        is FirWhenSubjectExpression -> whenSubject.whenExpression.let { it.subjectVariable ?: it.subject }?.unwrapElement() ?: this
+        is FirWhenSubjectExpression -> whenRef.value.let { it.subjectVariable ?: it.subject }?.unwrapElement() ?: this
         is FirExpressionWithSmartcast -> originalExpression.unwrapElement()
         is FirSafeCallExpression -> regularQualifiedAccess.unwrapElement()
-        is FirCheckedSafeCallSubject -> originalReceiverReference.value.unwrapElement()
+        is FirCheckedSafeCallSubject -> originalReceiverRef.value.unwrapElement()
         else -> this
     }
 
@@ -71,7 +71,7 @@ class VariableStorage(private val session: FirSession) {
         val isThisReference: Boolean
         val expression: FirQualifiedAccess? = when (originalFir) {
             is FirQualifiedAccessExpression -> originalFir
-            is FirWhenSubjectExpression -> originalFir.whenSubject.whenExpression.subject as? FirQualifiedAccessExpression
+            is FirWhenSubjectExpression -> originalFir.whenRef.value.subject as? FirQualifiedAccessExpression
             is FirVariableAssignment -> originalFir
             else -> null
         }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/util.kt

@@ -122,7 +122,7 @@ internal val FirElement.symbol: AbstractFirBasedSymbol<*>?
     get() = when (this) {
         is FirResolvable -> symbol
         is FirSymbolOwner<*> -> symbol
-        is FirWhenSubjectExpression -> whenSubject.whenExpression.subject?.symbol
+        is FirWhenSubjectExpression -> whenRef.value.subject?.symbol
         is FirSafeCallExpression -> regularQualifiedAccess.symbol
         else -> null
     }?.takeIf { this is FirThisReceiverExpression || (it !is FirFunctionSymbol<*> && it !is FirAccessorSymbol) }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/IntegerLiteralTypeApproximationTransformer.kt

@@ -136,7 +136,7 @@ class IntegerLiteralTypeApproximationTransformer(
         data: ConeKotlinType?
     ): CompositeTransformResult<FirStatement> {
         val newReceiver =
-            checkedSafeCallSubject.originalReceiverReference.value.transform<FirExpression, ConeKotlinType?>(this, data).single
+            checkedSafeCallSubject.originalReceiverRef.value.transform<FirExpression, ConeKotlinType?>(this, data).single
         checkedSafeCallSubject.propagateTypeFromOriginalReceiver(newReceiver, session)
         return super.transformCheckedSafeCallSubject(checkedSafeCallSubject, data)
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirControlFlowStatementsResolveTransformer.kt

@@ -125,7 +125,7 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
         whenSubjectExpression: FirWhenSubjectExpression,
         data: ResolutionMode
     ): CompositeTransformResult<FirStatement> {
-        val parentWhen = whenSubjectExpression.whenSubject.whenExpression
+        val parentWhen = whenSubjectExpression.whenRef.value
         val subjectType = parentWhen.subject?.resultType ?: parentWhen.subjectVariable?.returnTypeRef
         if (subjectType != null) {
             whenSubjectExpression.resultType = subjectType

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -177,7 +177,7 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
         dataFlowAnalyzer.enterSafeCallAfterNullCheck(safeCallExpression)
 
         safeCallExpression.apply {
-            checkedSubject.value.propagateTypeFromOriginalReceiver(receiver, components.session)
+            checkedSubjectRef.value.propagateTypeFromOriginalReceiver(receiver, components.session)
             transformRegularQualifiedAccess(this@FirExpressionsResolveTransformer, data)
             propagateTypeFromQualifiedAccessAfterNullCheck(receiver, session)
         }

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/FirCheckedSafeCallSubject.kt

@@ -5,7 +5,7 @@
 
 package org.jetbrains.kotlin.fir.expressions
 
-import org.jetbrains.kotlin.fir.FirSafeCallOriginalReceiverReference
+import org.jetbrains.kotlin.fir.FirExpressionRef
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.visitors.*
@@ -19,7 +19,7 @@ abstract class FirCheckedSafeCallSubject : FirExpression() {
     abstract override val source: FirSourceElement?
     abstract override val typeRef: FirTypeRef
     abstract override val annotations: List<FirAnnotationCall>
-    abstract val originalReceiverReference: FirSafeCallOriginalReceiverReference
+    abstract val originalReceiverRef: FirExpressionRef<FirExpression>
 
     override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R = visitor.visitCheckedSafeCallSubject(this, data)
 

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/FirSafeCallExpression.kt

@@ -5,7 +5,7 @@
 
 package org.jetbrains.kotlin.fir.expressions
 
-import org.jetbrains.kotlin.fir.FirSafeCallCheckedSubjectReference
+import org.jetbrains.kotlin.fir.FirExpressionRef
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.visitors.*
@@ -20,7 +20,7 @@ abstract class FirSafeCallExpression : FirExpression() {
     abstract override val typeRef: FirTypeRef
     abstract override val annotations: List<FirAnnotationCall>
     abstract val receiver: FirExpression
-    abstract val checkedSubject: FirSafeCallCheckedSubjectReference
+    abstract val checkedSubjectRef: FirExpressionRef<FirCheckedSafeCallSubject>
     abstract val regularQualifiedAccess: FirQualifiedAccess
 
     override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R = visitor.visitSafeCallExpression(this, data)

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/FirWhenSubjectExpression.kt

@@ -5,8 +5,8 @@
 
 package org.jetbrains.kotlin.fir.expressions
 
+import org.jetbrains.kotlin.fir.FirExpressionRef
 import org.jetbrains.kotlin.fir.FirSourceElement
-import org.jetbrains.kotlin.fir.FirWhenSubject
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.visitors.*
 
@@ -19,7 +19,7 @@ abstract class FirWhenSubjectExpression : FirExpression() {
     abstract override val source: FirSourceElement?
     abstract override val typeRef: FirTypeRef
     abstract override val annotations: List<FirAnnotationCall>
-    abstract val whenSubject: FirWhenSubject
+    abstract val whenRef: FirExpressionRef<FirWhenExpression>
 
     override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R = visitor.visitWhenSubjectExpression(this, data)
 

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/builder/FirCheckedSafeCallSubjectBuilder.kt

@@ -6,12 +6,13 @@
 package org.jetbrains.kotlin.fir.expressions.builder
 
 import kotlin.contracts.*
-import org.jetbrains.kotlin.fir.FirSafeCallOriginalReceiverReference
+import org.jetbrains.kotlin.fir.FirExpressionRef
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.builder.FirAnnotationContainerBuilder
 import org.jetbrains.kotlin.fir.builder.FirBuilderDsl
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirCheckedSafeCallSubject
+import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.builder.FirExpressionBuilder
 import org.jetbrains.kotlin.fir.expressions.impl.FirCheckedSafeCallSubjectImpl
 import org.jetbrains.kotlin.fir.types.FirTypeRef
@@ -28,14 +29,14 @@ class FirCheckedSafeCallSubjectBuilder : FirAnnotationContainerBuilder, FirExpre
     override var source: FirSourceElement? = null
     override var typeRef: FirTypeRef = FirImplicitTypeRefImpl(null)
     override val annotations: MutableList<FirAnnotationCall> = mutableListOf()
-    lateinit var originalReceiverReference: FirSafeCallOriginalReceiverReference
+    lateinit var originalReceiverRef: FirExpressionRef<FirExpression>
 
     override fun build(): FirCheckedSafeCallSubject {
         return FirCheckedSafeCallSubjectImpl(
             source,
             typeRef,
             annotations,
-            originalReceiverReference,
+            originalReceiverRef,
         )
     }
 

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/builder/FirSafeCallExpressionBuilder.kt

@@ -6,11 +6,12 @@
 package org.jetbrains.kotlin.fir.expressions.builder
 
 import kotlin.contracts.*
-import org.jetbrains.kotlin.fir.FirSafeCallCheckedSubjectReference
+import org.jetbrains.kotlin.fir.FirExpressionRef
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.builder.FirAnnotationContainerBuilder
 import org.jetbrains.kotlin.fir.builder.FirBuilderDsl
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
+import org.jetbrains.kotlin.fir.expressions.FirCheckedSafeCallSubject
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccess
 import org.jetbrains.kotlin.fir.expressions.FirSafeCallExpression
@@ -31,7 +32,7 @@ class FirSafeCallExpressionBuilder : FirAnnotationContainerBuilder, FirExpressio
     override var typeRef: FirTypeRef = FirImplicitTypeRefImpl(null)
     override val annotations: MutableList<FirAnnotationCall> = mutableListOf()
     lateinit var receiver: FirExpression
-    lateinit var checkedSubject: FirSafeCallCheckedSubjectReference
+    lateinit var checkedSubjectRef: FirExpressionRef<FirCheckedSafeCallSubject>
     lateinit var regularQualifiedAccess: FirQualifiedAccess
 
     override fun build(): FirSafeCallExpression {
@@ -40,7 +41,7 @@ class FirSafeCallExpressionBuilder : FirAnnotationContainerBuilder, FirExpressio
             typeRef,
             annotations,
             receiver,
-            checkedSubject,
+            checkedSubjectRef,
             regularQualifiedAccess,
         )
     }

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/builder/FirWhenSubjectExpressionBuilder.kt

@@ -6,11 +6,12 @@
 package org.jetbrains.kotlin.fir.expressions.builder
 
 import kotlin.contracts.*
+import org.jetbrains.kotlin.fir.FirExpressionRef
 import org.jetbrains.kotlin.fir.FirSourceElement
-import org.jetbrains.kotlin.fir.FirWhenSubject
 import org.jetbrains.kotlin.fir.builder.FirAnnotationContainerBuilder
 import org.jetbrains.kotlin.fir.builder.FirBuilderDsl
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
+import org.jetbrains.kotlin.fir.expressions.FirWhenExpression
 import org.jetbrains.kotlin.fir.expressions.FirWhenSubjectExpression
 import org.jetbrains.kotlin.fir.expressions.builder.FirExpressionBuilder
 import org.jetbrains.kotlin.fir.expressions.impl.FirWhenSubjectExpressionImpl
@@ -26,13 +27,13 @@ import org.jetbrains.kotlin.fir.visitors.*
 class FirWhenSubjectExpressionBuilder : FirAnnotationContainerBuilder, FirExpressionBuilder {
     override var source: FirSourceElement? = null
     override val annotations: MutableList<FirAnnotationCall> = mutableListOf()
-    lateinit var whenSubject: FirWhenSubject
+    lateinit var whenRef: FirExpressionRef<FirWhenExpression>
 
     override fun build(): FirWhenSubjectExpression {
         return FirWhenSubjectExpressionImpl(
             source,
             annotations,
-            whenSubject,
+            whenRef,
         )
     }
 

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/impl/FirCheckedSafeCallSubjectImpl.kt

@@ -5,10 +5,11 @@
 
 package org.jetbrains.kotlin.fir.expressions.impl
 
-import org.jetbrains.kotlin.fir.FirSafeCallOriginalReceiverReference
+import org.jetbrains.kotlin.fir.FirExpressionRef
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirCheckedSafeCallSubject
+import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.visitors.*
 
@@ -21,7 +22,7 @@ internal class FirCheckedSafeCallSubjectImpl(
     override val source: FirSourceElement?,
     override var typeRef: FirTypeRef,
     override val annotations: MutableList<FirAnnotationCall>,
-    override val originalReceiverReference: FirSafeCallOriginalReceiverReference,
+    override val originalReceiverRef: FirExpressionRef<FirExpression>,
 ) : FirCheckedSafeCallSubject() {
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
         typeRef.accept(visitor, data)

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/impl/FirSafeCallExpressionImpl.kt

@@ -5,9 +5,10 @@
 
 package org.jetbrains.kotlin.fir.expressions.impl
 
-import org.jetbrains.kotlin.fir.FirSafeCallCheckedSubjectReference
+import org.jetbrains.kotlin.fir.FirExpressionRef
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
+import org.jetbrains.kotlin.fir.expressions.FirCheckedSafeCallSubject
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccess
 import org.jetbrains.kotlin.fir.expressions.FirSafeCallExpression
@@ -24,7 +25,7 @@ internal class FirSafeCallExpressionImpl(
     override var typeRef: FirTypeRef,
     override val annotations: MutableList<FirAnnotationCall>,
     override var receiver: FirExpression,
-    override val checkedSubject: FirSafeCallCheckedSubjectReference,
+    override val checkedSubjectRef: FirExpressionRef<FirCheckedSafeCallSubject>,
     override var regularQualifiedAccess: FirQualifiedAccess,
 ) : FirSafeCallExpression() {
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/impl/FirWhenSubjectExpressionImpl.kt

@@ -5,9 +5,10 @@
 
 package org.jetbrains.kotlin.fir.expressions.impl
 
+import org.jetbrains.kotlin.fir.FirExpressionRef
 import org.jetbrains.kotlin.fir.FirSourceElement
-import org.jetbrains.kotlin.fir.FirWhenSubject
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
+import org.jetbrains.kotlin.fir.expressions.FirWhenExpression
 import org.jetbrains.kotlin.fir.expressions.FirWhenSubjectExpression
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.visitors.*
@@ -20,9 +21,9 @@ import org.jetbrains.kotlin.fir.visitors.*
 internal class FirWhenSubjectExpressionImpl(
     override val source: FirSourceElement?,
     override val annotations: MutableList<FirAnnotationCall>,
-    override val whenSubject: FirWhenSubject,
+    override val whenRef: FirExpressionRef<FirWhenExpression>,
 ) : FirWhenSubjectExpression() {
-    override val typeRef: FirTypeRef get() = whenSubject.whenExpression.subject!!.typeRef
+    override val typeRef: FirTypeRef get() = whenRef.value.subject!!.typeRef
 
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
         annotations.forEach { it.accept(visitor, data) }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirExpressionRef.kt

@@ -0,0 +1,23 @@
+/*
+ * Copyright 2010-2020 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.fir
+
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+
+@RequiresOptIn(level = RequiresOptIn.Level.ERROR)
+annotation class FirContractViolation
+
+/**
+ * This class is created to implicitly overcome FIR contract that some node may only be referenced once from its parents (not from any other node)
+ * Thus, these kinds of secondary references may not be common FIR nodes since it would lead to visiting them twice during usual traversal
+ * And this class is used to wrap such references
+ */
+class FirExpressionRef<T : FirExpression> @FirContractViolation constructor () {
+    lateinit var value: T
+    fun bind(value: T) {
+        this.value = value
+    }
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSafeCallCheckedSubjectReference.kt

@@ -1,14 +0,0 @@
-/*
- * Copyright 2010-2020 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.fir
-
-import org.jetbrains.kotlin.fir.expressions.FirCheckedSafeCallSubject
-
-class FirSafeCallCheckedSubjectReference(var value: FirCheckedSafeCallSubject) {
-    fun bind(subject: FirCheckedSafeCallSubject) {
-        this.value = subject
-    }
-}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSafeCallOriginalReceiverReference.kt

@@ -1,14 +0,0 @@
-/*
- * Copyright 2010-2020 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.fir
-
-import org.jetbrains.kotlin.fir.expressions.FirExpression
-
-class FirSafeCallOriginalReceiverReference(var value: FirExpression) {
-    fun bind(realReceiver: FirExpression) {
-        this.value = realReceiver
-    }
-}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirWhenSubject.kt

@@ -1,16 +0,0 @@
-/*
- * Copyright 2010-2019 JetBrains s.r.o. 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.fir
-
-import org.jetbrains.kotlin.fir.expressions.FirWhenExpression
-
-class FirWhenSubject {
-    lateinit var whenExpression: FirWhenExpression
-
-    fun bind(expression: FirWhenExpression) {
-        whenExpression = expression
-    }
-}

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/BuilderConfigurator.kt

@@ -269,6 +269,18 @@ object BuilderConfigurator : AbstractBuilderConfigurator<FirTreeBuilder>(FirTree
             defaultFalse("isNullableLHSForCallableReference")
         }
 
+//        builder(safeCallExpression) {
+//            useTypes(safeCallCheckedSubjectType)
+//        }
+//
+//        builder(checkedSafeCallSubject) {
+//            useTypes(expressionType)
+//        }
+//
+//        builder(whenSubjectExpression) {
+//            useTypes(whenExpressionType)
+//        }
+
         val elementsWithDefaultTypeRef = listOf(
             thisReceiverExpression,
             callableReferenceAccess,

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/ImplementationConfigurator.kt

@@ -286,9 +286,10 @@ object ImplementationConfigurator : AbstractFirTreeImplementationConfigurator()
 
         impl(whenSubjectExpression) {
             default("typeRef") {
-                value = "whenSubject.whenExpression.subject!!.typeRef"
+                value = "whenRef.value.subject!!.typeRef"
                 withGetter = true
             }
+            useTypes(whenExpressionType)
         }
 
         impl(wrappedDelegateExpression) {
@@ -443,6 +444,14 @@ object ImplementationConfigurator : AbstractFirTreeImplementationConfigurator()
             }
         }
 
+        impl(safeCallExpression) {
+            useTypes(safeCallCheckedSubjectType)
+        }
+
+        impl(checkedSafeCallSubject) {
+            useTypes(expressionType)
+        }
+
         noImpl(userTypeRef)
     }
 

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/NodeConfigurator.kt

@@ -441,13 +441,13 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
         safeCallExpression.configure {
             +field("receiver", expression).withTransform()
             // Special node that might be used as a reference to receiver of a safe call after null check
-            +field("checkedSubject", safeCallCheckedSubjectReferenceType)
+            +field("checkedSubjectRef", safeCallCheckedSubjectReferenceType)
             // One that uses checkedReceiver as a receiver
             +field("regularQualifiedAccess", qualifiedAccess, withReplace = true).withTransform()
         }
 
         checkedSafeCallSubject.configure {
-            +field("originalReceiverReference", safeCallReceiverReferenceType)
+            +field("originalReceiverRef", safeCallOriginalReceiverReferenceType)
         }
 
         callableReferenceAccess.configure {
@@ -498,7 +498,7 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
         }
 
         whenSubjectExpression.configure {
-            +field("whenSubject", whenSubjectType)
+            +field("whenRef", whenRefType)
         }
 
         wrappedExpression.configure {

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/Types.kt

@@ -37,9 +37,15 @@ val classIdType = type(ClassId::class)
 val annotationUseSiteTargetType = type(AnnotationUseSiteTarget::class)
 val operationKindType = type("fir.expressions", "LogicOperationKind")
 val coneKotlinTypeType = type(ConeKotlinType::class)
-val whenSubjectType = type("fir", "FirWhenSubject")
+
-val safeCallReceiverReferenceType = type("fir", "FirSafeCallOriginalReceiverReference")
+val whenExpressionType = generatedType("expressions", "FirWhenExpression")
-val safeCallCheckedSubjectReferenceType = type("fir", "FirSafeCallCheckedSubjectReference")
+val expressionType = generatedType("expressions", "FirExpression")
+val safeCallCheckedSubjectType = generatedType("expressions", "FirCheckedSafeCallSubject")
+
+val whenRefType = generatedType("", "FirExpressionRef<FirWhenExpression>")
+val safeCallOriginalReceiverReferenceType = generatedType("", "FirExpressionRef<FirExpression>")
+val safeCallCheckedSubjectReferenceType = generatedType("", "FirExpressionRef<FirCheckedSafeCallSubject>")
+
 val firSessionType = type("fir", "FirSession")
 val emptyCfgReferenceType = generatedType("references.impl", "FirEmptyControlFlowGraphReference")
 val noReceiverExpressionType = generatedType("expressions.impl", "FirNoReceiverExpression")