Support repeated 'invoke' calls on coroutines defined within inline lambdas

#KT-12782 Fixed

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/AnonymousObjectTransformer.java

@@ -67,6 +67,14 @@ public class AnonymousObjectTransformer extends ObjectTransformer<AnonymousObjec
             public void visit(int version, int access, @NotNull String name, String signature, String superName, String[] interfaces) {
                 InlineCodegenUtil.assertVersionNotGreaterThanGeneratedOne(version, name, inliningContext.state);
                 classBuilder.defineClass(null, version, access, name, signature, superName, interfaces);
+                if (interfaces != null) {
+                    for (String anInterface : interfaces) {
+                        if("kotlin/coroutines/Continuation".equals(anInterface)) {
+                            inliningContext.setContinuation(true);
+                            break;
+                        }
+                    }
+                }
             }
 
             @Override
@@ -139,6 +147,8 @@ public class AnonymousObjectTransformer extends ObjectTransformer<AnonymousObjec
                                                             transformationInfo, parentRemapper);
         List<MethodVisitor> deferringMethods = new ArrayList<MethodVisitor>();
 
+        generateConstructorAndFields(classBuilder, allCapturedParamBuilder, constructorParamBuilder, parentRemapper, additionalFakeParams);
+
         for (MethodNode next : methodsToTransform) {
             MethodVisitor deferringVisitor = newMethod(classBuilder, next);
             InlineResult funResult =
@@ -159,8 +169,6 @@ public class AnonymousObjectTransformer extends ObjectTransformer<AnonymousObjec
             method.visitEnd();
         }
 
-        generateConstructorAndFields(classBuilder, allCapturedParamBuilder, constructorParamBuilder, parentRemapper, additionalFakeParams);
-
         SourceMapper.Companion.flushToClassBuilder(sourceMapper, classBuilder);
 
         ClassVisitor visitor = classBuilder.getVisitor();
@@ -362,7 +370,7 @@ public class AnonymousObjectTransformer extends ObjectTransformer<AnonymousObjec
             @NotNull MethodNode constructor,
             @NotNull ParametersBuilder capturedParamBuilder,
             @NotNull ParametersBuilder constructorParamBuilder,
-            @NotNull final AnonymousObjectTransformationInfo transformationInfo,
+            @NotNull AnonymousObjectTransformationInfo transformationInfo,
             @NotNull FieldRemapper parentFieldRemapper
     ) {
         Set<LambdaInfo> capturedLambdas = new LinkedHashSet<LambdaInfo>(); //captured var of inlined parameter

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/InliningContext.java

@@ -35,6 +35,10 @@ public class InliningContext {
     public final ReifiedTypeInliner reifiedTypeInliner;
     public final boolean isInliningLambda;
     public final boolean classRegeneration;
+    public final Map<String, AnonymousObjectTransformationInfo> internalNameToAnonymousObjectTransformationInfo =
+            new HashMap<String, AnonymousObjectTransformationInfo>();
+
+    private boolean isContinuation;
 
     public InliningContext(
             @Nullable InliningContext parent,
@@ -117,4 +121,21 @@ public class InliningContext {
         assert parent != null : "At least root context should return proper value";
         return parent.getCallSiteInfo();
     }
+
+    @Nullable
+    public AnonymousObjectTransformationInfo findAnonymousObjectTransformationInfo(@NotNull String internalName) {
+        if (getRoot().internalNameToAnonymousObjectTransformationInfo.containsKey(internalName)) {
+            return getRoot().internalNameToAnonymousObjectTransformationInfo.get(internalName);
+        }
+
+        return null;
+    }
+
+    public boolean isContinuation() {
+        return isContinuation;
+    }
+
+    public void setContinuation(boolean continuation) {
+        isContinuation = continuation;
+    }
 }

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

@@ -260,16 +260,27 @@ public class MethodInliner {
                     assert transformationInfo instanceof AnonymousObjectTransformationInfo :
                             "<init> call doesn't correspond to object transformation info: " +
                             owner + "." + name + ", info " + transformationInfo;
-                    if (transformationInfo.shouldRegenerate(isSameModule)) {
+                    InliningContext parent = inliningContext.getParent();
+                    boolean shouldRegenerate = transformationInfo.shouldRegenerate(isSameModule);
+                    if (shouldRegenerate || (parent != null && parent.isContinuation())) {
+                        assert shouldRegenerate || inlineCallSiteInfo.getOwnerClassName().equals(transformationInfo.getOldClassName())
+                                : "Only coroutines can call their own constructors";
+
                         //put additional captured parameters on stack
                         AnonymousObjectTransformationInfo info = (AnonymousObjectTransformationInfo) transformationInfo;
+
+                        AnonymousObjectTransformationInfo oldInfo = inliningContext.findAnonymousObjectTransformationInfo(owner);
+                        if (oldInfo != null) {
+                            info = oldInfo;
+                        }
+
                         for (CapturedParamDesc capturedParamDesc : info.getAllRecapturedParameters()) {
                             visitFieldInsn(
                                     Opcodes.GETSTATIC, capturedParamDesc.getContainingLambdaName(),
                                     "$$$" + capturedParamDesc.getFieldName(), capturedParamDesc.getType().getDescriptor()
                             );
                         }
-                        super.visitMethodInsn(opcode, transformationInfo.getNewClassName(), name, info.getNewConstructorDescriptor(), itf);
+                        super.visitMethodInsn(opcode, info.getNewClassName(), name, info.getNewConstructorDescriptor(), itf);
 
                         //TODO: add new inner class also for other contexts
                         if (inliningContext.getParent() instanceof RegeneratedClassContext) {
@@ -565,7 +576,9 @@ public class MethodInliner {
             @NotNull Map<Integer, LambdaInfo> lambdaMapping,
             boolean needReification
     ) {
-        return new AnonymousObjectTransformationInfo(
+        boolean memoizeAnonymousObject = inliningContext.findAnonymousObjectTransformationInfo(anonymousType) == null;
+
+        AnonymousObjectTransformationInfo info = new AnonymousObjectTransformationInfo(
                 anonymousType, needReification, lambdaMapping,
                 inliningContext.classRegeneration,
                 isAlreadyRegenerated(anonymousType),
@@ -573,6 +586,11 @@ public class MethodInliner {
                 false,
                 inliningContext.nameGenerator
         );
+
+        if (memoizeAnonymousObject) {
+            inliningContext.getRoot().internalNameToAnonymousObjectTransformationInfo.put(anonymousType, info);
+        }
+        return info;
     }
 
     private boolean isAlreadyRegenerated(@NotNull String owner) {

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

@@ -0,0 +1,88 @@
+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")
+}
+
+inline fun run(b: () -> Unit) {
+    b()
+}
+
+fun box(): String {
+    // with capture and params
+
+    var x = "O"
+    var y = "K"
+
+    // inlined
+    run {
+        // 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/testData/codegen/box/coroutines/multipleInvokeCallsInsideInlineLambda2.kt

@@ -0,0 +1,109 @@
+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")
+}
+
+inline fun run(b: () -> Unit) {
+    b()
+}
+
+fun box(): String {
+    // with capture and params
+    var x = "O"
+
+    // inlined
+    run {
+        var y = "K"
+
+        {
+            // 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"
+            }
+
+            // 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/testData/codegen/box/coroutines/multipleInvokeCallsInsideInlineLambda3.kt

@@ -0,0 +1,107 @@
+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")
+}
+
+inline fun run(b: () -> Unit) {
+    b()
+}
+
+fun box(): String {
+    var x = "O"
+
+    {
+        var y = "K"
+        // inlined
+        run {
+            // 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"
+            }
+
+
+            // 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

@@ -4231,6 +4231,24 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("multipleInvokeCallsInsideInlineLambda1.kt")
+        public void testMultipleInvokeCallsInsideInlineLambda1() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/multipleInvokeCallsInsideInlineLambda1.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("multipleInvokeCallsInsideInlineLambda2.kt")
+        public void testMultipleInvokeCallsInsideInlineLambda2() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/multipleInvokeCallsInsideInlineLambda2.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("multipleInvokeCallsInsideInlineLambda3.kt")
+        public void testMultipleInvokeCallsInsideInlineLambda3() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/multipleInvokeCallsInsideInlineLambda3.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("nestedTryCatch.kt")
         public void testNestedTryCatch() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/nestedTryCatch.kt");