FIR: set missed source in various FirElements

This could be caught by debuggability tests, such as stepping tests,
which are not enabled for FIR yet.
Instead, for now, full pipeline tests will raise index out-of-bound
errors due to the undefined offsets, which stem from null source.

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/ConversionUtils.kt

@@ -30,8 +30,8 @@ import org.jetbrains.kotlin.fir.symbols.AccessorSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
 import org.jetbrains.kotlin.ir.descriptors.WrappedReceiverParameterDescriptor
 import org.jetbrains.kotlin.ir.expressions.IrConst
 import org.jetbrains.kotlin.ir.expressions.IrConstKind
@@ -46,17 +46,15 @@ import org.jetbrains.kotlin.ir.util.functions
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.psiUtil.startOffsetSkippingComments
-import org.jetbrains.kotlin.types.AbstractTypeChecker
-import org.jetbrains.kotlin.types.AbstractTypeCheckerContext
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.util.OperatorNameConventions
 
 internal fun <T : IrElement> FirElement.convertWithOffsets(
     f: (startOffset: Int, endOffset: Int) -> T
 ): T {
-    if (psi is PsiCompiledElement) return f(-1, -1)
-    val startOffset = psi?.startOffsetSkippingComments ?: -1
-    val endOffset = psi?.endOffset ?: -1
+    if (psi is PsiCompiledElement) return f(UNDEFINED_OFFSET, UNDEFINED_OFFSET)
+    val startOffset = psi?.startOffsetSkippingComments ?: UNDEFINED_OFFSET
+    val endOffset = psi?.endOffset ?: UNDEFINED_OFFSET
     return f(startOffset, endOffset)
 }
 

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt

@@ -550,7 +550,11 @@ class Fir2IrVisitor(
             return this
         }
         // TODO: Other conditions to check?
-        if (expression.typeRef.coneTypeSafe<ConeKotlinType>()?.hasEnhancedNullability != true) {
+        // [TypeOperatorLowering] will retrieve the source (from start offset to end offset) as an assertion message.
+        // Avoid type casting if we can't determine the source for some reasons, e.g., implicit `this` receiver.
+        if (expression.source == null ||
+            expression.typeRef.coneTypeSafe<ConeKotlinType>()?.hasEnhancedNullability != true
+        ) {
             return this
         }
         return IrTypeOperatorCallImpl(

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

@@ -1510,6 +1510,7 @@ class DeclarationsConverter(
                     }
                 rValue = firExpression!!
                 dispatchReceiver = buildThisReceiverExpression {
+                    source = firExpression!!.source
                     calleeReference = buildImplicitThisReference {
                         boundSymbol = containerSymbol
                     }

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

@@ -665,18 +665,21 @@ class ExpressionsConverter(
                     if (hasSubject) {
                         val firCondition = entry.toFirWhenCondition()
                         buildWhenBranch {
+                            source = branch.source
                             condition = firCondition
                             result = branch 
                         }
                     } else {
                         val firCondition = entry.toFirWhenConditionWithoutSubject()
                         buildWhenBranch {
+                            source = branch.source
                             condition = firCondition
                             result = branch
                         }
                     }
                 } else {
                     buildWhenBranch {
+                        source = branch.source
                         condition = buildElseIfTrueCondition()
                         result = branch
                     }

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

@@ -51,8 +51,9 @@ class ValueParameter(
             returnTypeRef = type.copyWithNewSourceKind(FirFakeSourceElementKind.PropertyFromParameter)
             this.name = name
             initializer = buildQualifiedAccessExpression {
+                source = firValueParameter.source
                 calleeReference = buildPropertyFromParameterResolvedNamedReference {
-                    this.name =  name
+                    this.name = name
                     resolvedSymbol = this@ValueParameter.firValueParameter.symbol
                 }
             }

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

@@ -561,11 +561,13 @@ class RawFirBuilder(
                                 source = delegateSource
                                 calleeReference =
                                     buildResolvedNamedReference {
+                                        source = delegateSource
                                         name = delegateName
                                         resolvedSymbol = delegateField.symbol
                                     }
                                 rValue = delegateExpression
                                 dispatchReceiver = buildThisReceiverExpression {
+                                    source = delegateSource
                                     calleeReference = buildImplicitThisReference {
                                         boundSymbol = containerSymbol
                                     }
@@ -1533,6 +1535,7 @@ class RawFirBuilder(
                 }
                 if (expression.elseKeyword != null) {
                     branches += buildWhenBranch {
+                        source = expression.elseKeyword?.toFirPsiSourceElement()
                         condition = buildElseIfTrueCondition()
                         result = expression.`else`.toFirBlock()
                     }

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

@@ -1038,6 +1038,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
                 source = parameterSource
                 typeRef = firPropertyReturnTypeRefWithCorrectSourceKind
                 dispatchReceiver = buildThisReceiverExpression {
+                    source = parameterSource
                     calleeReference = buildImplicitThisReference {
                         boundSymbol = classBuilder.symbol
                     }
@@ -1112,6 +1113,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
 
     private fun FirVariable<*>.toQualifiedAccess(): FirQualifiedAccessExpression = buildQualifiedAccessExpression {
         calleeReference = buildResolvedNamedReference {
+            source = this@toQualifiedAccess.source
             name = this@toQualifiedAccess.name
             resolvedSymbol = this@toQualifiedAccess.symbol
         }

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

@@ -190,7 +190,7 @@ fun FirExpression.generateComparisonExpression(
 
     val compareToCall = createConventionCall(
         operationReferenceSource,
-        baseSource?.fakeElement(FirFakeSourceElementKind.GeneratedCompararisonExpression),
+        baseSource?.fakeElement(FirFakeSourceElementKind.GeneratedComparisonExpression),
         argument,
         OperatorNameConventions.COMPARE_TO
     )
@@ -229,13 +229,13 @@ private fun FirExpression.createConventionCall(
 
 fun generateAccessExpression(
     qualifiedSource: FirSourceElement?,
-    calleReferenceSource: FirSourceElement?,
+    calleeReferenceSource: FirSourceElement?,
     name: Name
 ): FirQualifiedAccessExpression =
     buildQualifiedAccessExpression {
         this.source = qualifiedSource
         calleeReference = buildSimpleNamedReference {
-            this.source = calleReferenceSource
+            this.source = calleeReferenceSource
             this.name = name
         }
     }
@@ -379,6 +379,7 @@ fun FirPropertyBuilder.generateAccessorsByDelegate(
         val annotations = getter?.annotations
         val returnTarget = FirFunctionTarget(null, isLambda = false)
         getter = buildPropertyAccessor {
+            this.source = delegateBuilder.source
             this.session = session
             origin = FirDeclarationOrigin.Source
             returnTypeRef = buildImplicitTypeRef()
@@ -410,12 +411,14 @@ fun FirPropertyBuilder.generateAccessorsByDelegate(
     if (isVar && (setter == null || setter is FirDefaultPropertyAccessor)) {
         val annotations = setter?.annotations
         setter = buildPropertyAccessor {
+            this.source = delegateBuilder.source
             this.session = session
             origin = FirDeclarationOrigin.Source
             returnTypeRef = session.builtinTypes.unitType
             isGetter = false
             status = FirDeclarationStatusImpl(Visibilities.Unknown, Modality.FINAL)
             val parameter = buildValueParameter {
+                source = delegateBuilder.source
                 this.session = session
                 origin = FirDeclarationOrigin.Source
                 returnTypeRef = buildImplicitTypeRef()
@@ -429,6 +432,7 @@ fun FirPropertyBuilder.generateAccessorsByDelegate(
             symbol = FirPropertyAccessorSymbol()
             body = FirSingleExpressionBlock(
                 buildFunctionCall {
+                    source = delegateBuilder.source
                     explicitReceiver = delegateAccess()
                     calleeReference = buildSimpleNamedReference {
                         name = SET_VALUE
@@ -438,6 +442,7 @@ fun FirPropertyBuilder.generateAccessorsByDelegate(
                         arguments += propertyRef()
                         arguments += buildQualifiedAccessExpression {
                             calleeReference = buildResolvedNamedReference {
+                                source = delegateBuilder.source
                                 name = DELEGATED_SETTER_PARAM
                                 resolvedSymbol = parameter.symbol
                             }
@@ -493,6 +498,7 @@ fun FirQualifiedAccess.wrapWithSafeCall(receiver: FirExpression): FirSafeCallExp
         this.originalReceiverRef = FirExpressionRef<FirExpression>().apply {
             bind(receiver)
         }
+        this.source = receiver.source
     }
 
     replaceExplicitReceiver(checkedSafeCallSubject)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/FirReceivers.kt

@@ -94,6 +94,9 @@ sealed class ImplicitReceiverValue<S : AbstractFirBasedSymbol<*>>(
 
 private fun receiverExpression(symbol: AbstractFirBasedSymbol<*>, type: ConeKotlinType): FirThisReceiverExpression =
     buildThisReceiverExpression {
+        // NB: we can't use `symbol.fir.source` as the source of `this` receiver. For instance, if this is an implicit receiver for a class,
+        // the entire class itself will be set as a source. If combined with an implicit type operation, a certain assertion, like null
+        // check assertion, will retrieve source as an assertion message, which is literally the entire class (!).
         calleeReference = buildImplicitThisReference {
             boundSymbol = symbol
         }

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

@@ -57,6 +57,9 @@ internal fun remapArgumentsWithVararg(
                 (valueParameter.isVararg && arg !is FirNamedArgumentExpression)
             ) {
                 arguments += arg
+                if (this.source == null) {
+                    this.source = arg.source
+                }
             } else if (arguments.isEmpty()) {
                 // `arg` is BEFORE the vararg arguments.
                 newArgumentMapping[arg] = valueParameter

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

@@ -133,7 +133,7 @@ sealed class FirFakeSourceElementKind : FirSourceElementKind() {
     // `a > b` will be wrapped in FirComparisonExpression
     // with real source which points to initial `a > b` expression
     // and inner FirFunctionCall will refer to a fake source
-    object GeneratedCompararisonExpression : FirFakeSourceElementKind()
+    object GeneratedComparisonExpression : FirFakeSourceElementKind()
 
     // a ?: b --> when(val $subj = a) { .... }
     // where `val $subj = a` has a fake source