JVM_IR: prefer to move, not copy, suspend lambda bodies

Copying breaks reflection metadata.

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/AddContinuationLowering.kt

@@ -54,12 +54,17 @@ internal val addContinuationPhase = makeIrFilePhase(
 private class AddContinuationLowering(private val context: JvmBackendContext) : FileLoweringPass {
     override fun lower(irFile: IrFile) {
         val suspendLambdas = findSuspendAndInlineLambdas(irFile)
-        addContinuationObjectAndContinuationParameterToSuspendFunctions(irFile)
         transformSuspendLambdasIntoContinuations(irFile, suspendLambdas)
-        addContinuationParameterToSuspendCallsAndUpdateNonLocalReturns(irFile, suspendLambdas)
+        // This should be done after converting lambdas into classes to avoid breaking the invariant that
+        // each lambda is referenced at most once while creating `$$forInline` methods.
+        addContinuationObjectAndContinuationParameterToSuspendFunctions(irFile)
+        addContinuationParameterToSuspendCalls(irFile)
+        // This should be done after adding continuation parameters so that `attributeContainerId` links
+        // inside `$$forInline` copies of `invokeSuspend` do not confuse the previous passes.
+        fillInvokeSuspendForInlineBodies(irFile)
     }
 
-    private fun addContinuationParameterToSuspendCallsAndUpdateNonLocalReturns(irFile: IrFile, suspendLambdas: List<SuspendLambdaInfo>) {
+    private fun addContinuationParameterToSuspendCalls(irFile: IrFile) {
         irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
             val functionStack = mutableListOf<IrFunction>()
 
@@ -75,13 +80,6 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                 return (super.visitCall(expression) as IrCall)
                     .createSuspendFunctionCallViewIfNeeded(context, caller)
             }
-
-            override fun visitReturn(expression: IrReturn): IrExpression {
-                val ret = super.visitReturn(expression) as IrReturn
-                val irFunction = expression.returnTargetSymbol.owner as? IrFunction ?: return ret
-                val target = suspendLambdas.find { it.function == irFunction }?.invokeSuspend ?: irFunction
-                return IrReturnImpl(ret.startOffset, ret.endOffset, ret.type, target.symbol, ret.value)
-            }
         })
     }
 
@@ -168,9 +166,8 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             }
             val invokeSuspend = addInvokeSuspendForLambda(info.function, parametersFields, receiverField)
             if (info.capturesCrossinline) {
-                addInvokeSuspendForInlineForLambda(invokeSuspend, info.function, parametersFields, receiverField)
+                addInvokeSuspendForInlineForLambda(invokeSuspend)
             }
-            info.invokeSuspend = invokeSuspend
             info.function.parentAsClass.declarations.remove(info.function)
             if (info.arity <= 1) {
                 val singleParameterField = receiverField ?: parametersWithoutArguments.singleOrNull()
@@ -201,15 +198,23 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             it.owner.name.asString() == INVOKE_SUSPEND_METHOD_NAME && it.owner.valueParameters.size == 1 &&
                     it.owner.valueParameters[0].type.isKotlinResult()
         }.owner
-        return addFunctionOverride(superMethod).also { it.copySuspendLambdaBodyFrom(irFunction, receiverField, fields) }
+        return addFunctionOverride(superMethod).apply {
+            val parametersToFields = mutableMapOf<IrValueParameter, IrField>()
+            assert(irFunction.dispatchReceiverParameter == null) // LocalDeclarationsLowering-generated methods are static
+            irFunction.extensionReceiverParameter?.let { parametersToFields[it] = receiverField!! }
+            irFunction.valueParameters.zip(fields).toMap(parametersToFields)
+            body = irFunction.moveBodyTo(this, mapOf())
+            body?.transformChildrenVoid(object : IrElementTransformerVoid() {
+                override fun visitGetValue(expression: IrGetValue): IrExpression {
+                    val field = parametersToFields[expression.symbol.owner] ?: return expression
+                    val receiver = IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, dispatchReceiverParameter!!.symbol)
+                    return IrGetFieldImpl(expression.startOffset, expression.endOffset, field.symbol, field.type, receiver)
+                }
+            })
+        }
     }
 
-    private fun IrClass.addInvokeSuspendForInlineForLambda(
+    private fun IrClass.addInvokeSuspendForInlineForLambda(invokeSuspend: IrFunction): IrFunction {
-        invokeSuspend: IrFunction,
-        irFunction: IrFunction,
-        fields: List<IrField>,
-        receiverField: IrField?
-    ): IrFunction {
         return addFunction(
             INVOKE_SUSPEND_METHOD_NAME + FOR_INLINE_SUFFIX,
             context.irBuiltIns.anyNType,
@@ -217,47 +222,19 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             origin = JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE
         ).apply {
             valueParameters += invokeSuspend.valueParameters.map { it.copyTo(this) }
-        }.also { it.copySuspendLambdaBodyFrom(irFunction, receiverField, fields) }
+        }
     }
 
-    private fun IrSimpleFunction.copySuspendLambdaBodyFrom(
+    private fun fillInvokeSuspendForInlineBodies(irFile: IrFile) {
-        irFunction: IrFunction,
+        irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
-        receiverField: IrField?,
+            override fun visitClass(declaration: IrClass): IrStatement = declaration.transformPostfix {
-        fields: List<IrField>
+                if (origin == JvmLoweredDeclarationOrigin.SUSPEND_LAMBDA) {
-    ) {
+                    for (function in functions) {
-        body = irFunction.body?.deepCopyWithSymbols(this)
+                        if (function.origin == JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE) {
-        body?.transformChildrenVoid(object : IrElementTransformerVoid() {
+                            function.body = functions.single { it.name.asString() == "invokeSuspend" }.copyBodyTo(function)
-            override fun visitGetValue(expression: IrGetValue): IrExpression {
+                        }
-                if (expression.symbol.owner == irFunction.extensionReceiverParameter) {
-                    assert(receiverField != null)
-                    return IrGetFieldImpl(
-                        expression.startOffset,
-                        expression.endOffset,
-                        receiverField!!.symbol,
-                        receiverField.type
-                    ).also {
-                        it.receiver =
-                            IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, dispatchReceiverParameter!!.symbol)
                     }
-                } else if (expression.symbol.owner == irFunction.dispatchReceiverParameter) {
-                    return IrGetValueImpl(expression.startOffset, expression.endOffset, dispatchReceiverParameter!!.symbol)
                 }
-                val field = fields.find { it.name == expression.symbol.owner.name } ?: return expression
-                return IrGetFieldImpl(expression.startOffset, expression.endOffset, field.symbol, field.type).also {
-                    it.receiver = IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, dispatchReceiverParameter!!.symbol)
-                }
-            }
-
-            // If the suspend lambda body contains declarations of other classes (for other lambdas),
-            // do not rewrite those. In particular, that could lead to rewriting of returns in nested
-            // lambdas to unintended non-local returns.
-            override fun visitClass(declaration: IrClass): IrStatement {
-                return declaration
-            }
-
-            override fun visitReturn(expression: IrReturn): IrExpression {
-                val ret = super.visitReturn(expression) as IrReturn
-                return IrReturnImpl(ret.startOffset, ret.endOffset, ret.type, symbol, ret.value)
             }
         })
     }
@@ -619,12 +596,9 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                             else JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
                     }.apply {
                         annotations += view.annotations.map { it.deepCopyWithSymbols(this) }
-                        copyTypeParameters(view.typeParameters)
+                        copyParameterDeclarationsFrom(view)
-                        dispatchReceiverParameter = view.dispatchReceiverParameter?.copyTo(this)
-                        extensionReceiverParameter = view.extensionReceiverParameter?.copyTo(this)
-                        valueParameters += view.valueParameters.map { it.copyTo(this) }
-                        body = view.copyBodyTo(this)
                         copyAttributes(view)
+                        body = view.copyBodyTo(this)
                     }
                     registerNewFunction(newFunction)
                 }
@@ -728,7 +702,6 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         val capturesCrossinline: Boolean
     ) {
         lateinit var constructor: IrConstructor
-        lateinit var invokeSuspend: IrFunction
     }
 }
 

compiler/testData/codegen/box/reflection/lambdaClasses/reflectOnLambdaInSuspend.kt

@@ -0,0 +1,19 @@
+// IGNORE_BACKEND_FIR: JVM_IR
+// TARGET_BACKEND: JVM
+// WITH_REFLECT
+// WITH_COROUTINES
+// FILE: a.kt
+
+import helpers.*
+import kotlin.coroutines.*
+import kotlin.reflect.jvm.reflect
+
+suspend fun f() = { OK: String -> }
+
+fun box(): String {
+    lateinit var x: (String) -> Unit
+    suspend {
+        x = f()
+    }.startCoroutine(EmptyContinuation)
+    return x.reflect()?.parameters?.singleOrNull()?.name ?: "null"
+}

compiler/testData/codegen/box/reflection/lambdaClasses/reflectOnLambdaInSuspendLambda.kt

@@ -0,0 +1,17 @@
+// IGNORE_BACKEND_FIR: JVM_IR
+// TARGET_BACKEND: JVM
+// WITH_REFLECT
+// WITH_COROUTINES
+// FILE: a.kt
+
+import helpers.*
+import kotlin.coroutines.*
+import kotlin.reflect.jvm.reflect
+
+fun box(): String {
+    lateinit var x: (String) -> Unit
+    suspend {
+        x = { OK: String -> }
+    }.startCoroutine(EmptyContinuation)
+    return x.reflect()?.parameters?.singleOrNull()?.name ?: "null"
+}

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -24445,6 +24445,16 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             public void testReflectOnLambdaInStaticField() throws Exception {
                 runTest("compiler/testData/codegen/box/reflection/lambdaClasses/reflectOnLambdaInStaticField.kt");
             }
+
+            @TestMetadata("reflectOnLambdaInSuspend.kt")
+            public void testReflectOnLambdaInSuspend() throws Exception {
+                runTest("compiler/testData/codegen/box/reflection/lambdaClasses/reflectOnLambdaInSuspend.kt");
+            }
+
+            @TestMetadata("reflectOnLambdaInSuspendLambda.kt")
+            public void testReflectOnLambdaInSuspendLambda() throws Exception {
+                runTest("compiler/testData/codegen/box/reflection/lambdaClasses/reflectOnLambdaInSuspendLambda.kt");
+            }
         }
 
         @TestMetadata("compiler/testData/codegen/box/reflection/mapping")

compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -23262,6 +23262,16 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             public void testReflectOnLambdaInStaticField() throws Exception {
                 runTest("compiler/testData/codegen/box/reflection/lambdaClasses/reflectOnLambdaInStaticField.kt");
             }
+
+            @TestMetadata("reflectOnLambdaInSuspend.kt")
+            public void testReflectOnLambdaInSuspend() throws Exception {
+                runTest("compiler/testData/codegen/box/reflection/lambdaClasses/reflectOnLambdaInSuspend.kt");
+            }
+
+            @TestMetadata("reflectOnLambdaInSuspendLambda.kt")
+            public void testReflectOnLambdaInSuspendLambda() throws Exception {
+                runTest("compiler/testData/codegen/box/reflection/lambdaClasses/reflectOnLambdaInSuspendLambda.kt");
+            }
         }
 
         @TestMetadata("compiler/testData/codegen/box/reflection/mapping")

compiler/tests/org/jetbrains/kotlin/codegen/ir/FirBlackBoxCodegenTestGenerated.java

@@ -22954,6 +22954,16 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
             public void testReflectOnLambdaInStaticField() throws Exception {
                 runTest("compiler/testData/codegen/box/reflection/lambdaClasses/reflectOnLambdaInStaticField.kt");
             }
+
+            @TestMetadata("reflectOnLambdaInSuspend.kt")
+            public void testReflectOnLambdaInSuspend() throws Exception {
+                runTest("compiler/testData/codegen/box/reflection/lambdaClasses/reflectOnLambdaInSuspend.kt");
+            }
+
+            @TestMetadata("reflectOnLambdaInSuspendLambda.kt")
+            public void testReflectOnLambdaInSuspendLambda() throws Exception {
+                runTest("compiler/testData/codegen/box/reflection/lambdaClasses/reflectOnLambdaInSuspendLambda.kt");
+            }
         }
 
         @TestMetadata("compiler/testData/codegen/box/reflection/mapping")

compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -22954,6 +22954,16 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             public void testReflectOnLambdaInStaticField() throws Exception {
                 runTest("compiler/testData/codegen/box/reflection/lambdaClasses/reflectOnLambdaInStaticField.kt");
             }
+
+            @TestMetadata("reflectOnLambdaInSuspend.kt")
+            public void testReflectOnLambdaInSuspend() throws Exception {
+                runTest("compiler/testData/codegen/box/reflection/lambdaClasses/reflectOnLambdaInSuspend.kt");
+            }
+
+            @TestMetadata("reflectOnLambdaInSuspendLambda.kt")
+            public void testReflectOnLambdaInSuspendLambda() throws Exception {
+                runTest("compiler/testData/codegen/box/reflection/lambdaClasses/reflectOnLambdaInSuspendLambda.kt");
+            }
         }
 
         @TestMetadata("compiler/testData/codegen/box/reflection/mapping")