Support repeated 'invoke' calls on coroutines

#KT-12782 In Progress

compiler/backend/src/org/jetbrains/kotlin/codegen/ClosureCodegen.java

@@ -75,8 +75,8 @@ public class ClosureCodegen extends MemberCodegen<KtElement> {
     private final KotlinType superClassType;
     private final List<KotlinType> superInterfaceTypes;
     private final FunctionDescriptor functionReferenceTarget;
-    protected final FunctionGenerationStrategy strategy;
-    private final CalculatedClosure closure;
+    private final FunctionGenerationStrategy strategy;
+    protected final CalculatedClosure closure;
     private final Type asmType;
     private final int visibilityFlag;
 

compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/CoroutineCodegen.kt

@@ -37,9 +37,11 @@ import org.jetbrains.kotlin.resolve.jvm.diagnostics.OtherOrigin
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.util.OperatorNameConventions
+import org.jetbrains.org.objectweb.asm.Label
 import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter
+import org.jetbrains.org.objectweb.asm.commons.Method
 
 
 class CoroutineCodegen(
@@ -126,13 +128,23 @@ class CoroutineCodegen(
     }
 
     private fun generateInvokeMethod(codegen: ExpressionCodegen, signature: JvmMethodSignature) {
-        AsmUtil.genAssignInstanceFieldFromParam(
-                FieldInfo.createForHiddenField(
-                        typeMapper.mapClass(classDescriptor),
-                        typeMapper.mapType(controllerType), COROUTINE_CONTROLLER_FIELD_NAME),
-                1, codegen.v)
+        val classDescriptor = closureContext.contextDescriptor
+        val owner = typeMapper.mapClass(classDescriptor)
+        val controllerFieldInfo = FieldInfo.createForHiddenField(
+                owner,
+                typeMapper.mapType(controllerType), COROUTINE_CONTROLLER_FIELD_NAME)
+
+        val thisInstance = StackValue.thisOrOuter(codegen, classDescriptor, false, false)
 
         with(codegen.v) {
+            // if (controller != null)
+            StackValue.field(controllerFieldInfo, thisInstance).put(AsmTypes.OBJECT_TYPE, this)
+            val repeated = Label()
+            ifnonnull(repeated)
+
+            // first call
+            AsmUtil.genAssignInstanceFieldFromParam(controllerFieldInfo, 1, this)
+
             setLabelValue(LABEL_VALUE_BEFORE_FIRST_SUSPENSION)
 
             // Save lambda parameters to fields
@@ -148,6 +160,31 @@ class CoroutineCodegen(
 
             load(0, AsmTypes.OBJECT_TYPE)
             areturn(AsmTypes.OBJECT_TYPE)
+
+            // repeated call
+            visitLabel(repeated)
+            anew(owner)
+            dup()
+
+            // pass closure parameters to constructor
+            val constructorParameters = calculateConstructorParameters(typeMapper, closure, owner)
+            for (parameter in constructorParameters) {
+                StackValue.field(parameter, thisInstance).put(parameter.fieldType, this)
+            }
+
+            val constructor = Method("<init>", Type.VOID_TYPE, constructorParameters.map { it.fieldType }.toTypedArray())
+            invokespecial(owner.internalName, constructor.name, constructor.descriptor, false)
+
+            // Pass lambda parameters to 'invoke' call on newly constructed object
+            index = 1
+            for (parameter in signature.valueParameters) {
+                load(index, parameter.asmType)
+                index += parameter.asmType.size
+            }
+
+            // 'invoke' call on freshly constructed coroutine returns receiver itself
+            invokevirtual(owner.internalName, signature.asmMethod.name, signature.asmMethod.descriptor, false)
+            areturn(AsmTypes.OBJECT_TYPE)
         }
     }
 

compiler/testData/codegen/box/coroutines/multipleInvokeCalls.kt

@@ -0,0 +1,99 @@
+class Controller {
+    var lastSuspension: Continuation<String>? = null
+    var result = "fail"
+    suspend fun suspendHere(x: Continuation<String>) {
+        lastSuspension = x
+    }
+
+    fun hasNext() = lastSuspension != null
+    fun next() {
+        val x = lastSuspension!!
+        lastSuspension = null
+        x.resume("56")
+    }
+}
+
+fun builder(coroutine c: Controller.() -> Continuation<Unit>) {
+    val controller1 = Controller()
+    val controller2 = Controller()
+
+    c(controller1).resume(Unit)
+    c(controller2).resume(Unit)
+
+    runControllers(controller1, controller2)
+}
+
+fun builder2(coroutine c: Controller.(Long, String) -> Continuation<Unit>) {
+    val controller1 = Controller()
+    val controller2 = Controller()
+
+    c(controller1, 1234567890123456789L, "Q").resume(Unit)
+    c(controller2, 1234567890123456789L, "Q").resume(Unit)
+
+    runControllers(controller1, controller2)
+}
+
+
+private fun runControllers(controller1: Controller, controller2: Controller) {
+    while (controller1.hasNext()) {
+        if (!controller2.hasNext()) throw RuntimeException("fail 1")
+
+        if (controller1.lastSuspension === controller2.lastSuspension) throw RuntimeException("equal references")
+
+        controller1.next()
+        controller2.next()
+    }
+
+    if (controller2.hasNext()) throw RuntimeException("fail 2")
+
+    if (controller1.result != "OK") throw RuntimeException("fail 3")
+    if (controller2.result != "OK") throw RuntimeException("fail 4")
+}
+
+fun box(): String {
+    // no suspension
+    builder {
+        result = "OK"
+    }
+
+    // 1 suspension
+    builder {
+        if (suspendHere() != "56") return@builder
+        result = "OK"
+    }
+
+    // 2 suspensions
+    builder {
+        if (suspendHere() != "56") return@builder
+        suspendHere()
+        result = "OK"
+    }
+
+    // with capture and params
+
+    var x = "O"
+    var y = "K"
+
+    // no suspension
+    builder2 { a, b ->
+        if (a != 1234567890123456789L || b != "Q" ) return@builder2
+        result = x + y
+    }
+
+    // 1 suspension
+    builder2 { a, b ->
+        if (a != 1234567890123456789L || b != "Q" ) return@builder2
+        if (suspendHere() != "56") return@builder2
+        result = x + y
+    }
+
+    // 2 suspensions
+    builder2 { a, b ->
+        if (a != 1234567890123456789L || b != "Q" ) return@builder2
+        if (suspendHere() != "56") return@builder2
+        suspendHere()
+        result = x + y
+    }
+
+    return "OK"
+}

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

@@ -4225,6 +4225,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("multipleInvokeCalls.kt")
+        public void testMultipleInvokeCalls() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/multipleInvokeCalls.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("nestedTryCatch.kt")
         public void testNestedTryCatch() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/nestedTryCatch.kt");