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
2020-02-26 20:52:41 +01:00
parent 91ee63432a
commit 536e0e23a0
11 changed files with 36 additions and 85 deletions
@@ -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)
@@ -181,9 +181,7 @@ open class ClassCodegen protected constructor(

        generateKotlinMetadataAnnotation()

-        if (irClass.origin == JvmLoweredDeclarationOrigin.CONTINUATION_CLASS) {
-            context.continuationClassBuilders[irClass.attributeOwnerId as IrSimpleFunction] = visitor
-        } else {
+        if (irClass.origin != JvmLoweredDeclarationOrigin.CONTINUATION_CLASS) {
            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

@@ -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.IrGetField
-import org.jetbrains.kotlin.ir.expressions.IrGetValue
-import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
+import org.jetbrains.kotlin.ir.expressions.*
 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,
-    classBuilder: ClassBuilder,
+    classCodegen: ClassCodegen,
+    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>()
@@ -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(
-            mv,
-            classCodegen.visitor,
-            context.continuationClassBuilders[(continuationClass as IrClass).attributeOwnerId]!!,
-            irFunction
-        )
+        generateFakeContinuationConstructorCall(mv, classCodegen, continuationClass, 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 {
-            closureReifiedMarkers[declaration] = it
+        if (declaration.origin != JvmLoweredDeclarationOrigin.CONTINUATION_CLASS) {
+            classCodegen.generateLocalClass(declaration, generateSequence(this) { it.inlinedInto }.last().irFunction).also {
+                closureReifiedMarkers[declaration] = it
+            }
        }
        return immaterialUnitValue
    }
@@ -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 = {
-                            context.continuationClassBuilders[continuationClass().attributeOwnerId]!!
-                        },
+                        obtainContinuationClassBuilder = { continuationClassBuilder },
                        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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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 and dummy's continuation
-// 3 LOCALVARIABLE \$result Ljava/lang/Object; L0 L.* 1
+// c's lambda and foo's continuation
+// 2 LOCALVARIABLE \$result Ljava/lang/Object; L0 L.* 1

 // foo and <init>
 // 2 LOCALVARIABLE this LA; L0 L.* 0