JVM_IR: Lazyly generate continuation classes of suspend functions

Otherwise, they will be generated, but unused, if the function is tail-call.
 #KT-36795 Fixed

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmBackendContext.kt

@@ -20,7 +20,6 @@ import org.jetbrains.kotlin.backend.jvm.intrinsics.IrIntrinsicMethods
 import org.jetbrains.kotlin.backend.jvm.lower.inlineclasses.InlineClassAbi
 import org.jetbrains.kotlin.backend.jvm.lower.inlineclasses.MemoizedInlineClassReplacements
 import org.jetbrains.kotlin.backend.jvm.lower.suspendFunctionOriginal
-import org.jetbrains.kotlin.codegen.ClassBuilder
 import org.jetbrains.kotlin.codegen.inline.NameGenerator
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.descriptors.ClassConstructorDescriptor
@@ -114,7 +113,6 @@ class JvmBackendContext(
     override val internalPackageFqn = FqName("kotlin.jvm")
 
     val suspendLambdaToOriginalFunctionMap = mutableMapOf<IrFunctionReference, IrFunction>()
-    val continuationClassBuilders = mutableMapOf<IrSimpleFunction, ClassBuilder>()
     val suspendFunctionOriginalToView = mutableMapOf<IrFunction, IrFunction>()
     val suspendFunctionOriginalToStub = mutableMapOf<IrFunction, IrFunction>()
     val fakeContinuation: IrExpression = createFakeContinuation(this)

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ClassCodegen.kt

@@ -181,9 +181,7 @@ open class ClassCodegen protected constructor(
 
         generateKotlinMetadataAnnotation()
 
-        if (irClass.origin == JvmLoweredDeclarationOrigin.CONTINUATION_CLASS) {
+        if (irClass.origin != JvmLoweredDeclarationOrigin.CONTINUATION_CLASS) {
-            context.continuationClassBuilders[irClass.attributeOwnerId as IrSimpleFunction] = visitor
-        } else {
             done()
         }
         return reifiedTypeParametersUsages
@@ -341,9 +339,12 @@ open class ClassCodegen protected constructor(
     }
 
     fun generateLocalClass(klass: IrClass, parentFunction: IrFunction): ReifiedTypeParametersUsages {
-        return ClassCodegen(klass, context, this, parentFunction, withinInline = withinInline || parentFunction.isInline).generate()
+        return createLocalClassCodegen(klass, parentFunction).generate()
     }
 
+    fun createLocalClassCodegen(klass: IrClass, parentFunction: IrFunction): ClassCodegen =
+            ClassCodegen(klass, context, this, parentFunction, withinInline = withinInline || parentFunction.isInline)
+
     private fun generateField(field: IrField) {
         if (field.origin == IrDeclarationOrigin.FAKE_OVERRIDE) return
 

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/CoroutineCodegen.kt

@@ -20,10 +20,7 @@ import org.jetbrains.kotlin.codegen.inline.addFakeContinuationConstructorCallMar
 import org.jetbrains.kotlin.config.isReleaseCoroutines
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.*
-import org.jetbrains.kotlin.ir.expressions.IrGetField
-import org.jetbrains.kotlin.ir.expressions.IrGetValue
-import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
 import org.jetbrains.kotlin.ir.expressions.impl.IrErrorExpressionImpl
 import org.jetbrains.kotlin.ir.types.createType
 import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
@@ -36,6 +33,7 @@ import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodGenericSignature
 import org.jetbrains.kotlin.types.Variance
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
 import org.jetbrains.org.objectweb.asm.MethodVisitor
 import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter
@@ -161,13 +159,13 @@ internal fun createFakeContinuation(context: JvmBackendContext): IrExpression =
 
 internal fun generateFakeContinuationConstructorCall(
     v: InstructionAdapter,
-    containingClassBuilder: ClassBuilder,
+    classCodegen: ClassCodegen,
-    classBuilder: ClassBuilder,
+    continuationClass: IrClass,
     irFunction: IrFunction
 ) {
-    val continuationType = Type.getObjectType(classBuilder.thisName)
+    val continuationType = classCodegen.typeMapper.mapClass(continuationClass)
     // TODO: This is different in case of DefaultImpls
-    val thisNameType = Type.getObjectType(containingClassBuilder.thisName.replace(".", "/"))
+    val thisNameType = Type.getObjectType(classCodegen.visitor.thisName.replace(".", "/"))
     val continuationIndex = listOfNotNull(irFunction.dispatchReceiverParameter, irFunction.extensionReceiverParameter).size +
             irFunction.valueParameters.size - 1
     with(v) {
@@ -185,4 +183,6 @@ internal fun generateFakeContinuationConstructorCall(
         addFakeContinuationConstructorCallMarker(this, false)
         pop()
     }
-}
+}
+
+internal fun IrFunction.continuationClass(): IrClass = (body as IrBlockBody).statements.firstIsInstance<IrClass>()

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ExpressionCodegen.kt

@@ -223,12 +223,7 @@ class ExpressionCodegen(
             it.attributeOwnerId == (irFunction as? IrSimpleFunction)?.attributeOwnerId &&
                     it.name.asString() == irFunction.name.asString().removeSuffix(FOR_INLINE_SUFFIX)
         }?.body?.statements?.firstIsInstance<IrClass>() ?: error("could not find continuation for ${irFunction.render()}")
-        generateFakeContinuationConstructorCall(
+        generateFakeContinuationConstructorCall(mv, classCodegen, continuationClass, irFunction)
-            mv,
-            classCodegen.visitor,
-            context.continuationClassBuilders[(continuationClass as IrClass).attributeOwnerId]!!,
-            irFunction
-        )
     }
 
     private fun generateNonNullAssertions() {
@@ -618,8 +613,10 @@ class ExpressionCodegen(
         )
 
     override fun visitClass(declaration: IrClass, data: BlockInfo): PromisedValue {
-        classCodegen.generateLocalClass(declaration, generateSequence(this) { it.inlinedInto }.last().irFunction).also {
+        if (declaration.origin != JvmLoweredDeclarationOrigin.CONTINUATION_CLASS) {
-            closureReifiedMarkers[declaration] = it
+            classCodegen.generateLocalClass(declaration, generateSequence(this) { it.inlinedInto }.last().irFunction).also {
+                closureReifiedMarkers[declaration] = it
+            }
         }
         return immaterialUnitValue
     }

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/FunctionCodegen.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.backend.common.lower.BOUND_VALUE_PARAMETER
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.codegen.AsmUtil
+import org.jetbrains.kotlin.codegen.inline.coroutines.FOR_INLINE_SUFFIX
 import org.jetbrains.kotlin.codegen.mangleNameIfNeeded
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.codegen.visitAnnotableParameterCount
@@ -44,6 +45,10 @@ open class FunctionCodegen(
     val context = classCodegen.context
     val state = classCodegen.state
 
+    val continuationClassBuilder by lazy {
+        classCodegen.createLocalClassCodegen(irFunction.continuationClass(), irFunction).also { it.generate() }.visitor
+    }
+
     fun generate(): JvmMethodGenericSignature =
         try {
             doGenerate()
@@ -90,13 +95,20 @@ open class FunctionCodegen(
             val frameMap = createFrameMapWithReceivers()
             methodVisitor = when {
                 irFunction.hasContinuation() -> {
+                    if (irFunction.parentAsClass.declarations.any {
+                        (it as? IrAttributeContainer)?.attributeOwnerId == (irFunction as IrAttributeContainer).attributeOwnerId &&
+                                it.nameForIrSerialization.asString() == irFunction.nameForIrSerialization.asString() + FOR_INLINE_SUFFIX &&
+                                it.origin == JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
+                        }
+                    ) {
+                        // force generation of fake continuation for inliner.
+                        continuationClassBuilder
+                    }
                     generateStateMachineForNamedFunction(
                         irFunction, classCodegen, methodVisitor,
                         access = flags,
                         signature = signature,
-                        obtainContinuationClassBuilder = {
+                        obtainContinuationClassBuilder = { continuationClassBuilder },
-                            context.continuationClassBuilders[continuationClass().attributeOwnerId]!!
-                        },
                         element = psiElement()
                     )
                 }
@@ -107,11 +119,6 @@ open class FunctionCodegen(
             }
             ExpressionCodegen(irFunction, signature, frameMap, InstructionAdapter(methodVisitor), classCodegen, inlinedInto).generate()
             methodVisitor.visitMaxs(-1, -1)
-            if (irFunction.hasContinuation()) {
-                context.continuationClassBuilders[continuationClass().attributeOwnerId].sure {
-                    "Could not find continuation class builder for ${continuationClass().render()}"
-                }.done()
-            }
         }
         methodVisitor.visitEnd()
 
@@ -143,9 +150,6 @@ open class FunctionCodegen(
             origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE &&
             origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
 
-    private fun continuationClass() =
-        irFunction.body!!.statements.firstIsInstance<IrClass>()
-
     private fun IrFunction.getVisibilityForDefaultArgumentStub(): Int =
         when (visibility) {
             Visibilities.PUBLIC -> Opcodes.ACC_PUBLIC

compiler/testData/codegen/box/coroutines/bridges/lambdaWithLongReceiver_ir.txt

@@ -25,14 +25,6 @@ final class LambdaWithLongReceiverKt$box$1 {
     public final @org.jetbrains.annotations.Nullable method invokeSuspend(@org.jetbrains.annotations.NotNull p0: java.lang.Object): java.lang.Object
 }
 
-@kotlin.Metadata
-final class LambdaWithLongReceiverKt$call$1 {
-    field label: int
-    @org.jetbrains.annotations.NotNull field result: java.lang.Object
-    public method <init>(p0: kotlin.coroutines.Continuation): void
-    public final @org.jetbrains.annotations.Nullable method invokeSuspend(@org.jetbrains.annotations.NotNull p0: java.lang.Object): java.lang.Object
-}
-
 @kotlin.Metadata
 public final class LambdaWithLongReceiverKt {
     inner class LambdaWithLongReceiverKt$box$1

compiler/testData/codegen/box/coroutines/bridges/lambdaWithMultipleParameters_ir.txt

@@ -29,14 +29,6 @@ final class LambdaWithMultipleParametersKt$box$1 {
     public final @org.jetbrains.annotations.Nullable method invokeSuspend(@org.jetbrains.annotations.NotNull p0: java.lang.Object): java.lang.Object
 }
 
-@kotlin.Metadata
-final class LambdaWithMultipleParametersKt$call$1 {
-    field label: int
-    @org.jetbrains.annotations.NotNull field result: java.lang.Object
-    public method <init>(p0: kotlin.coroutines.Continuation): void
-    public final @org.jetbrains.annotations.Nullable method invokeSuspend(@org.jetbrains.annotations.NotNull p0: java.lang.Object): java.lang.Object
-}
-
 @kotlin.Metadata
 public final class LambdaWithMultipleParametersKt {
     inner class LambdaWithMultipleParametersKt$box$1

compiler/testData/codegen/box/coroutines/tailCallOptimizations/crossinline_ir.txt

@@ -70,14 +70,6 @@ final class CrossinlineKt$box$1 {
     public final @org.jetbrains.annotations.Nullable method invokeSuspend(@org.jetbrains.annotations.NotNull p0: java.lang.Object): java.lang.Object
 }
 
-@kotlin.Metadata
-final class CrossinlineKt$consumeEach$1 {
-    field label: int
-    @org.jetbrains.annotations.NotNull field result: java.lang.Object
-    public method <init>(p0: kotlin.coroutines.Continuation): void
-    public final @org.jetbrains.annotations.Nullable method invokeSuspend(@org.jetbrains.annotations.NotNull p0: java.lang.Object): java.lang.Object
-}
-
 @kotlin.Metadata
 final class CrossinlineKt$consumeEach$2$send$1 {
     field label: int

compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithoutStateMachine_ir.txt

@@ -23,22 +23,6 @@ final class InlineWithoutStateMachineKt$complexSuspend$1 {
     public final @org.jetbrains.annotations.Nullable method invokeSuspend(@org.jetbrains.annotations.NotNull p0: java.lang.Object): java.lang.Object
 }
 
-@kotlin.Metadata
-final class InlineWithoutStateMachineKt$suspendHere$1 {
-    field label: int
-    @org.jetbrains.annotations.NotNull field result: java.lang.Object
-    public method <init>(p0: kotlin.coroutines.Continuation): void
-    public final @org.jetbrains.annotations.Nullable method invokeSuspend(@org.jetbrains.annotations.NotNull p0: java.lang.Object): java.lang.Object
-}
-
-@kotlin.Metadata
-final class InlineWithoutStateMachineKt$suspendThere$1 {
-    field label: int
-    @org.jetbrains.annotations.NotNull field result: java.lang.Object
-    public method <init>(p0: kotlin.coroutines.Continuation): void
-    public final @org.jetbrains.annotations.Nullable method invokeSuspend(@org.jetbrains.annotations.NotNull p0: java.lang.Object): java.lang.Object
-}
-
 @kotlin.Metadata
 public final class InlineWithoutStateMachineKt {
     inner class InlineWithoutStateMachineKt$box$1

compiler/testData/codegen/box/coroutines/tailCallOptimizations/innerObjectRetransformation_ir.txt

@@ -51,14 +51,6 @@ final class flow/InnerObjectRetransformationKt$check$1 {
     public final @org.jetbrains.annotations.Nullable method invokeSuspend(@org.jetbrains.annotations.NotNull p0: java.lang.Object): java.lang.Object
 }
 
-@kotlin.Metadata
-final class flow/InnerObjectRetransformationKt$collect$1 {
-    field label: int
-    @org.jetbrains.annotations.NotNull field result: java.lang.Object
-    public method <init>(p0: kotlin.coroutines.Continuation): void
-    public final @org.jetbrains.annotations.Nullable method invokeSuspend(@org.jetbrains.annotations.NotNull p0: java.lang.Object): java.lang.Object
-}
-
 @kotlin.Metadata
 final class flow/InnerObjectRetransformationKt$collect$2$emit$1 {
     field label: int

compiler/testData/codegen/bytecodeText/coroutines/debug/thisAndResultInLvt.kt

@@ -30,9 +30,8 @@ class A {
 // JVM_IR_TEMPLATES
 // <init>, invoke, invokeSuspend, create
 // 4 LOCALVARIABLE this LThisAndResultInLvtKt\$c\$1; L0 L.* 0
-// TODO: Continuation are generated even for tail-call suspend functions: KT-36795
+// c's lambda and foo's continuation
-// c's lambda and foo's and dummy's continuation
+// 2 LOCALVARIABLE \$result Ljava/lang/Object; L0 L.* 1
-// 3 LOCALVARIABLE \$result Ljava/lang/Object; L0 L.* 1
 
 // foo and <init>
 // 2 LOCALVARIABLE this LA; L0 L.* 0