JVM_IR: Move suspend function views creation to lowering

Now AddContinuationLowering is responsible for both adding continuation classes to suspend functions and adding continuation parameters to them. Because we cannot create a view if inline suspend function is defined in another file, we generate a stub without body when we encounter call to it. And then, when we lower the file containing the function we add the body. This way we have no unlowered views after the lowering. Thus, after the lowering there should be no suspend function, which are not views, therefore, remove VIEW origins. Because transformations of suspend functions can copy them into another object, use attribute as a key inside function to view map.
2020-01-31 12:35:11 +01:00
parent 5e7343624b
commit fc70455877
18 changed files with 291 additions and 238 deletions
@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */

@@ -33,7 +33,6 @@ interface JvmLoweredDeclarationOrigin : IrDeclarationOrigin {
    object INLINE_CLASS_GENERATED_IMPL_METHOD : IrDeclarationOriginImpl("INLINE_CLASS_GENERATED_IMPL_METHOD")
    object GENERATED_ASSERTION_ENABLED_FIELD : IrDeclarationOriginImpl("GENERATED_ASSERTION_ENABLED_FIELD", isSynthetic = true)
    object GENERATED_MAIN_FOR_PARAMETERLESS_MAIN_METHOD : IrDeclarationOriginImpl("GENERATED_MAIN_FOR_PARAMETERLESS_MAIN_METHOD", isSynthetic = true)
-    object SUSPEND_FUNCTION_VIEW : IrDeclarationOriginImpl("SUSPEND_FUNCTION_VIEW")
    object SUSPEND_IMPL_STATIC_FUNCTION : IrDeclarationOriginImpl("SUSPEND_IMPL_STATIC_FUNCTION", isSynthetic = true)
    object INTERFACE_COMPANION_PRIVATE_INSTANCE : IrDeclarationOriginImpl("INTERFACE_COMPANION_PRIVATE_INSTANCE", isSynthetic = true)
    object POLYMORPHIC_SIGNATURE_INSTANTIATION : IrDeclarationOriginImpl("POLYMORPHIC_SIGNATURE_INSTANTIATION", isSynthetic = true)
@@ -43,6 +42,5 @@ interface JvmLoweredDeclarationOrigin : IrDeclarationOrigin {
    object SUSPEND_LAMBDA : IrDeclarationOriginImpl("SUSPEND_LAMBDA")
    object FOR_INLINE_STATE_MACHINE_TEMPLATE : IrDeclarationOriginImpl("FOR_INLINE_TEMPLATE")
    object FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE : IrDeclarationOriginImpl("FOR_INLINE_TEMPLATE_CROSSINLINE")
-    object FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE_VIEW : IrDeclarationOriginImpl("FOR_INLINE_TEMPLATE_CROSSINLINE_VIEW")
    object COMPANION_PROPERTY_BACKING_FIELD : IrDeclarationOriginImpl("COMPANION_MOVED_PROPERTY_BACKING_FIELD")
 }
@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.backend.jvm.descriptors.JvmSharedVariablesManager
 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
@@ -108,7 +109,7 @@ class JvmBackendContext(
    val suspendLambdaToOriginalFunctionMap = mutableMapOf<IrFunctionReference, IrFunction>()
    val continuationClassBuilders = mutableMapOf<IrSimpleFunction, ClassBuilder>()
    val suspendFunctionOriginalToView = mutableMapOf<IrFunction, IrFunction>()
-    val suspendFunctionViewToOriginal = mutableMapOf<IrFunction, IrFunction>()
+    val suspendFunctionOriginalToStub = mutableMapOf<IrFunction, IrFunction>()
    val suspendTailCallsWithUnitReplacement = mutableSetOf<IrAttributeContainer>()
    val fakeContinuation: IrExpression = createFakeContinuation(this)

@@ -117,8 +118,13 @@ class JvmBackendContext(
    val inlineClassReplacements = MemoizedInlineClassReplacements()

    internal fun recordSuspendFunctionView(function: IrFunction, view: IrFunction) {
-        suspendFunctionOriginalToView[function] = view
-        suspendFunctionViewToOriginal[view] = function
+        val attribute = function.suspendFunctionOriginal()
+        suspendFunctionOriginalToStub.remove(attribute)
+        suspendFunctionOriginalToView[attribute] = view
+    }
+
+    internal fun recordSuspendFunctionViewStub(function: IrFunction, stub: IrFunction) {
+        suspendFunctionOriginalToStub[function.suspendFunctionOriginal()] = stub
    }

    internal fun referenceClass(descriptor: ClassDescriptor): IrClassSymbol =
@@ -6,36 +6,29 @@
 package org.jetbrains.kotlin.backend.jvm.codegen

 import org.jetbrains.kotlin.backend.common.CodegenUtil
-import org.jetbrains.kotlin.backend.common.ir.copyTypeParametersFrom
-import org.jetbrains.kotlin.backend.common.ir.copyValueParametersInsertingContinuationFrom
-import org.jetbrains.kotlin.backend.common.ir.isSuspend
-import org.jetbrains.kotlin.backend.common.lower.VariableRemapper
 import org.jetbrains.kotlin.backend.common.lower.allOverridden
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
+import org.jetbrains.kotlin.backend.jvm.lower.suspendFunctionOriginal
 import org.jetbrains.kotlin.codegen.ClassBuilder
 import org.jetbrains.kotlin.codegen.coroutines.CoroutineTransformerMethodVisitor
 import org.jetbrains.kotlin.codegen.coroutines.INVOKE_SUSPEND_METHOD_NAME
-import org.jetbrains.kotlin.codegen.coroutines.SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME
 import org.jetbrains.kotlin.codegen.coroutines.reportSuspensionPointInsideMonitor
 import org.jetbrains.kotlin.codegen.inline.addFakeContinuationConstructorCallMarker
 import org.jetbrains.kotlin.config.isReleaseCoroutines
-import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
-import org.jetbrains.kotlin.ir.builders.declarations.addValueParameter
-import org.jetbrains.kotlin.ir.builders.declarations.buildFunWithDescriptorForInlining
-import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
+import org.jetbrains.kotlin.ir.declarations.name
 import org.jetbrains.kotlin.ir.expressions.IrExpression
-import org.jetbrains.kotlin.ir.expressions.copyTypeArgumentsFrom
-import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrErrorExpressionImpl
-import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
 import org.jetbrains.kotlin.ir.types.createType
 import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
 import org.jetbrains.kotlin.ir.types.isUnit
-import org.jetbrains.kotlin.ir.util.*
-import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+import org.jetbrains.kotlin.ir.util.file
+import org.jetbrains.kotlin.ir.util.isSuspend
+import org.jetbrains.kotlin.ir.util.parentAsClass
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodGenericSignature
@@ -71,8 +64,9 @@ internal fun generateStateMachineForNamedFunction(
                || irFunction.origin == JvmLoweredDeclarationOrigin.SUSPEND_IMPL_STATIC_FUNCTION,
        internalNameForDispatchReceiver = classCodegen.visitor.thisName,
        putContinuationParameterToLvt = false,
-        disableTailCallOptimizationForFunctionReturningUnit = irFunction.returnType.isUnit() &&
-                irFunction.anyOfOverriddenFunctionsReturnsNonUnit()
+        disableTailCallOptimizationForFunctionReturningUnit = irFunction.suspendFunctionOriginal().let {
+            it.returnType.isUnit() && it.anyOfOverriddenFunctionsReturnsNonUnit()
+        }
    )
 }

@@ -134,109 +128,6 @@ internal fun IrFunction.isKnownToBeTailCall(): Boolean =
 internal fun IrFunction.shouldNotContainSuspendMarkers(): Boolean =
    isInvokeSuspendOfContinuation() || isKnownToBeTailCall()

-// Transform `suspend fun foo(params): RetType` into `fun foo(params, $completion: Continuation<RetType>): Any?`
-// the result is called 'view', just to be consistent with old backend.
-internal fun IrFunction.getOrCreateSuspendFunctionViewIfNeeded(context: JvmBackendContext): IrFunction {
-    if (!isSuspend || origin == JvmLoweredDeclarationOrigin.SUSPEND_FUNCTION_VIEW ||
-        origin == JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE_VIEW
-    ) return this
-    return context.suspendFunctionOriginalToView[this] ?: suspendFunctionView(context, true)
-}
-
-fun IrFunction.suspendFunctionView(context: JvmBackendContext, generateBody: Boolean): IrFunction {
-    require(this.isSuspend && this is IrSimpleFunction)
-    // For SuspendFunction{N}.invoke we need to generate INVOKEINTERFACE Function{N+1}.invoke(...Ljava/lang/Object;)...
-    // instead of INVOKEINTERFACE Function{N+1}.invoke(...Lkotlin/coroutines/Continuation;)...
-    val isInvokeOfNumberedSuspendFunction = (parent as? IrClass)?.defaultType?.isSuspendFunction() == true
-    // And we need to generate this function for callable references
-    val isBridgeInvokeOfCallableReference = origin == IrDeclarationOrigin.BRIDGE &&
-            (parent as? IrClass)?.origin == JvmLoweredDeclarationOrigin.FUNCTION_REFERENCE_IMPL
-    val continuationType = if (isInvokeOfNumberedSuspendFunction || isBridgeInvokeOfCallableReference)
-        context.irBuiltIns.anyNType
-    else
-        context.ir.symbols.continuationClass.createType(false, listOf(makeTypeProjection(returnType, Variance.INVARIANT)))
-    return buildFunWithDescriptorForInlining(descriptor) {
-        updateFrom(this@suspendFunctionView)
-        name = this@suspendFunctionView.name
-        origin = if (origin == JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE)
-            JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE_VIEW
-        else
-            JvmLoweredDeclarationOrigin.SUSPEND_FUNCTION_VIEW
-        returnType = context.irBuiltIns.anyNType
-    }.also {
-        it.parent = parent
-
-        it.annotations.addAll(annotations)
-
-        it.copyAttributes(this)
-        it.copyTypeParametersFrom(this)
-
-        // Copy the value parameters and insert the continuation parameter. The continuation parameter
-        // goes before the default argument mask(s) and handler for default argument stubs.
-        // TODO: It would be nice if AddContinuationLowering could insert the continuation argument before default stub generation.
-        // That would avoid the reshuffling both here and in createSuspendFunctionClassViewIfNeeded.
-        var continuationValueParam: IrValueParameter? = null
-        it.copyValueParametersInsertingContinuationFrom(this) {
-            continuationValueParam = it.addValueParameter(SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME, continuationType)
-        }
-
-        if (generateBody) {
-            // Add the suspend function view to the map before transforming the body to make sure
-            // that recursive suspend functions do not lead to unbounded recursion at compile time.
-            context.recordSuspendFunctionView(this, it)
-
-            val valueParametersMapping = explicitParameters.zip(it.explicitParameters.filter { it != continuationValueParam }).toMap()
-            it.body = body?.deepCopyWithSymbols(this)
-            it.body?.transformChildrenVoid(object : VariableRemapper(valueParametersMapping) {
-                // Do not cross class boundaries inside functions. Otherwise, callable references will try to access wrong $completion.
-                override fun visitClass(declaration: IrClass): IrStatement = declaration
-
-                override fun visitCall(expression: IrCall): IrExpression =
-                    super.visitCall(expression.createSuspendFunctionCallViewIfNeeded(context, it, callerIsInlineLambda = false))
-            })
-        }
-    }
-}
-
-internal fun IrCall.createSuspendFunctionCallViewIfNeeded(
-    context: JvmBackendContext,
-    caller: IrFunction,
-    callerIsInlineLambda: Boolean
-): IrCall {
-    if (!isSuspend) return this
-    val view = (symbol.owner as IrSimpleFunction).getOrCreateSuspendFunctionViewIfNeeded(context)
-    if (view == symbol.owner) return this
-    return IrCallImpl(startOffset, endOffset, view.returnType, view.symbol, superQualifierSymbol = superQualifierSymbol).also {
-        it.copyTypeArgumentsFrom(this)
-        it.dispatchReceiver = dispatchReceiver
-        it.extensionReceiver = extensionReceiver
-        // Locate the caller continuation parameter. The continuation parameter is before default argument mask(s) and handler params.
-        val callerNumberOfMasks = caller.valueParameters.count { it.origin == IrDeclarationOrigin.MASK_FOR_DEFAULT_FUNCTION }
-        val callerContinuationIndex = caller.valueParameters.size - 1 - (if (callerNumberOfMasks != 0) callerNumberOfMasks + 1 else 0)
-        val continuationParameter =
-            when {
-                caller.isInvokeSuspendOfLambda() || caller.isInvokeSuspendOfContinuation() ||
-                        caller.isInvokeSuspendForInlineOfLambda() ->
-                    IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, caller.dispatchReceiverParameter!!.symbol)
-                callerIsInlineLambda -> context.fakeContinuation
-                else -> IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, caller.valueParameters[callerContinuationIndex].symbol)
-            }
-        // If the suspend function view that we are calling has default parameters, we need to make sure to pass the
-        // continuation before the default parameter mask(s) and handler.
-        val numberOfMasks = view.valueParameters.count { it.origin == IrDeclarationOrigin.MASK_FOR_DEFAULT_FUNCTION }
-        val continuationIndex = valueArgumentsCount - (if (numberOfMasks != 0) numberOfMasks + 1 else 0)
-        for (i in 0 until continuationIndex) {
-            it.putValueArgument(i, getValueArgument(i))
-        }
-        it.putValueArgument(continuationIndex, continuationParameter)
-        if (numberOfMasks != 0) {
-            for (i in 0 until numberOfMasks + 1) {
-                it.putValueArgument(valueArgumentsCount + i - 1, getValueArgument(continuationIndex + i))
-            }
-        }
-    }
-}
-
 internal fun createFakeContinuation(context: JvmBackendContext): IrExpression = IrErrorExpressionImpl(
    UNDEFINED_OFFSET,
    UNDEFINED_OFFSET,
@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.backend.jvm.codegen
 import org.jetbrains.kotlin.backend.common.lower.BOUND_RECEIVER_PARAMETER
 import org.jetbrains.kotlin.backend.common.lower.LocalDeclarationsLowering
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
+import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin.*
 import org.jetbrains.kotlin.backend.jvm.intrinsics.IrIntrinsicMethods
 import org.jetbrains.kotlin.backend.jvm.intrinsics.JavaClassProperty
 import org.jetbrains.kotlin.backend.jvm.lower.MultifileFacadeFileEntry
@@ -52,6 +53,7 @@ import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeSystemCommonBackendContext
 import org.jetbrains.kotlin.types.model.TypeParameterMarker
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.kotlin.utils.keysToMap
@@ -166,7 +168,7 @@ class ExpressionCodegen(

    // TODO remove
    fun gen(expression: IrExpression, type: Type, irType: IrType, data: BlockInfo): StackValue {
-        if (expression === context.fakeContinuation) {
+        if (expression.attributeOwnerId === context.fakeContinuation) {
            addFakeContinuationMarker(mv)
        } else {
            expression.accept(this, data).coerce(type, irType).materialize()
@@ -219,10 +221,11 @@ class ExpressionCodegen(
    }

    private fun generateFakeContinuationConstructorIfNeeded() {
-        if (irFunction.origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE_VIEW) return
+        if (irFunction.origin != FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE) return
        val continuationClass = classCodegen.irClass.functions.find {
-            it.name.asString() == irFunction.name.asString().removeSuffix(FOR_INLINE_SUFFIX)
-        }?.body?.statements?.get(0) ?: error("could not find continuation for ${irFunction.render()}")
+            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,
@@ -254,12 +257,10 @@ class ExpressionCodegen(
        if (notCallableFromJava)
            return

-        // Do not generate non-null checks for suspend function views. When resumed the arguments
+        // Do not generate non-null checks for suspend functions. When resumed the arguments
        // will be null and the actual values are taken from the continuation.
-        val isSuspendFunctionView = irFunction.origin == JvmLoweredDeclarationOrigin.SUSPEND_FUNCTION_VIEW ||
-                irFunction.origin == JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE_VIEW

-        if (isSuspendFunctionView)
+        if (irFunction.isSuspend)
            return

        irFunction.extensionReceiverParameter?.let { generateNonNullAssertion(it) }
@@ -361,10 +362,6 @@ class ExpressionCodegen(
    override fun visitContainerExpression(expression: IrContainerExpression, data: BlockInfo) =
        visitStatementContainer(expression, data).coerce(expression.type)

-    override fun visitCall(expression: IrCall, data: BlockInfo): PromisedValue {
-        return visitFunctionAccess(expression.createSuspendFunctionCallViewIfNeeded(context, irFunction, inlinedInto != null), data)
-    }
-
    override fun visitFunctionAccess(expression: IrFunctionAccessExpression, data: BlockInfo): PromisedValue {
        classCodegen.context.irIntrinsics.getIntrinsic(expression.symbol)
            ?.invoke(expression, this, data)?.let { return it }
@@ -445,7 +442,7 @@ class ExpressionCodegen(
            // Check return type of non-lowered suspend call, in order to replace the result of the call with Unit,
            // otherwise, it would seem like the call returns non-unit upon resume.
            // See box/coroutines/tailCallOptimization/unit tests.
-            if (context.suspendFunctionViewToOriginal[expression.symbol.owner]?.returnType?.isUnit() == true) {
+            if (((expression.symbol.owner as? IrSimpleFunction)?.attributeOwnerId as? IrSimpleFunction)?.returnType?.isUnit() == true) {
                addReturnsUnitMarker(mv)
            }

@@ -486,11 +483,11 @@ class ExpressionCodegen(
            // Inside $$forInline functions crossinline calls are (usually) not suspension points, otherwise, flow will be pessimized
            (dispatchReceiver as? IrGetField)?.symbol?.owner?.origin ==
                    LocalDeclarationsLowering.DECLARATION_ORIGIN_FIELD_FOR_CROSSINLINE_CAPTURED_VALUE ->
-                irFunction.origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
+                irFunction.origin != FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
            // The same goes for inline lambdas
            (dispatchReceiver as? IrGetValue)?.let {
                it.origin == IrStatementOrigin.VARIABLE_AS_FUNCTION && (it.symbol.owner as? IrValueParameter)?.isNoinline != true
-            } == true -> irFunction.origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE
+            } == true -> irFunction.origin != FOR_INLINE_STATE_MACHINE_TEMPLATE
            // Otherwise, this is normal inline call, ergo not a suspension point
            else -> false
        }
@@ -636,10 +633,21 @@ class ExpressionCodegen(

    override fun visitReturn(expression: IrReturn, data: BlockInfo): PromisedValue {
        val returnTarget = expression.returnTargetSymbol.owner
-        val owner =
-            (returnTarget as? IrFunction
+        var owner =
+            returnTarget as? IrFunction
                ?: (returnTarget as? IrReturnableBlock)?.inlineFunctionSymbol?.owner
-                ?: error("Unsupported IrReturnTarget: $returnTarget")).getOrCreateSuspendFunctionViewIfNeeded(context)
+                ?: error("Unsupported IrReturnTarget: $returnTarget")
+        // $$forInline and $suspendImpl functions share the same attributes as originals, but the name is different
+        // fix the return target. TODO: attributes seem to be overloaded. We rely on their uniqueness across transformations to find views
+        // TODO: but in this case, they are not unique enough.
+        if ((irFunction.origin == FOR_INLINE_STATE_MACHINE_TEMPLATE ||
+                    irFunction.origin == SUSPEND_IMPL_STATIC_FUNCTION ||
+                    irFunction.origin == FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE) &&
+            owner.isSuspend && irFunction is IrSimpleFunction && owner is IrSimpleFunction &&
+            irFunction.attributeOwnerId == owner.attributeOwnerId
+        ) {
+            owner = irFunction
+        }
        //TODO: should be owner != irFunction
        val isNonLocalReturn =
            methodSignatureMapper.mapFunctionName(owner) != methodSignatureMapper.mapFunctionName(irFunction)
@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */

@@ -53,10 +53,9 @@ open class FunctionCodegen(
        }

    private fun doGenerate(): JvmMethodGenericSignature {
-        val functionView = irFunction.getOrCreateSuspendFunctionViewIfNeeded(context)
-        val signature = classCodegen.methodSignatureMapper.mapSignatureWithGeneric(functionView)
+        val signature = classCodegen.methodSignatureMapper.mapSignatureWithGeneric(irFunction)

-        val flags = calculateMethodFlags(functionView.isStatic)
+        val flags = calculateMethodFlags(irFunction.isStatic)
        var methodVisitor = createMethod(flags, signature)

        if (state.generateParametersMetadata && flags.and(Opcodes.ACC_SYNTHETIC) == 0) {
@@ -75,14 +74,14 @@ open class FunctionCodegen(
                    )
                }
            }.genAnnotations(
-                functionView,
+                irFunction,
                signature.asmMethod.returnType,
                irFunction.returnType
            )
            // Not generating parameter annotations for default stubs fixes KT-7892, though
            // this certainly looks like a workaround for a javac bug.
            if (irFunction !is IrConstructor || !irFunction.parentAsClass.shouldNotGenerateConstructorParameterAnnotations()) {
-                generateParameterAnnotations(functionView, methodVisitor, signature, classCodegen, context)
+                generateParameterAnnotations(irFunction, methodVisitor, signature, classCodegen, context)
            }
        }

@@ -107,7 +106,7 @@ open class FunctionCodegen(
                )
                else -> methodVisitor
            }
-            ExpressionCodegen(functionView, signature, frameMap, InstructionAdapter(methodVisitor), classCodegen, inlinedInto).generate()
+            ExpressionCodegen(irFunction, signature, frameMap, InstructionAdapter(methodVisitor), classCodegen, inlinedInto).generate()
            methodVisitor.visitMaxs(-1, -1)
            if (irFunction.hasContinuation()) {
                context.continuationClassBuilders[continuationClass().attributeOwnerId].sure {
@@ -151,7 +150,8 @@ open class FunctionCodegen(
            // $$forInline functions never have a continuation
            origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE

-    private fun continuationClass() = irFunction.body!!.statements.firstIsInstance<IrClass>()
+    private fun continuationClass() =
+        irFunction.body!!.statements.firstIsInstance<IrClass>()

    private fun IrFunction.getVisibilityForDefaultArgumentStub(): Int =
        when (visibility) {
@@ -260,17 +260,16 @@ open class FunctionCodegen(

    private fun createFrameMapWithReceivers(): IrFrameMap {
        val frameMap = IrFrameMap()
-        val functionView = irFunction.getOrCreateSuspendFunctionViewIfNeeded(context)

        if (irFunction is IrConstructor) {
            frameMap.enterDispatchReceiver(irFunction.constructedClass.thisReceiver!!)
-        } else if (functionView.dispatchReceiverParameter != null) {
-            frameMap.enterDispatchReceiver(functionView.dispatchReceiverParameter!!)
+        } else if (irFunction.dispatchReceiverParameter != null) {
+            frameMap.enterDispatchReceiver(irFunction.dispatchReceiverParameter!!)
        }
-        functionView.extensionReceiverParameter?.let {
+        irFunction.extensionReceiverParameter?.let {
            frameMap.enter(it, classCodegen.typeMapper.mapType(it))
        }
-        for (parameter in functionView.valueParameters) {
+        for (parameter in irFunction.valueParameters) {
            frameMap.enter(parameter, classCodegen.typeMapper.mapType(parameter.type))
        }

@@ -1,24 +1,22 @@
 /*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */

 package org.jetbrains.kotlin.backend.jvm.codegen

 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.codegen.*
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.backend.jvm.lower.suspendFunctionViewOrStub
 import org.jetbrains.kotlin.backend.jvm.ir.isInlineParameter
 import org.jetbrains.kotlin.backend.jvm.ir.isLambda
-import org.jetbrains.kotlin.codegen.*
 import org.jetbrains.kotlin.codegen.inline.*
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
-import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
-import org.jetbrains.kotlin.ir.declarations.IrFunction
-import org.jetbrains.kotlin.ir.declarations.IrValueParameter
-import org.jetbrains.kotlin.ir.declarations.IrVariable
+import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.descriptors.WrappedSimpleFunctionDescriptor
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.types.IrType
@@ -251,7 +249,7 @@ class IrExpressionLambdaImpl(
            capturedParamDesc(param.name.asString(), typeMapper.mapType(param.type))
        }

-    private val loweredMethod = methodSignatureMapper.mapAsmMethod(function.getOrCreateSuspendFunctionViewIfNeeded(context))
+    private val loweredMethod = methodSignatureMapper.mapAsmMethod(function.suspendFunctionViewOrStub(context))

    val capturedParamsInDesc: List<Type> = if (isBoundCallableReference) {
        loweredMethod.argumentTypes.take(1)
@@ -274,7 +272,7 @@ class IrExpressionLambdaImpl(
    override val hasDispatchReceiver: Boolean = false

    override fun getInlineSuspendLambdaViewDescriptor(): FunctionDescriptor {
-        return function.getOrCreateSuspendFunctionViewIfNeeded(context).descriptor
+        return function.suspendFunctionViewOrStub(context).descriptor
    }
 }

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */

@@ -9,6 +9,7 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import org.jetbrains.kotlin.backend.common.CodegenUtil
 import org.jetbrains.kotlin.backend.common.ir.ir2string
+import org.jetbrains.kotlin.backend.jvm.lower.getOrCreateSuspendFunctionViewIfNeeded
 import org.jetbrains.kotlin.codegen.BaseExpressionCodegen
 import org.jetbrains.kotlin.codegen.ClassBuilder
 import org.jetbrains.kotlin.codegen.OwnerKind
@@ -86,7 +87,7 @@ class IrSourceCompilerForInline(
                function.parentAsClass.functions.find {
                    it.name.asString() == function.name.asString() + FOR_INLINE_SUFFIX &&
                            it.attributeOwnerId == (function as? IrAttributeContainer)?.attributeOwnerId
-                } ?: function
+                } ?: function.getOrCreateSuspendFunctionViewIfNeeded(classCodegen.context)
            else function
        val functionCodegen = object : FunctionCodegen(forInlineFunction, classCodegen, codegen.takeIf { isLambda }) {
            override fun createMethod(flags: Int, signature: JvmMethodGenericSignature): MethodVisitor {
@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.backend.jvm.ir.getJvmNameFromAnnotation
 import org.jetbrains.kotlin.backend.jvm.ir.hasJvmDefault
 import org.jetbrains.kotlin.backend.jvm.ir.propertyIfAccessor
+import org.jetbrains.kotlin.backend.jvm.lower.*
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.codegen.AsmUtil
 import org.jetbrains.kotlin.codegen.JvmCodegenUtil
@@ -208,13 +209,6 @@ class MethodSignatureMapper(private val context: JvmBackendContext) {
            }
        }

-        if (function.isSuspend && function.origin != JvmLoweredDeclarationOrigin.SUSPEND_FUNCTION_VIEW &&
-            function.origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE_VIEW &&
-            (function.parent as? IrClass)?.origin != JvmLoweredDeclarationOrigin.FUNCTION_REFERENCE_IMPL
-        ) {
-            return mapSignature(function.suspendFunctionView(context, false), skipGenericSignature)
-        }
-
        val sw = if (skipGenericSignature) JvmSignatureWriter() else BothSignatureWriter(BothSignatureWriter.Mode.METHOD)

        typeMapper.writeFormalTypeParameters(function.typeParameters, sw)
@@ -300,7 +294,7 @@ class MethodSignatureMapper(private val context: JvmBackendContext) {
        }

    fun mapToCallableMethod(caller: IrFunction, expression: IrFunctionAccessExpression): IrCallableMethod {
-        val callee = expression.symbol.owner.getOrCreateSuspendFunctionViewIfNeeded(context)
+        val callee = expression.symbol.owner
        val calleeParent = callee.parentAsClass
        val owner = typeMapper.mapClass(calleeParent)

@@ -14,7 +14,9 @@ import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.backend.jvm.ir.IrInlineReferenceLocator
 import org.jetbrains.kotlin.backend.jvm.ir.defaultValue
-import org.jetbrains.kotlin.backend.jvm.codegen.isInlineIrBlock
+import org.jetbrains.kotlin.backend.jvm.codegen.isInvokeSuspendForInlineOfLambda
+import org.jetbrains.kotlin.backend.jvm.codegen.isInvokeSuspendOfContinuation
+import org.jetbrains.kotlin.backend.jvm.codegen.isInvokeSuspendOfLambda
 import org.jetbrains.kotlin.backend.jvm.localDeclarationsPhase
 import org.jetbrains.kotlin.codegen.coroutines.*
 import org.jetbrains.kotlin.codegen.inline.coroutines.FOR_INLINE_SUFFIX
@@ -28,9 +30,7 @@ import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.builders.declarations.*
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
-import org.jetbrains.kotlin.ir.expressions.impl.IrGetFieldImpl
-import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
-import org.jetbrains.kotlin.ir.expressions.impl.IrReturnImpl
+import org.jetbrains.kotlin.ir.expressions.impl.*
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
@@ -47,19 +47,45 @@ import org.jetbrains.kotlin.util.OperatorNameConventions
 internal val addContinuationPhase = makeIrFilePhase(
    ::AddContinuationLowering,
    "AddContinuation",
-    "Add continuation classes to suspend functions and transform suspend lambdas into continuations",
+    "Add continuation classes and parameters to suspend functions and transform suspend lambdas into continuations",
    prerequisite = setOf(localDeclarationsPhase, tailCallOptimizationPhase)
 )

 private class AddContinuationLowering(private val context: JvmBackendContext) : FileLoweringPass {
-
    override fun lower(irFile: IrFile) {
-        val (suspendLambdas, inlineLambdas) = markSuspendLambdas(irFile)
-        transformSuspendFunctions(irFile, (suspendLambdas.map { it.function } + inlineLambdas).toSet())
-        transformReferencesToSuspendLambdas(irFile, suspendLambdas)
+        val suspendLambdas = findSuspendAndInlineLambdas(irFile)
+        addContinuationObjectAndContinuationParameterToSuspendFunctions(irFile)
+        transformSuspendLambdasIntoContinuations(irFile, suspendLambdas)
+        addContinuationParameterToSuspendCallsAndUpdateNonLocalReturns(irFile)
    }

-    private fun transformReferencesToSuspendLambdas(irFile: IrFile, suspendLambdas: List<SuspendLambdaInfo>) {
+    private fun addContinuationParameterToSuspendCallsAndUpdateNonLocalReturns(irFile: IrFile) {
+        irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
+            val functionStack = mutableListOf<IrFunction>()
+
+            override fun visitFunction(declaration: IrFunction): IrStatement {
+                functionStack.push(declaration)
+                return super.visitFunction(declaration).also { functionStack.pop() }
+            }
+
+            override fun visitCall(expression: IrCall): IrExpression {
+                // This is a property, no need to add continuation parameter, since this cannot be suspend call
+                if (functionStack.isEmpty()) return super.visitCall(expression)
+                val caller = functionStack.peek()!!
+                return (super.visitCall(expression) as IrCall)
+                    .createSuspendFunctionCallViewIfNeeded(context, caller)
+            }
+
+            override fun visitReturn(expression: IrReturn): IrExpression {
+                val ret = super.visitReturn(expression) as IrReturn
+                val irFunction = expression.returnTargetSymbol.owner as? IrFunction ?: return ret
+                val targetViewOrStub = irFunction.suspendFunctionViewOrStub(context)
+                return IrReturnImpl(ret.startOffset, ret.endOffset, ret.type, targetViewOrStub.symbol, ret.value)
+            }
+        })
+    }
+
+    private fun transformSuspendLambdasIntoContinuations(irFile: IrFile, suspendLambdas: List<SuspendLambdaInfo>) {
        for (lambda in suspendLambdas) {
            (lambda.function.parent as IrDeclarationContainer).declarations.remove(lambda.function)
        }
@@ -508,9 +534,13 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                            UNDEFINED_OFFSET, UNDEFINED_OFFSET, backendContext
                        )
                    }
-                    for ((i, parameter) in irFunction.valueParameters.withIndex()) {
+                    for ((i, parameter) in irFunction.valueParameters.dropLast(1).withIndex()) {
                        it.putValueArgument(i, parameter.type.defaultValue(UNDEFINED_OFFSET, UNDEFINED_OFFSET, backendContext))
                    }
+                    it.putValueArgument(
+                        irFunction.valueParameters.size - 1,
+                        IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, function.dispatchReceiverParameter!!.symbol)
+                    )
                    if (isStaticSuspendImpl) {
                        it.putValueArgument(0, capturedThisValue)
                    }
@@ -529,6 +559,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
            copyMetadata = false
        )
        static.body = irFunction.moveBodyTo(static)
+        static.copyAttributes(irFunction)
        // Rewrite the body of the original suspend method to forward to the new static method.
        irFunction.body = context.createIrBuilder(irFunction.symbol).irBlockBody {
            +irReturn(irCall(static).also {
@@ -547,11 +578,10 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                }
            })
        }
-        static.copyAttributes(irFunction)
        return static
    }

-    private fun transformSuspendFunctions(irFile: IrFile, suspendAndInlineLambdas: Set<IrFunction>) {
+    private fun addContinuationObjectAndContinuationParameterToSuspendFunctions(irFile: IrFile) {
        irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
            private val functionsStack = arrayListOf<IrFunction>()
            private val suspendFunctionsCapturingCrossinline = mutableSetOf<IrFunction>()
@@ -573,39 +603,45 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                functionsStack.pop()
                if (skip(function)) return function

-                function as IrSimpleFunction
+                // TODO: This does not work for DEFAULT_IMPLS
+                val view =
+                    (if (function.body == null) function.suspendFunctionStub(context, false)
+                    else function.getOrCreateSuspendFunctionViewIfNeeded(context)) as IrSimpleFunction
+
+                if (withoutContinuationClass(function)) return view
+
                if (function in suspendFunctionsCapturingCrossinline || function.isInline) {
-                    val newFunction = buildFunWithDescriptorForInlining(function.descriptor) {
-                        name = Name.identifier(function.name.asString() + FOR_INLINE_SUFFIX)
-                        returnType = function.returnType
-                        modality = function.modality
-                        isSuspend = function.isSuspend
-                        isInline = function.isInline
+                    val newFunction = buildFunWithDescriptorForInlining(view.descriptor) {
+                        name = Name.identifier(view.name.asString() + FOR_INLINE_SUFFIX)
+                        returnType = view.returnType
+                        modality = view.modality
+                        isSuspend = view.isSuspend
+                        isInline = view.isInline
                        origin =
-                            if (function.isInline) JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE
+                            if (view.isInline) JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE
                            else JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
                    }.apply {
-                        copyTypeParameters(function.typeParameters)
-                        dispatchReceiverParameter = function.dispatchReceiverParameter?.copyTo(this)
-                        extensionReceiverParameter = function.extensionReceiverParameter?.copyTo(this)
-                        function.valueParameters.mapTo(valueParameters) { it.copyTo(this) }
-                        body = function.copyBodyTo(this)
-                        copyAttributes(function)
+                        copyTypeParameters(view.typeParameters)
+                        dispatchReceiverParameter = view.dispatchReceiverParameter?.copyTo(this)
+                        extensionReceiverParameter = view.extensionReceiverParameter?.copyTo(this)
+                        view.valueParameters.mapTo(valueParameters) { it.copyTo(this) }
+                        body = view.copyBodyTo(this)
+                        copyAttributes(view)
                    }
                    registerNewFunction(newFunction)
                }

-                val newFunction = if (function.isOverridable) {
+                val newFunction = if ((function as IrSimpleFunction).isOverridable) {
                    // Create static method for the suspend state machine method so that reentering the method
                    // does not lead to virtual dispatch to the wrong method.
-                    registerNewFunction(function)
-                    createStaticSuspendImpl(function)
-                } else function
+                    registerNewFunction(view)
+                    createStaticSuspendImpl(view)
+                } else view

                newFunction.body = context.createIrBuilder(newFunction.symbol).irBlockBody {
                    +generateContinuationClassForNamedFunction(
                        newFunction,
-                        function.dispatchReceiverParameter,
+                        view.dispatchReceiverParameter,
                        declaration as IrAttributeContainer
                    )
                    for (statement in newFunction.body!!.statements) {
@@ -615,11 +651,16 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                return newFunction
            }

-            private fun skip(function: IrFunction) =
-                !function.isSuspend || function in suspendAndInlineLambdas || function.body == null ||
-                        function.origin == JvmLoweredDeclarationOrigin.DEFAULT_IMPLS_BRIDGE ||
+            // TODO: Merge with `knownToBeTailCall`
+            private fun withoutContinuationClass(function: IrFunction): Boolean =
+                function.body == null ||
+                        function.parentAsClass.origin == JvmLoweredDeclarationOrigin.FUNCTION_REFERENCE_IMPL ||
                        function.origin == IrDeclarationOrigin.FUNCTION_FOR_DEFAULT_PARAMETER ||
-                        function.parentAsClass.origin == JvmLoweredDeclarationOrigin.FUNCTION_REFERENCE_IMPL
+                        function.origin == JvmLoweredDeclarationOrigin.DEFAULT_IMPLS_BRIDGE
+
+            private fun skip(function: IrFunction) =
+                !function.isSuspend ||
+                        function.origin == IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA

            private fun registerNewFunction(function: IrFunction) {
                functionsToAdd.peek()!!.add(function)
@@ -638,7 +679,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
        })
    }

-    private fun markSuspendLambdas(irElement: IrElement): Pair<List<SuspendLambdaInfo>, List<IrFunction>> {
+    private fun findSuspendAndInlineLambdas(irElement: IrElement): List<SuspendLambdaInfo> {
        val suspendLambdas = arrayListOf<SuspendLambdaInfo>()
        val capturesCrossinline = mutableSetOf<IrCallableReference>()
        val inlineReferences = mutableSetOf<IrCallableReference>()
@@ -678,7 +719,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                }
            }
        })
-        return suspendLambdas to inlineReferences.map { it.symbol.owner as IrFunction }
+        return suspendLambdas
    }

    private class SuspendLambdaInfo(
@@ -689,4 +730,126 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
    ) {
        lateinit var constructor: IrConstructor
    }
+}
+
+// Transform `suspend fun foo(params): RetType` into `fun foo(params, $completion: Continuation<RetType>): Any?`
+// the result is called 'view', just to be consistent with old backend.
+internal fun IrFunction.getOrCreateSuspendFunctionViewIfNeeded(context: JvmBackendContext): IrFunction {
+    if (!isSuspend) return this
+    return context.suspendFunctionOriginalToView[suspendFunctionOriginal()] ?: suspendFunctionView(context)
+}
+
+private fun IrFunction.getOrCreateSuspendFunctionStub(context: JvmBackendContext): IrFunction {
+    if (!isSuspend) return this
+    return context.suspendFunctionOriginalToStub[suspendFunctionOriginal()] ?: suspendFunctionStub(context)
+}
+
+internal fun IrFunction.suspendFunctionOriginal(): IrFunction {
+    require(isSuspend && this is IrSimpleFunction)
+    return attributeOwnerId as IrFunction
+}
+
+private fun IrFunction.suspendFunctionStub(context: JvmBackendContext): IrFunction {
+    require(this.isSuspend && this is IrSimpleFunction)
+    return buildFunWithDescriptorForInlining(descriptor) {
+        updateFrom(this@suspendFunctionStub)
+        name = this@suspendFunctionStub.name
+        origin = this@suspendFunctionStub.origin
+        returnType = context.irBuiltIns.anyNType
+    }.also { function ->
+        function.parent = parent
+
+        function.annotations.addAll(annotations)
+        function.metadata = metadata
+
+        function.copyAttributes(this)
+        function.copyTypeParametersFrom(this)
+
+        if (origin != JvmLoweredDeclarationOrigin.DEFAULT_IMPLS_BRIDGE) {
+            function.overriddenSymbols
+                .addAll(overriddenSymbols.map { it.owner.suspendFunctionViewOrStub(context).symbol as IrSimpleFunctionSymbol })
+        }
+
+        // Copy the value parameters and insert the continuation parameter. The continuation parameter
+        // goes before the default argument mask(s) and handler for default argument stubs.
+        // TODO: It would be nice if AddContinuationLowering could insert the continuation argument before default stub generation.
+        // That would avoid the reshuffling both here and in createSuspendFunctionCallViewIfNeeded.
+        function.copyValueParametersInsertingContinuationFrom(this) {
+            function.addValueParameter(SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME, continuationType(context))
+        }
+
+        context.recordSuspendFunctionViewStub(this, function)
+    }
+}
+
+private fun IrFunction.continuationType(context: JvmBackendContext): IrType {
+    // For SuspendFunction{N}.invoke we need to generate INVOKEINTERFACE Function{N+1}.invoke(...Ljava/lang/Object;)...
+    // instead of INVOKEINTERFACE Function{N+1}.invoke(...Lkotlin/coroutines/Continuation;)...
+    val isInvokeOfNumberedSuspendFunction = (parent as? IrClass)?.defaultType?.isSuspendFunction() == true
+    val isInvokeOfNumberedFunction = (parent as? IrClass)?.fqNameWhenAvailable?.asString()?.let {
+        it.startsWith("kotlin.jvm.functions.Function") && it.removePrefix("kotlin.jvm.functions.Function").all { c -> c.isDigit() }
+    } == true
+    return if (isInvokeOfNumberedSuspendFunction || isInvokeOfNumberedFunction)
+        context.irBuiltIns.anyNType
+    else
+        context.ir.symbols.continuationClass.typeWith(returnType)
+}
+
+private fun IrFunction.suspendFunctionView(context: JvmBackendContext): IrFunction {
+    require(isSuspend && this is IrSimpleFunction)
+    return getOrCreateSuspendFunctionStub(context).also { function ->
+        context.recordSuspendFunctionView(this, function)
+
+        val continuationParameter =
+            function.valueParameters.find { it.origin == IrDeclarationOrigin.MASK_FOR_DEFAULT_FUNCTION }
+                ?.let { function.valueParameters[it.index - 1] } ?: function.valueParameters.last()
+
+        function.body =
+            moveBodyTo(function, explicitParameters.zip(function.explicitParameters.filter { it != continuationParameter }).toMap())
+    }
+}
+
+fun IrFunction.suspendFunctionViewOrStub(context: JvmBackendContext): IrFunction {
+    if (!isSuspend) return this
+    return context.suspendFunctionOriginalToView[suspendFunctionOriginal()] ?: getOrCreateSuspendFunctionStub(context)
+}
+
+private fun IrCall.createSuspendFunctionCallViewIfNeeded(context: JvmBackendContext, caller: IrFunction): IrCall {
+    // Calls inside continuation are already generated with continuation parameter as well as calls to suspendImpls
+    if (!isSuspend || caller.isInvokeSuspendOfContinuation()
+        || symbol.owner.origin == JvmLoweredDeclarationOrigin.SUSPEND_IMPL_STATIC_FUNCTION
+        || (symbol.owner.valueParameters.lastOrNull()?.name?.asString() == SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME)
+    ) return this
+    val view = (symbol.owner as IrSimpleFunction).suspendFunctionViewOrStub(context)
+    if (view == symbol.owner) return this
+
+    return IrCallImpl(startOffset, endOffset, view.returnType, view.symbol, superQualifierSymbol = superQualifierSymbol).also {
+        it.copyTypeArgumentsFrom(this)
+        it.copyAttributes(this)
+        it.dispatchReceiver = dispatchReceiver
+        it.extensionReceiver = extensionReceiver
+        val callerNumberOfMasks = caller.valueParameters.count { it.origin == IrDeclarationOrigin.MASK_FOR_DEFAULT_FUNCTION }
+        val callerContinuationIndex = caller.valueParameters.size - 1 - (if (callerNumberOfMasks != 0) callerNumberOfMasks + 1 else 0)
+        val continuationParameter =
+            when {
+                caller.isInvokeSuspendOfLambda() || caller.isInvokeSuspendForInlineOfLambda() ->
+                    IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, caller.dispatchReceiverParameter!!.symbol)
+                // At this point the only LOCAL_FUNCTION_FOR_LAMBDAs are inline and crossinline lambdas
+                caller.origin == IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA -> context.fakeContinuation
+                else -> IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, caller.valueParameters[callerContinuationIndex].symbol)
+            }
+        // If the suspend function view that we are calling has default parameters, we need to make sure to pass the
+        // continuation before the default parameter mask(s) and handler.
+        val numberOfMasks = view.valueParameters.count { it.origin == IrDeclarationOrigin.MASK_FOR_DEFAULT_FUNCTION }
+        val continuationIndex = valueArgumentsCount - (if (numberOfMasks != 0) numberOfMasks + 1 else 0)
+        for (i in 0 until continuationIndex) {
+            it.putValueArgument(i, getValueArgument(i))
+        }
+        it.putValueArgument(continuationIndex, continuationParameter)
+        if (numberOfMasks != 0) {
+            for (i in 0 until numberOfMasks + 1) {
+                it.putValueArgument(valueArgumentsCount + i - 1, getValueArgument(continuationIndex + i))
+            }
+        }
+    }
 }
@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2018 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */

@@ -37,10 +37,13 @@ import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
 import org.jetbrains.kotlin.ir.symbols.impl.IrSimpleFunctionSymbolImpl
 import org.jetbrains.kotlin.ir.symbols.impl.IrValueParameterSymbolImpl
 import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.types.getClass
+import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
 import org.jetbrains.kotlin.ir.types.isNullable
 import org.jetbrains.kotlin.ir.types.makeNullable
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.org.objectweb.asm.Type
@@ -424,7 +427,11 @@ private class BridgeLowering(val context: JvmBackendContext) : ClassLoweringPass
            IrValueParameterSymbolImpl(descriptor),
            name,
            index,
-            type.eraseTypeParameters(),
+            // SuspendFunction{N} is Function{N+1} at runtime, thus, when we generate a bridge for suspend callable reference,
+            // we need to replace type of its continuation parameter with Any?
+            if (target.isSuspend && type.eraseTypeParameters().getClass()
+                    ?.fqNameWhenAvailable == DescriptorUtils.CONTINUATION_INTERFACE_FQ_NAME_RELEASE
+            ) context.irBuiltIns.anyNType else type.eraseTypeParameters(),
            varargElementType?.eraseTypeParameters(),
            isCrossinline,
            isNoinline
@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */

@@ -88,7 +88,7 @@ private class FunctionNVarargBridgeLowering(val context: JvmBackendContext) :

        // Add vararg invoke bridge
        val invokeFunction = declaration.functions.single {
-            it.name.asString() == "invoke" && it.valueParameters.size == superType.arguments.size - 1
+            it.name.asString() == "invoke" && it.valueParameters.size == superType.arguments.size - if (it.isSuspend) 0 else 1
        }
        invokeFunction.overriddenSymbols.clear()
        declaration.addBridge(invokeFunction, functionNInvokeFun.owner)
@@ -1,17 +1,6 @@
 /*
- * Copyright 2010-2016 JetBrains s.r.o.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */

 package org.jetbrains.kotlin.ir.util
@@ -179,6 +168,7 @@ open class DeepCopyIrTreeWithSymbols(
            declaration.overriddenSymbols.mapTo(overriddenSymbols) {
                symbolRemapper.getReferencedFunction(it) as IrSimpleFunctionSymbol
            }
+            copyAttributes(declaration)
            transformFunctionChildren(declaration)
        }

@@ -1,4 +1,5 @@
 // IGNORE_BACKEND_FIR: JVM_IR
+// IGNORE_BACKEND: JVM_IR
 // WITH_RUNTIME
 // WITH_COROUTINES
 // COMMON_COROUTINES_TEST
@@ -1,4 +1,5 @@
 // IGNORE_BACKEND_FIR: JVM_IR
+// IGNORE_BACKEND: JVM_IR
 // WITH_RUNTIME
 // WITH_COROUTINES
 // COMMON_COROUTINES_TEST
@@ -6,6 +6,7 @@ final class CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1$1 {
    field L$2: java.lang.Object
    field L$3: java.lang.Object
    field L$4: java.lang.Object
+    field L$5: java.lang.Object
    field label: int
    @org.jetbrains.annotations.NotNull field result: java.lang.Object
    synthetic final field this$0: CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1
@@ -121,6 +122,7 @@ final class CrossinlineKt$filter$$inlined$source$1$1 {
    field L$2: java.lang.Object
    field L$3: java.lang.Object
    field L$4: java.lang.Object
+    field L$5: java.lang.Object
    field label: int
    @org.jetbrains.annotations.NotNull field result: java.lang.Object
    synthetic final field this$0: CrossinlineKt$filter$$inlined$source$1
@@ -146,6 +148,7 @@ final class CrossinlineKt$filter$$inlined$source$2$1 {
    field L$2: java.lang.Object
    field L$3: java.lang.Object
    field L$4: java.lang.Object
+    field L$5: java.lang.Object
    field label: int
    @org.jetbrains.annotations.NotNull field result: java.lang.Object
    synthetic final field this$0: CrossinlineKt$filter$$inlined$source$2
@@ -310,7 +313,6 @@ final class CrossinlineKt$range$$inlined$source$1$1 {
    field L$1: java.lang.Object
    field L$2: java.lang.Object
    field L$3: java.lang.Object
-    field L$4: java.lang.Object
    field label: int
    @org.jetbrains.annotations.NotNull field result: java.lang.Object
    synthetic final field this$0: CrossinlineKt$range$$inlined$source$1
@@ -1,5 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
 // FILE: test.kt
 // COMMON_COROUTINES_TEST
 // WITH_RUNTIME
@@ -1,5 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
 // FILE: test.kt
 // COMMON_COROUTINES_TEST
 // WITH_RUNTIME
@@ -1,5 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
 // FILE: test.kt
 // COMMON_COROUTINES_TEST
 // WITH_RUNTIME