From 255d5c994292481cbb890e0065e99be028623e4a Mon Sep 17 00:00:00 2001
From: Tamas Vajk <tamasvajk@github.com>
Date: Tue, 15 Oct 2024 11:27:56 +0200
Subject: [PATCH] KE2: Extract binary operators on numeric types

---
 .../src/main/kotlin/KotlinFileExtractor.kt    | 25 ------
 .../src/main/kotlin/entities/Expression.kt    | 81 ++++++++++++-------
 2 files changed, 52 insertions(+), 54 deletions(-)

diff --git a/java/kotlin-extractor2/src/main/kotlin/KotlinFileExtractor.kt b/java/kotlin-extractor2/src/main/kotlin/KotlinFileExtractor.kt
index 74b49a4993b..7a94685ba6a 100644
--- a/java/kotlin-extractor2/src/main/kotlin/KotlinFileExtractor.kt
+++ b/java/kotlin-extractor2/src/main/kotlin/KotlinFileExtractor.kt
@@ -2981,11 +2981,6 @@ OLD: KE1
                 when {
                     isNumericFunction(
                         target,
-                        "plus",
-                        "minus",
-                        "times",
-                        "div",
-                        "rem",
                         "and",
                         "or",
                         "xor",
@@ -2996,26 +2991,6 @@ OLD: KE1
                         val type = useType(c.type)
                         val id: Label<out DbExpr> =
                             when (val targetName = target.name.asString()) {
-                                "minus" -> {
-                                    val id = tw.getFreshIdLabel<DbSubexpr>()
-                                    tw.writeExprs_subexpr(id, type.javaResult.id, parent, idx)
-                                    id
-                                }
-                                "times" -> {
-                                    val id = tw.getFreshIdLabel<DbMulexpr>()
-                                    tw.writeExprs_mulexpr(id, type.javaResult.id, parent, idx)
-                                    id
-                                }
-                                "div" -> {
-                                    val id = tw.getFreshIdLabel<DbDivexpr>()
-                                    tw.writeExprs_divexpr(id, type.javaResult.id, parent, idx)
-                                    id
-                                }
-                                "rem" -> {
-                                    val id = tw.getFreshIdLabel<DbRemexpr>()
-                                    tw.writeExprs_remexpr(id, type.javaResult.id, parent, idx)
-                                    id
-                                }
                                 "and" -> {
                                     val id = tw.getFreshIdLabel<DbAndbitexpr>()
                                     tw.writeExprs_andbitexpr(id, type.javaResult.id, parent, idx)
diff --git a/java/kotlin-extractor2/src/main/kotlin/entities/Expression.kt b/java/kotlin-extractor2/src/main/kotlin/entities/Expression.kt
index 15b8028a9f0..3c50a3ac855 100644
--- a/java/kotlin-extractor2/src/main/kotlin/entities/Expression.kt
+++ b/java/kotlin-extractor2/src/main/kotlin/entities/Expression.kt
@@ -297,37 +297,60 @@ private fun KotlinFileExtractor.extractBinaryExpression(
     val op = expression.operationToken
     val target = expression.resolveCallTarget()?.symbol
 
-    when (op) {
-        KtTokens.PLUS -> {
-            if (target == null) {
-                TODO()
-            }
-
-            if (target.isNumericWithName("plus") ||
-                target.hasName("kotlin", "String", "plus") ||
-                target.hasMatchingNames(
-                    CallableId(FqName("kotlin"), null, Name.identifier("plus")),
-                    ClassId(FqName("kotlin"), Name.identifier("String")),
-                    nullability = KaTypeNullability.NULLABLE,
-                )
-            ) {
-                val id = tw.getFreshIdLabel<DbAddexpr>()
-                val type = useType(expression.expressionType)
-                val exprParent = parent.expr(expression, callable)
-                tw.writeExprs_addexpr(id, type.javaResult.id, exprParent.parent, exprParent.idx)
-                tw.writeExprsKotlinType(id, type.kotlinResult.id)
-
-                extractExprContext(id, tw.getLocation(expression), callable, exprParent.enclosingStmt)
-                extractExpressionExpr(expression.left!!, callable, id, 0, exprParent.enclosingStmt)
-                extractExpressionExpr(expression.right!!, callable, id, 1, exprParent.enclosingStmt)
-            } else {
-                TODO("Extract as method call")
-            }
-        }
-
-        else -> TODO()
+    if (target == null) {
+        TODO()
     }
 
+    if (op == KtTokens.PLUS && target.isBinaryPlus()) {
+        extractBinaryExpression(expression, callable, parent, tw::writeExprs_addexpr)
+    } else if (op == KtTokens.MINUS && target.isNumericWithName("minus")) {
+        extractBinaryExpression(expression, callable, parent, tw::writeExprs_subexpr)
+    } else if (op == KtTokens.MUL && target.isNumericWithName("times")) {
+        extractBinaryExpression(expression, callable, parent, tw::writeExprs_mulexpr)
+    } else if (op == KtTokens.DIV && target.isNumericWithName("div")) {
+        extractBinaryExpression(expression, callable, parent, tw::writeExprs_divexpr)
+    } else if (op == KtTokens.PERC && target.isNumericWithName("rem")) {
+        extractBinaryExpression(expression, callable, parent, tw::writeExprs_remexpr)
+    } else {
+        if (op !in listOf(KtTokens.PLUS, KtTokens.MINUS, KtTokens.MUL, KtTokens.DIV, KtTokens.PERC)) {
+            TODO("Unhandled binary op")
+        }
+
+        TODO("Extract as method call")
+    }
+}
+
+private fun KaFunctionSymbol.isBinaryPlus(): Boolean {
+    return this.isNumericWithName("plus") ||
+            this.hasName("kotlin", "String", "plus") ||
+            this.hasMatchingNames(
+                CallableId(FqName("kotlin"), null, Name.identifier("plus")),
+                ClassId(FqName("kotlin"), Name.identifier("String")),
+                nullability = KaTypeNullability.NULLABLE,
+            )
+}
+
+context(KaSession)
+private fun <T : DbBinaryexpr> KotlinFileExtractor.extractBinaryExpression(
+    expression: KtBinaryExpression,
+    callable: Label<out DbCallable>,
+    parent: StmtExprParent,
+    extractExpression: (
+        id: Label<out T>,
+        typeid: Label<out DbType>,
+        parent: Label<out DbExprparent>,
+        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)
+
+    extractExprContext(id, tw.getLocation(expression), callable, exprParent.enclosingStmt)
+    extractExpressionExpr(expression.left!!, callable, id, 0, exprParent.enclosingStmt)
+    extractExpressionExpr(expression.right!!, callable, id, 1, exprParent.enclosingStmt)
 }
 
 context(KaSession)