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) {