Replace ALOAD 0 with fake continuation if suspension point is default

method

 #KT-25922 Fixed
 #KT-26608 Fixed
 #KT-26658 Fixed

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

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.codegen.inline.FieldRemapper.Companion.foldName
 import org.jetbrains.kotlin.codegen.intrinsics.IntrinsicMethods
 import org.jetbrains.kotlin.codegen.optimization.ApiVersionCallsPreprocessingMethodTransformer
 import org.jetbrains.kotlin.codegen.optimization.FixStackWithLabelNormalizationMethodTransformer
+import org.jetbrains.kotlin.codegen.optimization.common.ControlFlowGraph
 import org.jetbrains.kotlin.codegen.optimization.common.InsnSequence
 import org.jetbrains.kotlin.codegen.optimization.common.asSequence
 import org.jetbrains.kotlin.codegen.optimization.common.isMeaningful
@@ -23,7 +24,7 @@ import org.jetbrains.kotlin.codegen.optimization.fixStack.top
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.ParameterDescriptor
 import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
-import org.jetbrains.kotlin.resolve.isInlineClassType
+import org.jetbrains.kotlin.load.java.JvmAbi
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes.OBJECT_TYPE
 import org.jetbrains.kotlin.types.KotlinType
@@ -640,26 +641,78 @@ class MethodInliner(
     private fun replaceContinuationAccessesWithFakeContinuationsIfNeeded(processingNode: MethodNode) {
         val lambdaInfo = inliningContext.lambdaInfo ?: return
         if (!lambdaInfo.invokeMethodDescriptor.isSuspend) return
+        val sources = analyzeMethodNodeBeforeInline(processingNode)
+        val cfg = ControlFlowGraph.build(processingNode)
         val aload0s = processingNode.instructions.asSequence().filter { it.opcode == Opcodes.ALOAD && it.safeAs<VarInsnNode>()?.`var` == 0 }
-        // Expected pattern here:
+
-        //     ALOAD 0
+        val visited = hashSetOf<AbstractInsnNode>()
-        //     ICONST_0
+        fun findMeaningfulSuccs(insn: AbstractInsnNode): Collection<AbstractInsnNode> {
-        //     INVOKESTATIC InlineMarker.mark
+            if (!visited.add(insn)) return emptySet()
-        //     INVOKE* suspendingFunction(..., Continuation;)Ljava/lang/Object;
+            val res = hashSetOf<AbstractInsnNode>()
-        val continuationAsParameterAload0s =
+            for (succIndex in cfg.getSuccessorsIndices(insn)) {
-            aload0s.filter { it.next?.next?.let(::isBeforeSuspendMarker) == true && isSuspendCall(it.next?.next?.next) }
+                val succ = processingNode.instructions[succIndex]
-        replaceContinuationsWithFakeOnes(continuationAsParameterAload0s, processingNode)
+                if (succ.isMeaningful) res.add(succ)
+                else res.addAll(findMeaningfulSuccs(succ))
+            }
+            return res
+        }
+
+        // After inlining suspendCoroutineUninterceptedOrReturn there will be suspension point, which is not a MethodInsnNode.
+        // So, it is incorrect to expect MethodInsnNodes only
+        val suspensionPoints = processingNode.instructions.asSequence()
+            .filter { isBeforeSuspendMarker(it) }
+            .flatMap { findMeaningfulSuccs(it).asSequence() }
+            .filter { it is MethodInsnNode }
+
+        val toReplace = hashSetOf<AbstractInsnNode>()
+        for (suspensionPoint in suspensionPoints) {
+            assert(suspensionPoint is MethodInsnNode) {
+                "suspensionPoint shall be MethodInsnNode, but instead $suspensionPoint"
+            }
+            suspensionPoint as MethodInsnNode
+            assert(Type.getReturnType(suspensionPoint.desc) == OBJECT_TYPE) {
+                "suspensionPoint shall return $OBJECT_TYPE, but returns ${Type.getReturnType(suspensionPoint.desc)}"
+            }
+            val frame = sources[processingNode.instructions.indexOf(suspensionPoint)] ?: continue
+            val paramTypes = Type.getArgumentTypes(suspensionPoint.desc)
+            if (suspensionPoint.name.endsWith(JvmAbi.DEFAULT_PARAMS_IMPL_SUFFIX)) {
+                // Expected pattern here:
+                //     ALOAD 0
+                //     (ICONST or other integers creating instruction)
+                //     (ACONST_NULL or ALOAD)
+                //     ICONST_0
+                //     INVOKESTATIC InlineMarker.mark
+                //     INVOKE* suspendingFunction$default(..., Continuation;ILjava/lang/Object)Ljava/lang/Object;
+                assert(paramTypes.size >= 3) {
+                    "${suspensionPoint.name}${suspensionPoint.desc} shall have 3+ parameters"
+                }
+            } else {
+                // Expected pattern here:
+                //     ALOAD 0
+                //     ICONST_0
+                //     INVOKESTATIC InlineMarker.mark
+                //     INVOKE* suspendingFunction(..., Continuation;)Ljava/lang/Object;
+                assert(paramTypes.isNotEmpty()) {
+                    "${suspensionPoint.name}${suspensionPoint.desc} shall have 1+ parameters"
+                }
+            }
+            paramTypes.reversed().asSequence().withIndex()
+                .filter { it.value == languageVersionSettings.continuationAsmType() || it.value == OBJECT_TYPE }
+                .flatMap { frame.getStack(frame.stackSize - it.index - 1).insns.asSequence() }
+                .filter { it in aload0s }.let { toReplace.addAll(it) }
+        }
+
         // Expected pattern here:
         //     ALOAD 0
         //     ASTORE N
         // This pattern may occur after multiple inlines
-        val continuationToStoreAload0s = aload0s.filter { it.next?.opcode == Opcodes.ASTORE }
+        // Note, that this is not a suspension point, thus we check it separately
-        replaceContinuationsWithFakeOnes(continuationToStoreAload0s, processingNode)
+        toReplace.addAll(aload0s.filter { it.next?.opcode == Opcodes.ASTORE })
         // Expected pattern here:
         //     ALOAD 0
         //     INVOKEINTERFACE kotlin/jvm/functions/FunctionN.invoke (...,Ljava/lang/Object;)Ljava/lang/Object;
-        val continuationAsLambdaParameterAload0s = aload0s.filter { isLambdaCall(it.next) }
+        toReplace.addAll(aload0s.filter { isLambdaCall(it.next) })
-        replaceContinuationsWithFakeOnes(continuationAsLambdaParameterAload0s, processingNode)
+        replaceContinuationsWithFakeOnes(toReplace, processingNode)
     }
 
     private fun isLambdaCall(invoke: AbstractInsnNode?): Boolean {
@@ -672,7 +725,7 @@ class MethodInliner(
     }
 
     private fun replaceContinuationsWithFakeOnes(
-        continuations: Sequence<AbstractInsnNode>,
+        continuations: Collection<AbstractInsnNode>,
         node: MethodNode
     ) {
         for (toReplace in continuations) {

compiler/testData/codegen/boxInline/suspend/kt26658.kt

@@ -0,0 +1,189 @@
+// IGNORE_BACKEND: JVM_IR
+// FILE: inlined.kt
+// LANGUAGE_VERSION: 1.3
+// WITH_RUNTIME
+// NO_CHECK_LAMBDA_INLINING
+import kotlin.coroutines.*
+
+class Controller(val s: String)
+
+fun builder(c: suspend () -> Unit) {
+    c.startCoroutine(object : Continuation<Unit>{
+        override fun resumeWith(result: Result<Unit>) {
+            result.getOrThrow()
+        }
+
+        override val context: CoroutineContext
+            get() = EmptyCoroutineContext
+
+    })
+}
+
+inline fun execute(crossinline action: suspend () -> Unit) {
+    builder { action() }
+}
+
+fun builder(controller: Controller, c: suspend Controller.() -> Unit) {
+    c.startCoroutine(controller, object : Continuation<Unit>{
+        override fun resumeWith(result: Result<Unit>) {
+            result.getOrThrow()
+        }
+
+        override val context: CoroutineContext
+            get() = EmptyCoroutineContext
+    })
+}
+
+inline fun execute(controller: Controller, crossinline action: suspend Controller.() -> Unit) {
+    builder(controller) { action() }
+}
+
+// FILE: inlineSite.kt
+import kotlin.coroutines.*
+
+suspend fun withDefaultParameter(s: String, s1: String = "") = s + s1
+
+fun launch(s: String = "", block: suspend () -> String): String {
+    var res = s
+    builder { res += block() }
+    return res
+}
+
+inline fun launchCrossinline(s: String = "", crossinline block: suspend () -> String): String {
+    var res = s
+    builder { res += block() }
+    return res
+}
+
+suspend fun Controller.withDefaultParameter(s: String, s1: String = "") = this.s + s + s1
+
+fun Controller.launch(s: String = "", block: suspend Controller.() -> String): String {
+    var res = s
+    builder(this) { res += block() }
+    return res
+}
+
+inline fun Controller.launchCrossinline(s: String = "", crossinline block: suspend Controller.() -> String): String {
+    var res = s
+    builder(this) { res += block() }
+    return res
+}
+
+fun box(): String {
+    var res = ""
+    execute {
+        res = withDefaultParameter("OK")
+    }
+    if (res != "OK") return "FAIL 1: $res"
+    execute {
+        res = withDefaultParameter("O", "K")
+    }
+    if (res != "OK") return "FAIL 2: $res"
+    execute {
+        res = launch {
+            withDefaultParameter("OK")
+        }
+    }
+    if (res != "OK") return "FAIL 3: $res"
+    execute {
+        res = launch {
+            withDefaultParameter("O", "K")
+        }
+    }
+    if (res != "OK") return "FAIL 4: $res"
+    execute {
+        res = launch("O") {
+            withDefaultParameter("K")
+        }
+    }
+    if (res != "OK") return "FAIL 5: $res"
+    execute {
+        res = launch("O") {
+            withDefaultParameter("", "K")
+        }
+    }
+    if (res != "OK") return "FAIL 6: $res"
+    execute {
+        res = launchCrossinline {
+            withDefaultParameter("OK")
+        }
+    }
+    if (res != "OK") return "FAIL 7: $res"
+    execute {
+        res = launchCrossinline {
+            withDefaultParameter("O", "K")
+        }
+    }
+    if (res != "OK") return "FAIL 8: $res"
+    execute {
+        res = launchCrossinline ("O") {
+            withDefaultParameter("K")
+        }
+    }
+    if (res != "OK") return "FAIL 9: $res"
+    execute {
+        res = launchCrossinline("O") {
+            withDefaultParameter("", "K")
+        }
+    }
+    if (res != "OK") return "FAIL 10: $res"
+
+    val controller = Controller("A")
+    execute(controller) {
+        res = withDefaultParameter("OK")
+    }
+    if (res != "AOK") return "FAIL 11: $res"
+    execute(controller) {
+        res = withDefaultParameter("O", "K")
+    }
+    if (res != "AOK") return "FAIL 12: $res"
+    execute(controller) {
+        res = launch {
+            withDefaultParameter("OK")
+        }
+    }
+    if (res != "AOK") return "FAIL 13: $res"
+    execute(controller) {
+        res = launch {
+            withDefaultParameter("O", "K")
+        }
+    }
+    if (res != "AOK") return "FAIL 14: $res"
+    execute(controller) {
+        res = launch("O") {
+            withDefaultParameter("K")
+        }
+    }
+    if (res != "OAK") return "FAIL 15: $res"
+    execute(controller) {
+        res = launch("O") {
+            withDefaultParameter("", "K")
+        }
+    }
+    if (res != "OAK") return "FAIL 16: $res"
+    execute(controller) {
+        res = launchCrossinline {
+            withDefaultParameter("OK")
+        }
+    }
+    if (res != "AOK") return "FAIL 17: $res"
+    execute(controller) {
+        res = launchCrossinline {
+            withDefaultParameter("O", "K")
+        }
+    }
+    if (res != "AOK") return "FAIL 18: $res"
+    execute(controller) {
+        res = launchCrossinline ("O") {
+            withDefaultParameter("K")
+        }
+    }
+    if (res != "OAK") return "FAIL 19: $res"
+    execute(controller) {
+        res = launchCrossinline("O") {
+            withDefaultParameter("", "K")
+        }
+    }
+    if (res != "OAK") return "FAIL 20: $res"
+    return "OK"
+}

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

@@ -3354,6 +3354,11 @@ public class BlackBoxInlineCodegenTestGenerated extends AbstractBlackBoxInlineCo
             runTestWithPackageReplacement("compiler/testData/codegen/boxInline/suspend/inlineSuspendOfSuspend.kt", "kotlin.coroutines");
         }
 
+        @TestMetadata("kt26658.kt")
+        public void testKt26658() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/suspend/kt26658.kt");
+        }
+
         @TestMetadata("multipleLocals.kt")
         public void testMultipleLocals_1_2() throws Exception {
             runTestWithPackageReplacement("compiler/testData/codegen/boxInline/suspend/multipleLocals.kt", "kotlin.coroutines.experimental");

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

@@ -3354,6 +3354,11 @@ public class CompileKotlinAgainstInlineKotlinTestGenerated extends AbstractCompi
             runTestWithPackageReplacement("compiler/testData/codegen/boxInline/suspend/inlineSuspendOfSuspend.kt", "kotlin.coroutines");
         }
 
+        @TestMetadata("kt26658.kt")
+        public void testKt26658() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/suspend/kt26658.kt");
+        }
+
         @TestMetadata("multipleLocals.kt")
         public void testMultipleLocals_1_2() throws Exception {
             runTestWithPackageReplacement("compiler/testData/codegen/boxInline/suspend/multipleLocals.kt", "kotlin.coroutines.experimental");

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

@@ -3354,6 +3354,11 @@ public class IrBlackBoxInlineCodegenTestGenerated extends AbstractIrBlackBoxInli
             runTestWithPackageReplacement("compiler/testData/codegen/boxInline/suspend/inlineSuspendOfSuspend.kt", "kotlin.coroutines");
         }
 
+        @TestMetadata("kt26658.kt")
+        public void testKt26658() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/suspend/kt26658.kt");
+        }
+
         @TestMetadata("multipleLocals.kt")
         public void testMultipleLocals_1_2() throws Exception {
             runTestWithPackageReplacement("compiler/testData/codegen/boxInline/suspend/multipleLocals.kt", "kotlin.coroutines.experimental");

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/InlineSuspendTestsGenerated.java

@@ -133,6 +133,11 @@ public class InlineSuspendTestsGenerated extends AbstractInlineSuspendTests {
         runTestWithPackageReplacement("compiler/testData/codegen/boxInline/suspend/inlineSuspendOfSuspend.kt", "kotlin.coroutines");
     }
 
+    @TestMetadata("kt26658.kt")
+    public void testKt26658() throws Exception {
+        runTest("compiler/testData/codegen/boxInline/suspend/kt26658.kt");
+    }
+
     @TestMetadata("multipleLocals.kt")
     public void testMultipleLocals_1_2() throws Exception {
         runTestWithPackageReplacement("compiler/testData/codegen/boxInline/suspend/multipleLocals.kt", "kotlin.coroutines.experimental");