JVM_IR KT-47984 don't move inplace arguments with suspension points

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/InplaceArgumentsMethodTransformer.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.codegen.inline
 
+import org.jetbrains.kotlin.codegen.optimization.boxing.isMethodInsnWith
 import org.jetbrains.kotlin.codegen.optimization.common.InsnSequence
 import org.jetbrains.kotlin.codegen.optimization.transformer.MethodTransformer
 import org.jetbrains.org.objectweb.asm.Label
@@ -19,6 +20,7 @@ class InplaceArgumentsMethodTransformer : MethodTransformer() {
 
             collectStartToEnd(methodContext)
             collectLvtEntryInstructions(methodContext)
+            collectSuspensionPoints(methodContext)
 
             transformMethod(methodContext)
             updateLvtEntriesForMovedInstructions(methodContext)
@@ -36,6 +38,7 @@ class InplaceArgumentsMethodTransformer : MethodTransformer() {
         val startArgToEndArg = HashMap<AbstractInsnNode, AbstractInsnNode>()
         val lvtEntryForInstruction = HashMap<AbstractInsnNode, LocalVariableNode>()
         val varInstructionMoved = HashMap<AbstractInsnNode, CallContext>()
+        val suspensionJumpLabels = HashSet<LabelNode>()
     }
 
     private class CallContext(
@@ -185,6 +188,37 @@ class InplaceArgumentsMethodTransformer : MethodTransformer() {
         }
     }
 
+    private fun collectSuspensionPoints(methodContext: MethodContext) {
+        val insnList = methodContext.methodNode.instructions
+        var insn = insnList.first
+        while (
+            !insn.isMethodInsnWith(Opcodes.INVOKESTATIC) {
+                owner == "kotlin/coroutines/intrinsics/IntrinsicsKt" &&
+                        name == "getCOROUTINE_SUSPENDED" &&
+                        desc == "()Ljava/lang/Object;"
+            }
+        ) {
+            insn = insn.next ?: return
+        }
+
+        // Find a first TABLESWITCH and record its jump destinations
+        while (insn != null) {
+            if (insn.opcode != Opcodes.TABLESWITCH || insn.previous.opcode != Opcodes.GETFIELD) {
+                insn = insn.next
+                continue
+            }
+            val getFiendInsn = insn.previous as FieldInsnNode
+            if (getFiendInsn.name != "label" || getFiendInsn.desc != "I") {
+                insn = insn.next
+                continue
+            }
+            val tableSwitchInsn = insn as TableSwitchInsnNode
+            methodContext.suspensionJumpLabels.addAll(tableSwitchInsn.labels)
+            methodContext.suspensionJumpLabels.add(tableSwitchInsn.dflt)
+            return
+        }
+    }
+
     private fun transformMethod(methodContext: MethodContext) {
         for (call in methodContext.calls) {
             transformCall(methodContext, call)
@@ -205,7 +239,7 @@ class InplaceArgumentsMethodTransformer : MethodTransformer() {
         // If an inplace argument contains a non-local jump,
         // moving such argument inside inline function body can interfere with stack normalization.
         // TODO investigate complex cases
-        if (callContext.args.any { it.hasNonLocalJump() }) {
+        if (callContext.args.any { it.isUnsafeToMove(methodContext) }) {
             // Do not transform such call, just strip call and argument markers.
             val insnList = methodContext.methodNode.instructions
             for (arg in callContext.args) {
@@ -220,10 +254,13 @@ class InplaceArgumentsMethodTransformer : MethodTransformer() {
         moveInplaceArgumentsFromStoresToLoads(methodContext, callContext)
     }
 
-    private fun ArgContext.hasNonLocalJump(): Boolean {
+    private fun ArgContext.isUnsafeToMove(methodContext: MethodContext): Boolean {
         val argInsns = InsnSequence(this.argStartMarker, this.argEndMarker)
         val localLabels = argInsns.filterTo(HashSet()) { it is LabelNode }
-        return argInsns.any { insn -> insn.opcode == Opcodes.GOTO && (insn as JumpInsnNode).label !in localLabels }
+        return argInsns.any { insn ->
+            insn in methodContext.suspensionJumpLabels ||
+                    insn.opcode == Opcodes.GOTO && (insn as JumpInsnNode).label !in localLabels
+        }
     }
 
     private fun moveInplaceArgumentsFromStoresToLoads(methodContext: MethodContext, callContext: CallContext) {

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -18251,6 +18251,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
         public void testMutableCollectionPlusAssign() throws Exception {
             runTest("compiler/testData/codegen/box/inlineArgsInPlace/mutableCollectionPlusAssign.kt");
         }
+
+        @Test
+        @TestMetadata("suspensionPointInsideArgument.kt")
+        public void testSuspensionPointInsideArgument() throws Exception {
+            runTest("compiler/testData/codegen/box/inlineArgsInPlace/suspensionPointInsideArgument.kt");
+        }
     }
 
     @Nested

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrInlineCodegen.kt

@@ -39,20 +39,33 @@ class IrInlineCodegen(
     InlineCodegen<ExpressionCodegen>(codegen, state, signature, typeParameterMappings, sourceCompiler, reifiedTypeInliner),
     IrInlineCallGenerator {
 
-    private val inlineArgumentsInPlace = run {
+    private val inlineArgumentsInPlace = canInlineArgumentsInPlace()
-        if (function.isInlineOnly() &&
+
-            !function.isSuspend &&
+    private fun canInlineArgumentsInPlace(): Boolean {
-            function.valueParameters.isNotEmpty() &&
+        if (!function.isInlineOnly()) return false
-            function.valueParameters.none { it.type.isFunction() || it.type.isSuspendFunction() }
+        if (function.isSuspend) return false
-        ) {
+
-            val methodNode = sourceCompiler.compileInlineFunction(signature).node
+        var actualParametersCount = function.valueParameters.size
-            if (canInlineArgumentsInPlace(methodNode)) {
+        function.dispatchReceiverParameter?.let { dispatchReceiverParameter ->
-                return@run true
+            ++actualParametersCount
-            }
+            if (dispatchReceiverParameter.isFunctionOrSuspendFunction())
+                return false
         }
-        false
+        function.extensionReceiverParameter?.let { extensionReceiverParameter ->
+            ++actualParametersCount
+            if (extensionReceiverParameter.isFunctionOrSuspendFunction())
+                return false
+        }
+        if (actualParametersCount == 0)
+            return false
+        if (function.valueParameters.any { it.isFunctionOrSuspendFunction() })
+            return false
+
+        return canInlineArgumentsInPlace(sourceCompiler.compileInlineFunction(jvmSignature).node)
     }
 
+    private fun IrValueParameter.isFunctionOrSuspendFunction() = type.isFunction() || type.isSuspendFunction()
+
     override fun beforeCallStart() {
         if (inlineArgumentsInPlace) {
             codegen.visitor.addInplaceCallStartMarker()

compiler/testData/codegen/box/inlineArgsInPlace/suspensionPointInsideArgument.kt

@@ -0,0 +1,25 @@
+// IGNORE_BACKEND: WASM
+// WITH_RUNTIME
+
+import kotlin.coroutines.*
+
+fun runs(f: suspend () -> String): String {
+    var result: String? = null
+    f.startCoroutine(
+        Continuation(EmptyCoroutineContext) {
+            result = it.getOrThrow()
+        }
+    )
+    return result ?: "Fail"
+}
+
+suspend fun suspendListOf(s: String) = listOf(s)
+
+val strings: MutableCollection<String> = ArrayList()
+
+fun box(): String {
+    return runs {
+        strings += suspendListOf("OK")
+        strings.iterator().next()
+    }
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java

@@ -18107,6 +18107,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
         public void testMutableCollectionPlusAssign() throws Exception {
             runTest("compiler/testData/codegen/box/inlineArgsInPlace/mutableCollectionPlusAssign.kt");
         }
+
+        @Test
+        @TestMetadata("suspensionPointInsideArgument.kt")
+        public void testSuspensionPointInsideArgument() throws Exception {
+            runTest("compiler/testData/codegen/box/inlineArgsInPlace/suspensionPointInsideArgument.kt");
+        }
     }
 
     @Nested

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -18251,6 +18251,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
         public void testMutableCollectionPlusAssign() throws Exception {
             runTest("compiler/testData/codegen/box/inlineArgsInPlace/mutableCollectionPlusAssign.kt");
         }
+
+        @Test
+        @TestMetadata("suspensionPointInsideArgument.kt")
+        public void testSuspensionPointInsideArgument() throws Exception {
+            runTest("compiler/testData/codegen/box/inlineArgsInPlace/suspensionPointInsideArgument.kt");
+        }
     }
 
     @Nested

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

@@ -15008,6 +15008,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
         public void testMutableCollectionPlusAssign() throws Exception {
             runTest("compiler/testData/codegen/box/inlineArgsInPlace/mutableCollectionPlusAssign.kt");
         }
+
+        @TestMetadata("suspensionPointInsideArgument.kt")
+        public void testSuspensionPointInsideArgument() throws Exception {
+            runTest("compiler/testData/codegen/box/inlineArgsInPlace/suspensionPointInsideArgument.kt");
+        }
     }
 
     @TestMetadata("compiler/testData/codegen/box/inlineClasses")

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/es6/semantics/IrJsCodegenBoxES6TestGenerated.java

@@ -13097,6 +13097,11 @@ public class IrJsCodegenBoxES6TestGenerated extends AbstractIrJsCodegenBoxES6Tes
         public void testMutableCollectionPlusAssign() throws Exception {
             runTest("compiler/testData/codegen/box/inlineArgsInPlace/mutableCollectionPlusAssign.kt");
         }
+
+        @TestMetadata("suspensionPointInsideArgument.kt")
+        public void testSuspensionPointInsideArgument() throws Exception {
+            runTest("compiler/testData/codegen/box/inlineArgsInPlace/suspensionPointInsideArgument.kt");
+        }
     }
 
     @TestMetadata("compiler/testData/codegen/box/inlineClasses")

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/ir/semantics/IrJsCodegenBoxTestGenerated.java

@@ -12503,6 +12503,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
         public void testMutableCollectionPlusAssign() throws Exception {
             runTest("compiler/testData/codegen/box/inlineArgsInPlace/mutableCollectionPlusAssign.kt");
         }
+
+        @TestMetadata("suspensionPointInsideArgument.kt")
+        public void testSuspensionPointInsideArgument() throws Exception {
+            runTest("compiler/testData/codegen/box/inlineArgsInPlace/suspensionPointInsideArgument.kt");
+        }
     }
 
     @TestMetadata("compiler/testData/codegen/box/inlineClasses")

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -12568,6 +12568,11 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
         public void testMutableCollectionPlusAssign() throws Exception {
             runTest("compiler/testData/codegen/box/inlineArgsInPlace/mutableCollectionPlusAssign.kt");
         }
+
+        @TestMetadata("suspensionPointInsideArgument.kt")
+        public void testSuspensionPointInsideArgument() throws Exception {
+            runTest("compiler/testData/codegen/box/inlineArgsInPlace/suspensionPointInsideArgument.kt");
+        }
     }
 
     @TestMetadata("compiler/testData/codegen/box/inlineClasses")

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/wasm/semantics/IrCodegenBoxWasmTestGenerated.java

@@ -6658,6 +6658,11 @@ public class IrCodegenBoxWasmTestGenerated extends AbstractIrCodegenBoxWasmTest
         public void testMutableCollectionPlusAssign() throws Exception {
             runTest("compiler/testData/codegen/box/inlineArgsInPlace/mutableCollectionPlusAssign.kt");
         }
+
+        @TestMetadata("suspensionPointInsideArgument.kt")
+        public void testSuspensionPointInsideArgument() throws Exception {
+            runTest("compiler/testData/codegen/box/inlineArgsInPlace/suspensionPointInsideArgument.kt");
+        }
     }
 
     @TestMetadata("compiler/testData/codegen/box/inlineClasses")