FIR: Rename FirSourceElement::withKind to fakeElement

And reinforce its contract

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

@@ -25,7 +25,6 @@ import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.*
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
-import org.jetbrains.kotlin.fir.types.impl.FirResolvedTypeRefImpl
 import org.jetbrains.kotlin.lexer.KtTokens.CLOSING_QUOTE
 import org.jetbrains.kotlin.lexer.KtTokens.OPEN_QUOTE
 import org.jetbrains.kotlin.name.FqName
@@ -341,7 +340,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
                         SHORT_STRING_TEMPLATE_ENTRY, LONG_STRING_TEMPLATE_ENTRY -> {
                             hasExpressions = true
                             val firExpression = entry.convertTemplateEntry("Incorrect template argument")
-                            val source = firExpression.source?.withKind(FirFakeSourceElementKind.GeneratedToStringCallOnTemplateEntry)
+                            val source = firExpression.source?.fakeElement(FirFakeSourceElementKind.GeneratedToStringCallOnTemplateEntry)
                             buildFunctionCall {
                                 this.source = source
                                 explicitReceiver = firExpression
@@ -408,7 +407,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
         }
         return buildBlock {
             val baseSource = baseExpression?.toFirSourceElement()
-            val desugaredSource = baseSource?.withKind(FirFakeSourceElementKind.DesugaredIncrementOrDecrement)
+            val desugaredSource = baseSource?.fakeElement(FirFakeSourceElementKind.DesugaredIncrementOrDecrement)
             source = desugaredSource
             val tempName = Name.special("<unary>")
             val temporaryVariable = generateTemporaryVariable(
@@ -744,7 +743,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
                 componentIndex++
                 val parameterSource = sourceNode?.toFirSourceElement()
                 val componentFunction = buildSimpleFunction {
-                    source = parameterSource?.withKind(FirFakeSourceElementKind.DataClassGeneratedMembers)
+                    source = parameterSource?.fakeElement(FirFakeSourceElementKind.DataClassGeneratedMembers)
                     session = this@DataClassMembersGenerator.session
                     origin = FirDeclarationOrigin.Source
                     returnTypeRef = firProperty.returnTypeRef

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

@@ -170,9 +170,9 @@ fun FirExpression.generateContainsOperation(
     if (!inverted) return containsCall
 
     return buildFunctionCall {
-        source = baseSource?.withKind(FirFakeSourceElementKind.DesugaredInvertedContains)
+        source = baseSource?.fakeElement(FirFakeSourceElementKind.DesugaredInvertedContains)
         calleeReference = buildSimpleNamedReference {
-            source = operationReferenceSource?.withKind(FirFakeSourceElementKind.DesugaredInvertedContains)
+            source = operationReferenceSource?.fakeElement(FirFakeSourceElementKind.DesugaredInvertedContains)
             name = OperatorNameConventions.NOT
         }
         explicitReceiver = containsCall
@@ -191,7 +191,7 @@ fun FirExpression.generateComparisonExpression(
 
     val compareToCall = createConventionCall(
         operationReferenceSource,
-        baseSource?.withKind(FirFakeSourceElementKind.GeneratedCompararisonExpression),
+        baseSource?.fakeElement(FirFakeSourceElementKind.GeneratedCompararisonExpression),
         argument,
         OperatorNameConventions.COMPARE_TO
     )
@@ -483,6 +483,6 @@ fun FirModifiableQualifiedAccess.wrapWithSafeCall(receiver: FirExpression): FirS
             bind(checkedSafeCallSubject)
         }
         this.regularQualifiedAccess = this@wrapWithSafeCall
-        this.source = this@wrapWithSafeCall.source?.withKind(FirFakeSourceElementKind.DesugaredSafeCallExpression)
+        this.source = this@wrapWithSafeCall.source?.fakeElement(FirFakeSourceElementKind.DesugaredSafeCallExpression)
     }
 }

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

@@ -105,7 +105,7 @@ internal fun generateDestructuringBlock(
                 returnTypeRef = entry.typeReference.toFirOrImplicitTypeRef()
                 this.name = name
                 initializer = buildComponentCall {
-                    val componentCallSource = entrySource.withKind(FirFakeSourceElementKind.DesugaredComponentFunctionCall)
+                    val componentCallSource = entrySource.fakeElement(FirFakeSourceElementKind.DesugaredComponentFunctionCall)
                     source = componentCallSource
                     explicitReceiver = generateResolvedAccessExpression(componentCallSource, container)
                     componentIndex = index + 1

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

@@ -311,7 +311,7 @@ class RawFirBuilder(
                 extractValueParametersTo(this, propertyTypeRef)
                 if (!isGetter && valueParameters.isEmpty()) {
                     valueParameters += buildDefaultSetterValueParameter {
-                        this.source = source.withKind(FirFakeSourceElementKind.DefaultAccessor)
+                        this.source = source.fakeElement(FirFakeSourceElementKind.DefaultAccessor)
                         session = baseSession
                         origin = FirDeclarationOrigin.Source
                         returnTypeRef = propertyTypeRef
@@ -384,7 +384,7 @@ class RawFirBuilder(
                 symbol = FirPropertySymbol(callableIdForName(propertyName))
                 isLocal = false
                 this.status = status
-                val defaultAccessorSource = propertySource.withKind(FirFakeSourceElementKind.DefaultAccessor)
+                val defaultAccessorSource = propertySource.fakeElement(FirFakeSourceElementKind.DefaultAccessor)
                 getter = FirDefaultPropertyGetter(defaultAccessorSource, baseSession, FirDeclarationOrigin.Source, type, visibility)
                 setter = if (isMutable) FirDefaultPropertySetter(
                     defaultAccessorSource,
@@ -574,7 +574,7 @@ class RawFirBuilder(
             val constructorSource = this?.toFirSourceElement()
                 ?: owner.toFirPsiSourceElement(FirFakeSourceElementKind.ImplicitConstructor)
             val firDelegatedCall = buildDelegatedConstructorCall {
-                source = constructorCallee ?: constructorSource.withKind(FirFakeSourceElementKind.DelegatingConstructorCall)
+                source = constructorCallee ?: constructorSource.fakeElement(FirFakeSourceElementKind.DelegatingConstructorCall)
                 constructedTypeRef = delegatedSuperTypeRef
                 isThis = false
                 if (!stubMode) {
@@ -995,7 +995,7 @@ class RawFirBuilder(
                 val expressionSource = expression.toFirSourceElement()
                 label = context.firLabels.pop() ?: context.calleeNamesForLambda.lastOrNull()?.let {
                     buildLabel {
-                        source = expressionSource.withKind(FirFakeSourceElementKind.GeneratedLambdaLabel)
+                        source = expressionSource.fakeElement(FirFakeSourceElementKind.GeneratedLambdaLabel)
                         name = it.asString()
                     }
                 }
@@ -1011,10 +1011,10 @@ class RawFirBuilder(
                         if (statements.isEmpty()) {
                             statements.add(
                                 buildReturnExpression {
-                                    source = expressionSource.withKind(FirFakeSourceElementKind.ImplicitReturn)
+                                    source = expressionSource.fakeElement(FirFakeSourceElementKind.ImplicitReturn)
                                     this.target = target
                                     result = buildUnitExpression {
-                                        source = expressionSource.withKind(FirFakeSourceElementKind.ImplicitUnit)
+                                        source = expressionSource.fakeElement(FirFakeSourceElementKind.ImplicitUnit)
                                     }
                                 }
                             )
@@ -1706,7 +1706,7 @@ class RawFirBuilder(
 
                 else -> {
                     buildSimpleNamedReference {
-                        source = defaultSource.withKind(FirFakeSourceElementKind.ImplicitInvokeCall)
+                        source = defaultSource.fakeElement(FirFakeSourceElementKind.ImplicitInvokeCall)
                         name = OperatorNameConventions.INVOKE
                     } to calleeExpression.toFirExpression("Incorrect invoke receiver")
                 }
@@ -1793,7 +1793,7 @@ class RawFirBuilder(
                 val sourceElement = expression.toFirSourceElement()
                 source = sourceElement
                 calleeReference = buildExplicitThisReference {
-                    source = sourceElement.withKind(FirFakeSourceElementKind.ExplicitThisOrSuperReference)
+                    source = sourceElement.fakeElement(FirFakeSourceElementKind.ExplicitThisOrSuperReference)
                     labelName = expression.getLabelName()
                 }
             }
@@ -1805,7 +1805,7 @@ class RawFirBuilder(
             return buildQualifiedAccessExpression {
                 this.source = source
                 calleeReference = buildExplicitSuperReference {
-                    source.withKind(FirFakeSourceElementKind.ExplicitThisOrSuperReference)
+                    source.fakeElement(FirFakeSourceElementKind.ExplicitThisOrSuperReference)
                     labelName = expression.getLabelName()
                     superTypeRef = superType.toFirOrImplicitType()
                 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt

@@ -222,7 +222,7 @@ fun FirFunction<*>.constructFunctionalTypeRef(isSuspend: Boolean = false): FirRe
     val functionalType = createFunctionalType(parameters, receiverTypeRef?.coneType, rawReturnType, isSuspend = isSuspend)
 
     return buildResolvedTypeRef {
-        source = this@constructFunctionalTypeRef.source?.withKind(FirFakeSourceElementKind.ImplicitTypeRef)
+        source = this@constructFunctionalTypeRef.source?.fakeElement(FirFakeSourceElementKind.ImplicitTypeRef)
         type = functionalType
         this.isSuspend = isSuspend
     }
@@ -325,7 +325,7 @@ fun <T : FirResolvable> BodyResolveComponents.typeFromCallee(access: T): FirReso
     return when (val newCallee = access.calleeReference) {
         is FirErrorNamedReference ->
             buildErrorTypeRef {
-                source = access.source?.withKind(FirFakeSourceElementKind.ErrorTypeRef)
+                source = access.source?.fakeElement(FirFakeSourceElementKind.ErrorTypeRef)
                 diagnostic = ConeStubDiagnostic(newCallee.diagnostic)
             }
         is FirNamedReferenceWithCandidate -> {
@@ -390,7 +390,7 @@ fun BodyResolveComponents.transformQualifiedAccessUsingSmartcastInfo(qualifiedAc
     val intersectedType = ConeTypeIntersector.intersectTypes(inferenceComponents.ctx, allTypes)
     // TODO: add check that intersectedType is not equal to original type
     val intersectedTypeRef = buildResolvedTypeRef {
-        source = qualifiedAccessExpression.resultType.source?.withKind(FirFakeSourceElementKind.SmartCastedTypeRef)
+        source = qualifiedAccessExpression.resultType.source?.fakeElement(FirFakeSourceElementKind.SmartCastedTypeRef)
         type = intersectedType
         annotations += qualifiedAccessExpression.resultType.annotations
     }

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

@@ -686,7 +686,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                     override fun <E : FirElement> transformElement(element: E, data: FirExpression): CompositeTransformResult<E> {
                         if (element == lastStatement) {
                             val returnExpression = buildReturnExpression {
-                                source = element.source?.withKind(FirFakeSourceElementKind.ImplicitReturn)
+                                source = element.source?.fakeElement(FirFakeSourceElementKind.ImplicitReturn)
                                 result = lastStatement
                                 target = FirFunctionTarget(null, isLambda = this@addReturn.isLambda).also {
                                     it.bind(this@addReturn)

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

@@ -317,7 +317,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
             rightArgument = rightArgument.transformSingle(transformer, ResolutionMode.ContextDependent)
 
             fun createFunctionCall(name: Name) = buildFunctionCall {
-                source = operatorCall.source?.withKind(FirFakeSourceElementKind.DesugaredCompoundAssignment)
+                source = operatorCall.source?.fakeElement(FirFakeSourceElementKind.DesugaredCompoundAssignment)
                 explicitReceiver = leftArgument
                 argumentList = buildUnaryArgumentList(rightArgument)
                 calleeReference = buildSimpleNamedReference {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirClassSubstitutionScope.kt

@@ -29,7 +29,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
-import org.jetbrains.kotlin.fir.withKind
+import org.jetbrains.kotlin.fir.fakeElement
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 
@@ -548,12 +548,18 @@ fun FirTypeRef.withReplacedReturnType(newType: ConeKotlinType?): FirTypeRef {
     }
 }
 
-fun FirTypeRef.withReplacedConeType(newType: ConeKotlinType?, sourceElementKind: FirFakeSourceElementKind? = null): FirResolvedTypeRef {
+fun FirTypeRef.withReplacedConeType(
+    newType: ConeKotlinType?,
+    firFakeSourceElementKind: FirFakeSourceElementKind? = null
+): FirResolvedTypeRef {
     require(this is FirResolvedTypeRef)
     if (newType == null) return this
 
     return buildResolvedTypeRef {
-        source = this@withReplacedConeType.source?.withKind(sourceElementKind)
+        source = if (firFakeSourceElementKind != null)
+            this@withReplacedConeType.source?.fakeElement(firFakeSourceElementKind)
+        else
+            this@withReplacedConeType.source
         type = newType
         annotations += this@withReplacedConeType.annotations
     }

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

@@ -17,10 +17,10 @@ import org.jetbrains.kotlin.fir.expressions.FirLoop
 import org.jetbrains.kotlin.fir.expressions.builder.FirExpressionBuilder
 import org.jetbrains.kotlin.fir.expressions.builder.FirLoopJumpBuilder
 import org.jetbrains.kotlin.fir.expressions.impl.FirBreakExpressionImpl
+import org.jetbrains.kotlin.fir.fakeElement
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitNothingTypeRef
 import org.jetbrains.kotlin.fir.visitors.*
-import org.jetbrains.kotlin.fir.withKind
 
 /*
  * This file was generated automatically

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

@@ -17,10 +17,10 @@ import org.jetbrains.kotlin.fir.expressions.FirLoop
 import org.jetbrains.kotlin.fir.expressions.builder.FirExpressionBuilder
 import org.jetbrains.kotlin.fir.expressions.builder.FirLoopJumpBuilder
 import org.jetbrains.kotlin.fir.expressions.impl.FirContinueExpressionImpl
+import org.jetbrains.kotlin.fir.fakeElement
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitNothingTypeRef
 import org.jetbrains.kotlin.fir.visitors.*
-import org.jetbrains.kotlin.fir.withKind
 
 /*
  * This file was generated automatically

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

@@ -14,10 +14,10 @@ import org.jetbrains.kotlin.fir.builder.FirBuilderDsl
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.impl.FirElseIfTrueCondition
+import org.jetbrains.kotlin.fir.fakeElement
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitBooleanTypeRef
 import org.jetbrains.kotlin.fir.visitors.*
-import org.jetbrains.kotlin.fir.withKind
 
 /*
  * This file was generated automatically

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

@@ -17,10 +17,10 @@ import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
 import org.jetbrains.kotlin.fir.expressions.builder.FirExpressionBuilder
 import org.jetbrains.kotlin.fir.expressions.impl.FirReturnExpressionImpl
+import org.jetbrains.kotlin.fir.fakeElement
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitNothingTypeRef
 import org.jetbrains.kotlin.fir.visitors.*
-import org.jetbrains.kotlin.fir.withKind
 
 /*
  * This file was generated automatically

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

@@ -16,10 +16,10 @@ import org.jetbrains.kotlin.fir.expressions.FirStringConcatenationCall
 import org.jetbrains.kotlin.fir.expressions.builder.FirCallBuilder
 import org.jetbrains.kotlin.fir.expressions.builder.FirExpressionBuilder
 import org.jetbrains.kotlin.fir.expressions.impl.FirStringConcatenationCallImpl
+import org.jetbrains.kotlin.fir.fakeElement
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitStringTypeRef
 import org.jetbrains.kotlin.fir.visitors.*
-import org.jetbrains.kotlin.fir.withKind
 
 /*
  * This file was generated automatically

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

@@ -15,10 +15,10 @@ import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirThrowExpression
 import org.jetbrains.kotlin.fir.expressions.builder.FirExpressionBuilder
 import org.jetbrains.kotlin.fir.expressions.impl.FirThrowExpressionImpl
+import org.jetbrains.kotlin.fir.fakeElement
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitNothingTypeRef
 import org.jetbrains.kotlin.fir.visitors.*
-import org.jetbrains.kotlin.fir.withKind
 
 /*
  * This file was generated automatically

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

@@ -14,10 +14,10 @@ import org.jetbrains.kotlin.fir.builder.FirBuilderDsl
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.impl.FirUnitExpression
+import org.jetbrains.kotlin.fir.fakeElement
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitUnitTypeRef
 import org.jetbrains.kotlin.fir.visitors.*
-import org.jetbrains.kotlin.fir.withKind
 
 /*
  * This file was generated automatically

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

@@ -11,9 +11,9 @@ import org.jetbrains.kotlin.fir.FirTarget
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirBreakExpression
 import org.jetbrains.kotlin.fir.expressions.FirLoop
+import org.jetbrains.kotlin.fir.fakeElement
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitNothingTypeRef
-import org.jetbrains.kotlin.fir.withKind
 import org.jetbrains.kotlin.fir.visitors.*
 
 /*
@@ -26,7 +26,7 @@ internal class FirBreakExpressionImpl(
     override val annotations: MutableList<FirAnnotationCall>,
     override val target: FirTarget<FirLoop>,
 ) : FirBreakExpression() {
-    override var typeRef: FirTypeRef = FirImplicitNothingTypeRef(source?.withKind(FirFakeSourceElementKind.ImplicitTypeRef))
+    override var typeRef: FirTypeRef = FirImplicitNothingTypeRef(source?.fakeElement(FirFakeSourceElementKind.ImplicitTypeRef))
 
     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/FirContinueExpressionImpl.kt

@@ -11,9 +11,9 @@ import org.jetbrains.kotlin.fir.FirTarget
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirContinueExpression
 import org.jetbrains.kotlin.fir.expressions.FirLoop
+import org.jetbrains.kotlin.fir.fakeElement
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitNothingTypeRef
-import org.jetbrains.kotlin.fir.withKind
 import org.jetbrains.kotlin.fir.visitors.*
 
 /*
@@ -26,7 +26,7 @@ internal class FirContinueExpressionImpl(
     override val annotations: MutableList<FirAnnotationCall>,
     override val target: FirTarget<FirLoop>,
 ) : FirContinueExpression() {
-    override var typeRef: FirTypeRef = FirImplicitNothingTypeRef(source?.withKind(FirFakeSourceElementKind.ImplicitTypeRef))
+    override var typeRef: FirTypeRef = FirImplicitNothingTypeRef(source?.fakeElement(FirFakeSourceElementKind.ImplicitTypeRef))
 
     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/FirElseIfTrueCondition.kt

@@ -10,9 +10,9 @@ import org.jetbrains.kotlin.fir.FirImplementationDetail
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.fakeElement
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitBooleanTypeRef
-import org.jetbrains.kotlin.fir.withKind
 import org.jetbrains.kotlin.fir.visitors.*
 
 /*
@@ -24,7 +24,7 @@ class FirElseIfTrueCondition @FirImplementationDetail constructor(
     override val source: FirSourceElement?,
     override val annotations: MutableList<FirAnnotationCall>,
 ) : FirExpression() {
-    override var typeRef: FirTypeRef = FirImplicitBooleanTypeRef(source?.withKind(FirFakeSourceElementKind.ImplicitTypeRef))
+    override var typeRef: FirTypeRef = FirImplicitBooleanTypeRef(source?.fakeElement(FirFakeSourceElementKind.ImplicitTypeRef))
 
     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/FirReturnExpressionImpl.kt

@@ -12,9 +12,9 @@ import org.jetbrains.kotlin.fir.declarations.FirFunction
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
+import org.jetbrains.kotlin.fir.fakeElement
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitNothingTypeRef
-import org.jetbrains.kotlin.fir.withKind
 import org.jetbrains.kotlin.fir.visitors.*
 
 /*
@@ -28,7 +28,7 @@ internal class FirReturnExpressionImpl(
     override val target: FirTarget<FirFunction<*>>,
     override var result: FirExpression,
 ) : FirReturnExpression() {
-    override var typeRef: FirTypeRef = FirImplicitNothingTypeRef(source?.withKind(FirFakeSourceElementKind.ImplicitTypeRef))
+    override var typeRef: FirTypeRef = FirImplicitNothingTypeRef(source?.fakeElement(FirFakeSourceElementKind.ImplicitTypeRef))
 
     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/FirStringConcatenationCallImpl.kt

@@ -10,9 +10,9 @@ import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirArgumentList
 import org.jetbrains.kotlin.fir.expressions.FirStringConcatenationCall
+import org.jetbrains.kotlin.fir.fakeElement
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitStringTypeRef
-import org.jetbrains.kotlin.fir.withKind
 import org.jetbrains.kotlin.fir.visitors.*
 
 /*
@@ -25,7 +25,7 @@ internal class FirStringConcatenationCallImpl(
     override val annotations: MutableList<FirAnnotationCall>,
     override var argumentList: FirArgumentList,
 ) : FirStringConcatenationCall() {
-    override var typeRef: FirTypeRef = FirImplicitStringTypeRef(source?.withKind(FirFakeSourceElementKind.ImplicitTypeRef))
+    override var typeRef: FirTypeRef = FirImplicitStringTypeRef(source?.fakeElement(FirFakeSourceElementKind.ImplicitTypeRef))
 
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
         annotations.forEach { it.accept(visitor, data) }

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

@@ -10,9 +10,9 @@ import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirThrowExpression
+import org.jetbrains.kotlin.fir.fakeElement
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitNothingTypeRef
-import org.jetbrains.kotlin.fir.withKind
 import org.jetbrains.kotlin.fir.visitors.*
 
 /*
@@ -25,7 +25,7 @@ internal class FirThrowExpressionImpl(
     override val annotations: MutableList<FirAnnotationCall>,
     override var exception: FirExpression,
 ) : FirThrowExpression() {
-    override var typeRef: FirTypeRef = FirImplicitNothingTypeRef(source?.withKind(FirFakeSourceElementKind.ImplicitTypeRef))
+    override var typeRef: FirTypeRef = FirImplicitNothingTypeRef(source?.fakeElement(FirFakeSourceElementKind.ImplicitTypeRef))
 
     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/FirUnitExpression.kt

@@ -10,9 +10,9 @@ import org.jetbrains.kotlin.fir.FirImplementationDetail
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.fakeElement
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitUnitTypeRef
-import org.jetbrains.kotlin.fir.withKind
 import org.jetbrains.kotlin.fir.visitors.*
 
 /*
@@ -24,7 +24,7 @@ class FirUnitExpression @FirImplementationDetail constructor(
     override val source: FirSourceElement?,
     override val annotations: MutableList<FirAnnotationCall>,
 ) : FirExpression() {
-    override var typeRef: FirTypeRef = FirImplicitUnitTypeRef(source?.withKind(FirFakeSourceElementKind.ImplicitTypeRef))
+    override var typeRef: FirTypeRef = FirImplicitUnitTypeRef(source?.fakeElement(FirFakeSourceElementKind.ImplicitTypeRef))
 
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
         typeRef.accept(visitor, data)

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

@@ -32,7 +32,7 @@ private val VALUE = Name.identifier("value")
 fun FirRegularClassBuilder.generateValuesFunction(
     session: FirSession, packageFqName: FqName, classFqName: FqName, makeExpect: Boolean = false
 ) {
-    val sourceElement = source?.withKind(FirFakeSourceElementKind.EnumGeneratedDeclaration)
+    val sourceElement = source?.fakeElement(FirFakeSourceElementKind.EnumGeneratedDeclaration)
     declarations += buildSimpleFunction {
         source = sourceElement
         origin = FirDeclarationOrigin.Source
@@ -61,7 +61,7 @@ fun FirRegularClassBuilder.generateValuesFunction(
 fun FirRegularClassBuilder.generateValueOfFunction(
     session: FirSession, packageFqName: FqName, classFqName: FqName, makeExpect: Boolean = false
 ) {
-    val sourceElement = source?.withKind(FirFakeSourceElementKind.EnumGeneratedDeclaration)
+    val sourceElement = source?.fakeElement(FirFakeSourceElementKind.EnumGeneratedDeclaration)
     declarations += buildSimpleFunction {
         source = sourceElement
         origin = FirDeclarationOrigin.Source

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

@@ -165,15 +165,10 @@ class FirRealPsiSourceElement<out P : PsiElement>(psi: P) : FirPsiSourceElement<
 
 class FirFakeSourceElement<out P : PsiElement>(psi: P, override val kind: FirFakeSourceElementKind) : FirPsiSourceElement<P>(psi)
 
-fun FirSourceElement.withKind(newKind: FirSourceElementKind?): FirSourceElement {
+fun FirSourceElement.fakeElement(newKind: FirFakeSourceElementKind): FirSourceElement {
-    if (newKind == null) return this
     return when (this) {
         is FirLightSourceElement -> this
-        is FirPsiSourceElement<*> -> when (newKind) {
+        is FirPsiSourceElement<*> -> FirFakeSourceElement(psi, newKind)
-            kind -> this
-            FirRealSourceElementKind -> FirRealPsiSourceElement(psi)
-            is FirFakeSourceElementKind -> FirFakeSourceElement(psi, newKind)
-        }
     }
 }
 

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

@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.*
 import org.jetbrains.kotlin.fir.types.impl.*
-import org.jetbrains.kotlin.fir.types.impl.FirFunctionTypeRefImpl
 
 fun ModuleInfo.dependenciesWithoutSelf(): Sequence<ModuleInfo> = dependencies().asSequence().filter { it != this }
 
@@ -25,7 +24,7 @@ fun <R : FirTypeRef> R.copyWithNewSourceKind(newKind: FirFakeSourceElementKind):
     if (source == null) return this
     if (source?.kind == newKind) return this
     if (this is FirDelegatedTypeRef) return this
-    val newSource = source?.withKind(newKind)
+    val newSource = source?.fakeElement(newKind)
 
     @Suppress("UNCHECKED_CAST")
     return when (val typeRef = this) {

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/impl/FirSingleExpressionBlock.kt

@@ -17,13 +17,13 @@ import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 import org.jetbrains.kotlin.fir.visitors.transformInplace
 import org.jetbrains.kotlin.fir.visitors.transformSingle
-import org.jetbrains.kotlin.fir.withKind
+import org.jetbrains.kotlin.fir.fakeElement
 
 class FirSingleExpressionBlock(
     var statement: FirStatement
 ) : FirBlock(), FirAnnotationContainer {
     override val source: FirSourceElement?
-        get() = statement.source?.withKind(FirFakeSourceElementKind.SingleExpressionBlock)
+        get() = statement.source?.fakeElement(FirFakeSourceElementKind.SingleExpressionBlock)
     override val annotations: MutableList<FirAnnotationCall> = mutableListOf()
     override val statements: List<FirStatement> get() = listOf(statement)
     override var typeRef: FirTypeRef = FirImplicitTypeRefImpl(null)

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/impl/FirImplicitBuiltinTypeRef.kt

@@ -17,7 +17,7 @@ import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
-import org.jetbrains.kotlin.fir.withKind
+import org.jetbrains.kotlin.fir.fakeElement
 import org.jetbrains.kotlin.name.ClassId
 
 sealed class FirImplicitBuiltinTypeRef(
@@ -138,7 +138,7 @@ class FirImplicitKMutableProperty2TypeRef(
 fun FirImplicitBuiltinTypeRef.withFakeSource(kind: FirFakeSourceElementKind): FirImplicitBuiltinTypeRef {
     val source = source ?: return this
     if (source.kind == kind) return this
-    val newSource = source.withKind(kind)
+    val newSource = source.fakeElement(kind)
     return when (this) {
         is FirImplicitUnitTypeRef -> FirImplicitUnitTypeRef(newSource)
         is FirImplicitAnyTypeRef -> FirImplicitAnyTypeRef(newSource)

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

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.fir.tree.generator.context
 
-import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.fir.tree.generator.model.*
 import org.jetbrains.kotlin.fir.tree.generator.noReceiverExpressionType
 import org.jetbrains.kotlin.fir.tree.generator.printer.call
@@ -123,9 +122,9 @@ abstract class AbstractFirTreeImplementationConfigurator {
         }
 
         fun defaultTypeRefWithSource(typeRefClass: String) {
-            default("typeRef", "$typeRefClass(source?.withKind(FirFakeSourceElementKind.ImplicitTypeRef))")
+            default("typeRef", "$typeRefClass(source?.fakeElement(FirFakeSourceElementKind.ImplicitTypeRef))")
             implementation.arbitraryImportables += ArbitraryImportable("org.jetbrains.kotlin.fir", "FirFakeSourceElementKind")
-            implementation.arbitraryImportables += ArbitraryImportable("org.jetbrains.kotlin.fir", "withKind")
+            implementation.arbitraryImportables += ArbitraryImportable("org.jetbrains.kotlin.fir", "fakeElement")
         }
 
         fun defaultTrue(field: String, withGetter: Boolean = false) {