From 14150ea78d61fa5a307c04ebdd372453ce376ad5 Mon Sep 17 00:00:00 2001
From: Tamas Vajk <tamasvajk@github.com>
Date: Fri, 8 Nov 2024 14:56:49 +0100
Subject: [PATCH] KE2: Extract `compareTo` calls for binary comparisons

---
 .../src/main/kotlin/entities/Expression.kt    | 68 +++++++++++++++++--
 .../src/main/kotlin/entities/MethodCall.kt    |  2 +-
 2 files changed, 64 insertions(+), 6 deletions(-)

diff --git a/java/kotlin-extractor2/src/main/kotlin/entities/Expression.kt b/java/kotlin-extractor2/src/main/kotlin/entities/Expression.kt
index 3cc081261a7..bd5fda7a91e 100644
--- a/java/kotlin-extractor2/src/main/kotlin/entities/Expression.kt
+++ b/java/kotlin-extractor2/src/main/kotlin/entities/Expression.kt
@@ -373,7 +373,50 @@ private fun KotlinFileExtractor.extractBinaryExpression(
                 else -> TODO("error")
             }
         } else {
-            TODO("Extract lowered equivalent call, such as `a.compareTo(b) < 0`")
+            if (target == null) {
+                TODO()
+            }
+
+            // Extract lowered equivalent call, such as `a.compareTo(b) < 0`
+            val exprParent = parent.expr(expression, callable)
+
+            val id: Label<out DbBinaryexpr> = when (op) {
+                KtTokens.LT -> extractRawBinaryExpression(builtinTypes.boolean, exprParent, tw::writeExprs_ltexpr)
+                KtTokens.GT -> extractRawBinaryExpression(builtinTypes.boolean, exprParent, tw::writeExprs_gtexpr)
+                KtTokens.LTEQ -> extractRawBinaryExpression(builtinTypes.boolean, exprParent, tw::writeExprs_leexpr)
+                KtTokens.GTEQ -> extractRawBinaryExpression(builtinTypes.boolean, exprParent, tw::writeExprs_geexpr)
+                else -> TODO("error")
+            }
+
+            extractExprContext(id, tw.getLocation(expression), callable, exprParent.enclosingStmt)
+
+            extractRawMethodAccess(
+                target,
+                tw.getLocation(expression),
+                target.returnType,
+                callable,
+                id,
+                0,
+                exprParent.enclosingStmt,
+                if (target.isExtension) null else expression.left!!,
+                if (target.isExtension) expression.left!! else null,
+                listOf(expression.right!!)
+            )
+
+            extractConstantInteger(
+                "0",
+                builtinTypes.int,
+                0,
+                tw.getLocation(expression),
+                id,
+                1,
+                callable,
+                exprParent.enclosingStmt,
+                /*
+                OLD: KE1
+                                            overrideId = overrideId
+                */
+            )
         }
 
     } else {
@@ -394,6 +437,24 @@ private fun KaFunctionSymbol?.isBinaryPlus(): Boolean {
                     ))
 }
 
+context(KaSession)
+private fun <T : DbBinaryexpr> KotlinFileExtractor.extractRawBinaryExpression(
+    expressionType: KaType,
+    exprParent: KotlinFileExtractor.ExprParent,
+    extractExpression: (
+        id: Label<out T>,
+        typeid: Label<out DbType>,
+        parent: Label<out DbExprparent>,
+        idx: Int
+    ) -> Unit
+): Label<T> {
+    val id = tw.getFreshIdLabel<T>()
+    val type = useType(expressionType)
+    extractExpression(id, type.javaResult.id, exprParent.parent, exprParent.idx)
+    tw.writeExprsKotlinType(id, type.kotlinResult.id)
+    return id
+}
+
 context(KaSession)
 private fun <T : DbBinaryexpr> KotlinFileExtractor.extractBinaryExpression(
     expression: KtBinaryExpression,
@@ -406,11 +467,8 @@ private fun <T : DbBinaryexpr> KotlinFileExtractor.extractBinaryExpression(
         idx: Int
     ) -> Unit
 ) {
-    val id = tw.getFreshIdLabel<T>()
-    val type = useType(expression.expressionType)
     val exprParent = parent.expr(expression, callable)
-    extractExpression(id, type.javaResult.id, exprParent.parent, exprParent.idx)
-    tw.writeExprsKotlinType(id, type.kotlinResult.id)
+    val id = extractRawBinaryExpression(expression.expressionType!!, exprParent, extractExpression)
 
     extractExprContext(id, tw.getLocation(expression), callable, exprParent.enclosingStmt)
     extractExpressionExpr(expression.left!!, callable, id, 0, exprParent.enclosingStmt)
diff --git a/java/kotlin-extractor2/src/main/kotlin/entities/MethodCall.kt b/java/kotlin-extractor2/src/main/kotlin/entities/MethodCall.kt
index 00339634c3a..b279ac7fcfc 100644
--- a/java/kotlin-extractor2/src/main/kotlin/entities/MethodCall.kt
+++ b/java/kotlin-extractor2/src/main/kotlin/entities/MethodCall.kt
@@ -116,7 +116,7 @@ private fun KotlinFileExtractor.getCalleeMethodId(callTarget: KaFunctionSymbol):
 }
 
 context(KaSession)
-private fun KotlinFileExtractor.extractRawMethodAccess(
+fun KotlinFileExtractor.extractRawMethodAccess(
     callTarget: KaFunctionSymbol,
     locId: Label<DbLocation>,
     returnType: KaType,