From fcd7677a3f8dd6e86a7299ab2484112548137b8b Mon Sep 17 00:00:00 2001
From: Denis Zharkov <denis.zharkov@jetbrains.com>
Date: Thu, 17 Aug 2017 16:36:02 +0700
Subject: [PATCH] Fix compatibility of suspend functions with strict bytecode
 analyzers

In short, some of the bytecode analyzers assume that there could be
no stores instructions into parameter vars with value of different
types (even when the value type is a subtype)

See the issue for details

 #KT-19713 Fixed
---
 .../CoroutineTransformerMethodVisitor.kt       | 15 ++++++++++++---
 .../coroutines/doNotReassignContinuation.kt    | 18 ++++++++++++++++++
 .../codegen/BytecodeTextTestGenerated.java     |  6 ++++++
 3 files changed, 36 insertions(+), 3 deletions(-)
 create mode 100644 compiler/testData/codegen/bytecodeText/coroutines/doNotReassignContinuation.kt

diff --git a/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/CoroutineTransformerMethodVisitor.kt b/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/CoroutineTransformerMethodVisitor.kt
index 77e460d6511..6d9f6fab627 100644
--- a/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/CoroutineTransformerMethodVisitor.kt
+++ b/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/CoroutineTransformerMethodVisitor.kt
@@ -63,7 +63,7 @@ class CoroutineTransformerMethodVisitor(
 
     private val classBuilderForCoroutineState: ClassBuilder by lazy(obtainClassBuilderForCoroutineState)
 
-    private val continuationIndex = if (isForNamedFunction) getLastParameterIndex(desc, access) else 0
+    private var continuationIndex = if (isForNamedFunction) -1 else 0
     private var dataIndex = if (isForNamedFunction) -1 else 1
     private var exceptionIndex = if (isForNamedFunction) -1 else 2
 
@@ -83,6 +83,7 @@ class CoroutineTransformerMethodVisitor(
 
             dataIndex = methodNode.maxLocals++
             exceptionIndex = methodNode.maxLocals++
+            continuationIndex = methodNode.maxLocals++
 
             prepareMethodNodePreludeForNamedFunction(methodNode)
         }
@@ -194,6 +195,13 @@ class CoroutineTransformerMethodVisitor(
 
     private fun prepareMethodNodePreludeForNamedFunction(methodNode: MethodNode) {
         val objectTypeForState = Type.getObjectType(classBuilderForCoroutineState.thisName)
+        val continuationArgumentIndex = getLastParameterIndex(methodNode.desc, methodNode.access)
+        methodNode.instructions.asSequence().filterIsInstance<VarInsnNode>().forEach {
+            if (it.`var` != continuationArgumentIndex) return@forEach
+            assert(it.opcode == Opcodes.ALOAD) { "Only ALOADs are allowed for continuation arguments" }
+            it.`var` = continuationIndex
+        }
+
         methodNode.instructions.insert(withInstructionAdapter {
             val createStateInstance = Label()
             val afterCoroutineStateCreated = Label()
@@ -212,11 +220,12 @@ class CoroutineTransformerMethodVisitor(
             // - Otherwise it's still can be a recursive call. To check it's not the case we set the last bit in the label in
             // `doResume` just before calling the suspend function (see kotlin.coroutines.experimental.jvm.internal.CoroutineImplForNamedFunction).
             // So, if it's set we're in continuation.
-            visitVarInsn(Opcodes.ALOAD, continuationIndex)
+
+            visitVarInsn(Opcodes.ALOAD, continuationArgumentIndex)
             instanceOf(objectTypeForState)
             ifeq(createStateInstance)
 
-            visitVarInsn(Opcodes.ALOAD, continuationIndex)
+            visitVarInsn(Opcodes.ALOAD, continuationArgumentIndex)
             checkcast(objectTypeForState)
             visitVarInsn(Opcodes.ASTORE, continuationIndex)
 
diff --git a/compiler/testData/codegen/bytecodeText/coroutines/doNotReassignContinuation.kt b/compiler/testData/codegen/bytecodeText/coroutines/doNotReassignContinuation.kt
new file mode 100644
index 00000000000..6185ede2818
--- /dev/null
+++ b/compiler/testData/codegen/bytecodeText/coroutines/doNotReassignContinuation.kt
@@ -0,0 +1,18 @@
+// WITH_RUNTIME
+// WITH_COROUTINES
+import helpers.*
+// TREAT_AS_ONE_FILE
+import kotlin.coroutines.experimental.*
+import kotlin.coroutines.experimental.intrinsics.*
+suspend fun suspendHere(): String = suspendCoroutineOrReturn { x ->
+    x.resume("OK")
+}
+
+suspend fun suspendThere(param: Int, param2: String, param3: Long): String {
+    val a = suspendHere()
+    val b = suspendHere()
+    return a + b
+}
+
+/* 2 stores happen because the continuation parameter is visible in debug and should be spilled */
+// 2 ASTORE 4
diff --git a/compiler/tests/org/jetbrains/kotlin/codegen/BytecodeTextTestGenerated.java b/compiler/tests/org/jetbrains/kotlin/codegen/BytecodeTextTestGenerated.java
index 1e68a49a0a2..54bf3934f7f 100644
--- a/compiler/tests/org/jetbrains/kotlin/codegen/BytecodeTextTestGenerated.java
+++ b/compiler/tests/org/jetbrains/kotlin/codegen/BytecodeTextTestGenerated.java
@@ -1037,6 +1037,12 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/bytecodeText/coroutines"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.ANY, true);
         }
 
+        @TestMetadata("doNotReassignContinuation.kt")
+        public void testDoNotReassignContinuation() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/coroutines/doNotReassignContinuation.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("varValueConflictsWithTable.kt")
         public void testVarValueConflictsWithTable() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/coroutines/varValueConflictsWithTable.kt");