FIR: Support FirComparisonOperator in builders

^KT-31163 In Progress

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

@@ -209,18 +209,23 @@ class ExpressionsConverter(
         }
 
         val operationToken = operationTokenName.getOperationSymbol()
+        val leftArgAsFir = getAsFirExpression<FirExpression>(leftArgNode, "No left operand")
         when (operationToken) {
             ELVIS ->
-                return getAsFirExpression<FirExpression>(leftArgNode, "No left operand").generateNotNullOrOther(
+                return leftArgAsFir.generateNotNullOrOther(
                     baseSession, rightArgAsFir, "elvis", null
                 )
             ANDAND, OROR ->
-                return getAsFirExpression<FirExpression>(leftArgNode, "No left operand").generateLazyLogicalOperation(
+                return leftArgAsFir.generateLazyLogicalOperation(
                     rightArgAsFir, operationToken == ANDAND, null
                 )
             in OperatorConventions.IN_OPERATIONS ->
                 return rightArgAsFir.generateContainsOperation(
-                    getAsFirExpression(leftArgNode, "No left operand"), operationToken == NOT_IN, null, null
+                    leftArgAsFir, operationToken == NOT_IN, null, null
+                )
+            in OperatorConventions.COMPARISON_OPERATIONS ->
+                return leftArgAsFir.generateComparisonExpression(
+                    rightArgAsFir, operationToken, null, null
                 )
         }
         val conventionCallName = operationToken.toBinaryName()
@@ -230,7 +235,7 @@ class ExpressionsConverter(
                 calleeReference = buildSimpleNamedReference {
                     name = conventionCallName ?: operationTokenName.nameAsSafeName()
                 }
-                explicitReceiver = getAsFirExpression(leftArgNode, "No left operand")
+                explicitReceiver = leftArgAsFir
                 arguments += rightArgAsFir
             }
         } else {
@@ -241,7 +246,7 @@ class ExpressionsConverter(
                 buildOperatorCall {
                     source = binaryExpression.toFirSourceElement()
                     operation = firOperation
-                    arguments += getAsFirExpression<FirExpression>(leftArgNode, "No left operand")
+                    arguments += leftArgAsFir
                     arguments += rightArgAsFir
                 }
             }

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

@@ -257,17 +257,10 @@ fun FirExpression.generateContainsOperation(
     operationReference: KtOperationReferenceExpression?,
 ): FirFunctionCall {
     val baseSource = base?.toFirSourceElement()
-    val operationReferenceSource = operationReference?.toFirSourceElement()
+    val containsCall = createConventionCall(operationReference, baseSource, argument, OperatorNameConventions.CONTAINS)
-    val containsCall = buildFunctionCall {
-        source = baseSource
-        calleeReference = buildSimpleNamedReference {
-            source = operationReferenceSource
-            name = OperatorNameConventions.CONTAINS
-        }
-        explicitReceiver = this@generateContainsOperation
-        arguments += argument
-    }
     if (!inverted) return containsCall
+
+    val operationReferenceSource = operationReference?.toFirSourceElement()
     return buildFunctionCall {
         source = baseSource
         calleeReference = buildSimpleNamedReference {
@@ -278,6 +271,52 @@ fun FirExpression.generateContainsOperation(
     }
 }
 
+fun FirExpression.generateComparisonExpression(
+    argument: FirExpression,
+    operatorToken: IElementType,
+    base: KtExpression?,
+    operationReference: KtOperationReferenceExpression?,
+): 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 firOperation = when (operatorToken) {
+        KtTokens.LT -> FirOperation.LT
+        KtTokens.GT -> FirOperation.GT
+        KtTokens.LTEQ -> FirOperation.LT_EQ
+        KtTokens.GTEQ -> FirOperation.GT_EQ
+        else -> error("Unknown $operatorToken")
+    }
+
+    return buildComparisonExpression {
+        this.source = baseSource
+        this.operation = firOperation
+        this.compareToCall = compareToCall
+    }
+}
+
+private fun FirExpression.createConventionCall(
+    operationReference: KtOperationReferenceExpression?,
+    baseSource: FirPsiSourceElement?,
+    argument: FirExpression,
+    conventionName: Name
+): FirFunctionCall {
+    val operationReferenceSource = operationReference?.toFirSourceElement()
+    return buildFunctionCall {
+        source = baseSource
+        calleeReference = buildSimpleNamedReference {
+            source = operationReferenceSource
+            name = conventionName
+        }
+        explicitReceiver = this@createConventionCall
+        arguments += argument
+    }
+}
+
 fun generateAccessExpression(source: FirSourceElement?, name: Name): FirQualifiedAccessExpression =
     buildQualifiedAccessExpression {
         this.source = source

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

@@ -1373,6 +1373,10 @@ class RawFirBuilder(
                     return rightArgument.generateContainsOperation(
                         leftArgument, operationToken == NOT_IN, expression, expression.operationReference,
                     )
+                in OperatorConventions.COMPARISON_OPERATIONS ->
+                    return leftArgument.generateComparisonExpression(
+                        rightArgument, operationToken, expression, expression.operationReference,
+                    )
             }
             val conventionCallName = operationToken.toBinaryName()
             return if (conventionCallName != null || operationToken == IDENTIFIER) {

compiler/fir/psi2fir/testData/rawBuilder/expressions/branches.txt

@@ -1,7 +1,7 @@
 FILE: branches.kt
     public? final? fun foo(a: Int, b: Int): <implicit> {
         ^foo when () {
-            >(a#, b#) ->  {
+            CMP(>, a#.compareTo#(b#)) ->  {
                 a#
             }
             else ->  {
@@ -12,7 +12,7 @@ FILE: branches.kt
     }
     public? final? fun bar(a: Double, b: Double): Double {
         when () {
-            >(a#, b#) ->  {
+            CMP(>, a#.compareTo#(b#)) ->  {
                 println#(a#)
                 ^bar a#
             }
@@ -25,7 +25,7 @@ FILE: branches.kt
     }
     public? final? fun baz(a: Long, b: Long): Long {
         when () {
-            >(a#, b#) ->  {
+            CMP(>, a#.compareTo#(b#)) ->  {
                 println#(a#)
                 ^baz a#
             }

compiler/fir/psi2fir/testData/rawBuilder/expressions/while.txt

@@ -1,7 +1,7 @@
 FILE: while.kt
     public? final? fun foo(limit: Int): R|kotlin/Unit| {
         lvar k: <implicit> = IntegerLiteral(0)
-        some@while(<(k#, limit#)) {
+        some@while(CMP(<, k#.compareTo#(limit#))) {
             lval <unary>: <implicit> = k#
             k# = R|<local>/<unary>|.inc#()
             R|<local>/<unary>|
@@ -11,13 +11,13 @@ FILE: while.kt
                 k# = R|<local>/<unary>|.inc#()
                 R|<local>/<unary>|
                 when () {
-                    <(k#, limit#) ->  {
+                    CMP(<, k#.compareTo#(limit#)) ->  {
-                        break@@@[<(k#, limit#)]
+                        break@@@[CMP(<, k#.compareTo#(limit#))]
                     }
                 }
 
                 when () {
-                    >(k#, limit#) ->  {
+                    CMP(>, k#.compareTo#(limit#)) ->  {
                         continue@@@[==(k#, IntegerLiteral(13))]
                     }
                 }
@@ -35,5 +35,5 @@ FILE: while.kt
             R|<local>/<unary>|
             println#(k#)
         }
-        while(>=(k#, IntegerLiteral(0)))
+        while(CMP(>=, k#.compareTo#(IntegerLiteral(0))))
     }