From 943dde0cbb48d0a01b1a939277839f6da2c2be3f Mon Sep 17 00:00:00 2001
From: Alexander Udalov <alexander.udalov@jetbrains.com>
Date: Thu, 15 Jun 2023 18:22:35 +0200
Subject: [PATCH] JVM IR: do not copy call in
 ReplaceKFunctionInvokeWithFunctionInvoke

This is just a refactoring/optimization that makes use of the fact that
IrCall.symbol is now mutable.
---
 ...eplaceKFunctionInvokeWithFunctionInvoke.kt | 47 +++++++++----------
 1 file changed, 21 insertions(+), 26 deletions(-)

diff --git a/compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/lower/ReplaceKFunctionInvokeWithFunctionInvoke.kt b/compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/lower/ReplaceKFunctionInvokeWithFunctionInvoke.kt
index 29ebff1dec2..2311abe7d74 100644
--- a/compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/lower/ReplaceKFunctionInvokeWithFunctionInvoke.kt
+++ b/compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/lower/ReplaceKFunctionInvokeWithFunctionInvoke.kt
@@ -8,20 +8,18 @@ package org.jetbrains.kotlin.backend.jvm.lower
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.expressions.IrCall
-import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrTypeOperator
-import org.jetbrains.kotlin.ir.expressions.copyTypeArgumentsFrom
-import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrTypeOperatorCallImpl
 import org.jetbrains.kotlin.ir.util.defaultType
 import org.jetbrains.kotlin.ir.util.isKFunction
 import org.jetbrains.kotlin.ir.util.isKSuspendFunction
 import org.jetbrains.kotlin.ir.util.parentAsClass
-import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
-import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
+import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 import org.jetbrains.kotlin.util.OperatorNameConventions
 
 internal val replaceKFunctionInvokeWithFunctionInvokePhase = makeIrFilePhase<JvmBackendContext>(
@@ -35,35 +33,32 @@ internal val replaceKFunctionInvokeWithFunctionInvokePhase = makeIrFilePhase<Jvm
  * with `SuspendFunction{n}.invoke`. This is needed because normally the type e.g. `kotlin.reflect.KFunction2` is mapped to
  * `kotlin.reflect.KFunction` (a real class, without arity), which doesn't have the corresponding `invoke`.
  */
-private class ReplaceKFunctionInvokeWithFunctionInvoke : FileLoweringPass, IrElementTransformerVoid() {
+private class ReplaceKFunctionInvokeWithFunctionInvoke : FileLoweringPass, IrElementVisitorVoid {
     override fun lower(irFile: IrFile) {
-        irFile.transformChildrenVoid(this)
+        irFile.acceptChildrenVoid(this)
     }
 
-    override fun visitCall(expression: IrCall): IrExpression {
-        val callee = expression.symbol.owner
-        if (callee.name != OperatorNameConventions.INVOKE) return super.visitCall(expression)
+    override fun visitElement(element: IrElement) {
+        element.acceptChildren(this, null)
+    }
 
-        val parentClass = callee.parent as? IrClass ?: return super.visitCall(expression)
+    override fun visitCall(expression: IrCall) {
+        expression.acceptChildren(this, null)
+
+        val callee = expression.symbol.owner
+        if (callee.name != OperatorNameConventions.INVOKE) return
+
+        val parentClass = callee.parent as? IrClass ?: return
         if (!parentClass.defaultType.isKFunction() && !parentClass.defaultType.isKSuspendFunction()) {
             implicitCastKFunctionReceiverIntoFunctionIfNeeded(expression, parentClass)
-            return super.visitCall(expression)
+            return
         }
 
         // The single overridden function of KFunction{n}.invoke must be Function{n}.invoke.
-        val newCallee = callee.overriddenSymbols.single()
-        return expression.run {
-            IrCallImpl.fromSymbolOwner(startOffset, endOffset, type, newCallee).apply {
-                copyTypeArgumentsFrom(expression)
-                dispatchReceiver = expression.dispatchReceiver?.transform(this@ReplaceKFunctionInvokeWithFunctionInvoke, null)?.let {
-                    val newType = newCallee.owner.parentAsClass.defaultType
-                    IrTypeOperatorCallImpl(startOffset, endOffset, newType, IrTypeOperator.IMPLICIT_CAST, newType, it)
-                }
-                extensionReceiver = expression.extensionReceiver?.transform(this@ReplaceKFunctionInvokeWithFunctionInvoke, null)
-                for (i in 0 until valueArgumentsCount) {
-                    putValueArgument(i, expression.getValueArgument(i)?.transform(this@ReplaceKFunctionInvokeWithFunctionInvoke, null))
-                }
-            }
+        expression.symbol = callee.overriddenSymbols.single()
+        expression.dispatchReceiver = expression.dispatchReceiver?.let {
+            val newType = expression.symbol.owner.parentAsClass.defaultType
+            IrTypeOperatorCallImpl(expression.startOffset, expression.endOffset, newType, IrTypeOperator.IMPLICIT_CAST, newType, it)
         }
     }
 
@@ -74,7 +69,7 @@ private class ReplaceKFunctionInvokeWithFunctionInvoke : FileLoweringPass, IrEle
             val newType = parentClass.defaultType
 
             expression.dispatchReceiver = IrTypeOperatorCallImpl(
-                expression.startOffset, expression.endOffset, newType, IrTypeOperator.IMPLICIT_CAST, newType, receiver.transform(this, null)
+                expression.startOffset, expression.endOffset, newType, IrTypeOperator.IMPLICIT_CAST, newType, receiver
             )
         }
     }