[FIR] Use source elements instead of psi in ConversionUtils.kt

compiler/fir/lightTree/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt

@@ -194,11 +194,13 @@ class ExpressionsConverter(
         var leftArgNode: LighterASTNode? = null
         var rightArgAsFir: FirExpression = buildErrorExpression(null, ConeSimpleDiagnostic("No right operand", DiagnosticKind.Syntax))
         var rightArg: LighterASTNode? = null
+        var operationReferenceSource: FirLightSourceElement? = null
         binaryExpression.forEachChildren {
             when (it.tokenType) {
                 OPERATION_REFERENCE -> {
                     isLeftArgument = false
                     operationTokenName = it.asText
+                    operationReferenceSource = it.toFirSourceElement()
                 }
                 else -> if (it.isExpression()) {
                     if (isLeftArgument) {
@@ -211,25 +213,18 @@ class ExpressionsConverter(
             }
         }
 
+        val baseSource = binaryExpression.toFirSourceElement()
         val operationToken = operationTokenName.getOperationSymbol()
         val leftArgAsFir = getAsFirExpression<FirExpression>(leftArgNode, "No left operand")
         when (operationToken) {
             ELVIS ->
-                return leftArgAsFir.generateNotNullOrOther(
-                    baseSession, rightArgAsFir, "elvis", null
-                )
+                return leftArgAsFir.generateNotNullOrOther(baseSession, rightArgAsFir, "elvis", baseSource)
             ANDAND, OROR ->
-                return leftArgAsFir.generateLazyLogicalOperation(
-                    rightArgAsFir, operationToken == ANDAND, null
-                )
+                return leftArgAsFir.generateLazyLogicalOperation(rightArgAsFir, operationToken == ANDAND, baseSource)
             in OperatorConventions.IN_OPERATIONS ->
-                return rightArgAsFir.generateContainsOperation(
-                    leftArgAsFir, operationToken == NOT_IN, null, null
-                )
+                return rightArgAsFir.generateContainsOperation(leftArgAsFir, operationToken == NOT_IN, baseSource, operationReferenceSource)
             in OperatorConventions.COMPARISON_OPERATIONS ->
-                return leftArgAsFir.generateComparisonExpression(
-                    rightArgAsFir, operationToken, null, null
-                )
+                return leftArgAsFir.generateComparisonExpression(rightArgAsFir, operationToken, baseSource, operationReferenceSource)
         }
         val conventionCallName = operationToken.toBinaryName()
         return if (conventionCallName != null || operationToken == IDENTIFIER) {
@@ -245,7 +240,7 @@ class ExpressionsConverter(
         } else {
             val firOperation = operationToken.toFirOperation()
             if (firOperation in FirOperation.ASSIGNMENTS) {
-                return leftArgNode.generateAssignment(null, rightArg, rightArgAsFir, firOperation) { getAsFirExpression(this) }
+                return leftArgNode.generateAssignment(binaryExpression.toFirSourceElement(), rightArg, rightArgAsFir, firOperation) { getAsFirExpression(this) }
             } else {
                 buildOperatorCall {
                     source = binaryExpression.toFirSourceElement()
@@ -688,9 +683,13 @@ class ExpressionsConverter(
     private fun convertWhenConditionInRange(whenCondition: LighterASTNode, subject: FirWhenSubject?): FirExpression {
         var isNegate = false
         var firExpression: FirExpression = buildErrorExpression(null, ConeSimpleDiagnostic("No range in condition with range", DiagnosticKind.Syntax))
+        var conditionSource: FirLightSourceElement? = null
         whenCondition.forEachChildren {
             when {
-                it.tokenType == OPERATION_REFERENCE && it.asText == NOT_IN.value -> isNegate = true
+                it.tokenType == OPERATION_REFERENCE && it.asText == NOT_IN.value -> {
+                    conditionSource = it.toFirSourceElement()
+                    isNegate = true
+                }
                 else -> if (it.isExpression()) firExpression = getAsFirExpression(it)
             }
         }
@@ -709,8 +708,8 @@ class ExpressionsConverter(
         return firExpression.generateContainsOperation(
             subjectExpression,
             inverted = isNegate,
-            base = null, // TODO: replace with FirSourceElement
-            operationReference = null
+            baseSource = whenCondition.toFirSourceElement(),
+            operationReferenceSource = conditionSource
         )
     }
 

compiler/fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt

@@ -217,7 +217,12 @@ internal fun KtWhenCondition.toFirWhenCondition(
         }
         is KtWhenConditionInRange -> {
             val firRange = rangeExpression.convert("No range in condition with range")
-            firRange.generateContainsOperation(firSubjectExpression, isNegated, rangeExpression, operationReference)
+            firRange.generateContainsOperation(
+                firSubjectExpression,
+                isNegated,
+                rangeExpression?.toFirSourceElement(),
+                operationReference.toFirSourceElement()
+            )
         }
         is KtWhenConditionIsPattern -> {
             buildTypeOperatorCall {
@@ -255,14 +260,12 @@ internal fun Array<KtWhenCondition>.toFirWhenCondition(
 fun FirExpression.generateContainsOperation(
     argument: FirExpression,
     inverted: Boolean,
-    base: KtExpression?,
-    operationReference: KtOperationReferenceExpression?,
+    baseSource: FirSourceElement?,
+    operationReferenceSource: FirSourceElement?
 ): FirFunctionCall {
-    val baseSource = base?.toFirSourceElement()
-    val containsCall = createConventionCall(operationReference, baseSource, argument, OperatorNameConventions.CONTAINS)
+    val containsCall = createConventionCall(operationReferenceSource, baseSource, argument, OperatorNameConventions.CONTAINS)
     if (!inverted) return containsCall
 
-    val operationReferenceSource = operationReference?.toFirSourceElement()
     return buildFunctionCall {
         source = baseSource
         calleeReference = buildSimpleNamedReference {
@@ -276,15 +279,14 @@ fun FirExpression.generateContainsOperation(
 fun FirExpression.generateComparisonExpression(
     argument: FirExpression,
     operatorToken: IElementType,
-    base: KtExpression?,
-    operationReference: KtOperationReferenceExpression?,
+    baseSource: FirSourceElement?,
+    operationReferenceSource: FirSourceElement?,
 ): FirComparisonExpression {
     require(operatorToken in OperatorConventions.COMPARISON_OPERATIONS) {
         "$operatorToken is not in ${OperatorConventions.COMPARISON_OPERATIONS}"
     }
 
-    val baseSource = base?.toFirSourceElement()
-    val compareToCall = createConventionCall(operationReference, baseSource, argument, OperatorNameConventions.COMPARE_TO)
+    val compareToCall = createConventionCall(operationReferenceSource, baseSource, argument, OperatorNameConventions.COMPARE_TO)
 
     val firOperation = when (operatorToken) {
         KtTokens.LT -> FirOperation.LT
@@ -302,12 +304,11 @@ fun FirExpression.generateComparisonExpression(
 }
 
 private fun FirExpression.createConventionCall(
-    operationReference: KtOperationReferenceExpression?,
-    baseSource: FirPsiSourceElement?,
+    operationReferenceSource: FirSourceElement?,
+    baseSource: FirSourceElement?,
     argument: FirExpression,
     conventionName: Name
 ): FirFunctionCall {
-    val operationReferenceSource = operationReference?.toFirSourceElement()
     return buildFunctionCall {
         source = baseSource
         calleeReference = buildSimpleNamedReference {

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

@@ -1383,11 +1383,11 @@ class RawFirBuilder(
                     return leftArgument.generateLazyLogicalOperation(rightArgument, operationToken == ANDAND, source)
                 in OperatorConventions.IN_OPERATIONS ->
                     return rightArgument.generateContainsOperation(
-                        leftArgument, operationToken == NOT_IN, expression, expression.operationReference,
+                        leftArgument, operationToken == NOT_IN, source, expression.operationReference.toFirSourceElement(),
                     )
                 in OperatorConventions.COMPARISON_OPERATIONS ->
                     return leftArgument.generateComparisonExpression(
-                        rightArgument, operationToken, expression, expression.operationReference,
+                        rightArgument, operationToken, source, expression.operationReference.toFirSourceElement(),
                     )
             }
             val conventionCallName = operationToken.toBinaryName()