JVM_IR: make continuation detection more consistent

* unify various checks for whether a suspend function needs a continuation; * mark the continuation parameter with an origin (this also allows correctly computing the offset of the local in codegen -- see the new test); * when wrapping `suspend fun main`, use a suspend reference instead of a synthetic non-suspend lambda (required to correctly implement the previous point, as previously the continuation parameter was passed to main() itself correctly only because of very lenient checks in AddContinuationLowering).
2020-02-26 10:37:46 +01:00
parent c57f466494
commit 159d292ea7
13 changed files with 138 additions and 215 deletions
@@ -185,7 +185,7 @@ fun IrTypeParameter.copyToWithoutSuperTypes(
    }
 }

-private fun IrFunction.copyReceiverParametersFrom(from: IrFunction) {
+fun IrFunction.copyReceiverParametersFrom(from: IrFunction) {
    dispatchReceiverParameter = from.dispatchReceiverParameter?.let {
        IrValueParameterImpl(it.startOffset, it.endOffset, it.origin, it.descriptor, it.type, it.varargElementType).also {
            it.parent = this
@@ -200,22 +200,6 @@ fun IrFunction.copyValueParametersFrom(from: IrFunction) {
    valueParameters += from.valueParameters.map { it.copyTo(this, index = it.index + shift) }
 }

-fun IrFunction.copyValueParametersInsertingContinuationFrom(from: IrFunction, insertContinuation: () -> Unit) {
-    copyReceiverParametersFrom(from)
-    val shift = valueParameters.size
-    var additionalShift = 0
-    from.valueParameters.forEach {
-        // The continuation parameter goes before the default argument mask and handler.
-        if (it.origin == IrDeclarationOrigin.MASK_FOR_DEFAULT_FUNCTION && additionalShift == 0) {
-            insertContinuation()
-            additionalShift = 1
-        }
-        valueParameters += it.copyTo(this, index = it.index + shift + additionalShift)
-    }
-    // If there was no default argument mask and handler, the continuation goes last.
-    if (additionalShift == 0) insertContinuation()
-}
-
 fun IrFunction.copyParameterDeclarationsFrom(from: IrFunction) {
    assert(typeParameters.isEmpty())
    copyTypeParametersFrom(from)
@@ -16,7 +16,6 @@ import org.jetbrains.kotlin.codegen.coroutines.CoroutineTransformerMethodVisitor
 import org.jetbrains.kotlin.codegen.coroutines.INVOKE_SUSPEND_METHOD_NAME
 import org.jetbrains.kotlin.codegen.coroutines.SUSPEND_IMPL_NAME_SUFFIX
 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.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.declarations.*
@@ -30,13 +29,9 @@ import org.jetbrains.kotlin.ir.util.functions
 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
 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

 internal fun generateStateMachineForNamedFunction(
    irFunction: IrFunction,
@@ -101,45 +96,49 @@ internal fun generateStateMachineForLambda(
    )
 }

+internal fun IrFunction.continuationParameter(): IrValueParameter? = when {
+    isInvokeSuspendOfLambda() || isInvokeSuspendForInlineOfLambda() -> dispatchReceiverParameter
+    else -> valueParameters.singleOrNull { it.origin == JvmLoweredDeclarationOrigin.CONTINUATION_CLASS }
+}
+
 internal fun IrFunction.isInvokeSuspendOfLambda(): Boolean =
    name.asString() == INVOKE_SUSPEND_METHOD_NAME && parentAsClass.origin == JvmLoweredDeclarationOrigin.SUSPEND_LAMBDA

-internal fun IrFunction.isInvokeSuspendForInlineOfLambda(): Boolean =
+private fun IrFunction.isInvokeSuspendForInlineOfLambda(): Boolean =
    origin == JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE
            && parentAsClass.origin == JvmLoweredDeclarationOrigin.SUSPEND_LAMBDA

 internal fun IrFunction.isInvokeSuspendOfContinuation(): Boolean =
    name.asString() == INVOKE_SUSPEND_METHOD_NAME && parentAsClass.origin == JvmLoweredDeclarationOrigin.CONTINUATION_CLASS

-internal fun IrFunction.isInvokeOfSuspendCallableReference(): Boolean = isSuspend && name.asString() == "invoke" &&
-        parentAsClass.origin == JvmLoweredDeclarationOrigin.FUNCTION_REFERENCE_IMPL
-
-// Wrapper of suspend main is always tail-call and it is not generated as suspend function.
-private fun IrFunction.isInvokeOfSuspendMainWrapper(): Boolean = !isSuspend && name.asString() == "invoke" &&
-        parentAsClass.origin == JvmLoweredDeclarationOrigin.LAMBDA_IMPL
+private fun IrFunction.isInvokeOfSuspendCallableReference(): Boolean =
+    isSuspend && name.asString() == "invoke" && parentAsClass.origin == JvmLoweredDeclarationOrigin.FUNCTION_REFERENCE_IMPL

 private fun IrFunction.isBridgeToSuspendImplMethod(): Boolean =
    isSuspend && this is IrSimpleFunction && parentAsClass.functions.any {
        it.name.asString() == name.asString() + SUSPEND_IMPL_NAME_SUFFIX && it.attributeOwnerId == attributeOwnerId
    }

-internal fun IrFunction.isKnownToBeTailCall(): Boolean =
-    when (origin) {
-        IrDeclarationOrigin.FUNCTION_FOR_DEFAULT_PARAMETER,
-        JvmLoweredDeclarationOrigin.JVM_OVERLOADS_WRAPPER,
-        JvmLoweredDeclarationOrigin.MULTIFILE_BRIDGE,
-        JvmLoweredDeclarationOrigin.SYNTHETIC_ACCESSOR,
-        JvmLoweredDeclarationOrigin.DEFAULT_IMPLS_BRIDGE,
-        JvmLoweredDeclarationOrigin.DEFAULT_IMPLS_BRIDGE_TO_SYNTHETIC,
-        JvmLoweredDeclarationOrigin.GENERATED_MEMBER_IN_CALLABLE_REFERENCE,
-        IrDeclarationOrigin.BRIDGE,
-        IrDeclarationOrigin.BRIDGE_SPECIAL,
-        IrDeclarationOrigin.DELEGATED_MEMBER -> true
-        else -> isInvokeOfSuspendMainWrapper() || isInvokeOfSuspendCallableReference() || isBridgeToSuspendImplMethod()
-    }
+internal fun IrFunction.shouldContainSuspendMarkers(): Boolean = !isInvokeSuspendOfContinuation() &&
+        // These are tail-call bridges and do not require any bytecode modifications.
+        origin != IrDeclarationOrigin.FUNCTION_FOR_DEFAULT_PARAMETER &&
+        origin != JvmLoweredDeclarationOrigin.JVM_OVERLOADS_WRAPPER &&
+        origin != JvmLoweredDeclarationOrigin.MULTIFILE_BRIDGE &&
+        origin != JvmLoweredDeclarationOrigin.SYNTHETIC_ACCESSOR &&
+        origin != JvmLoweredDeclarationOrigin.DEFAULT_IMPLS_BRIDGE &&
+        origin != JvmLoweredDeclarationOrigin.DEFAULT_IMPLS_BRIDGE_TO_SYNTHETIC &&
+        origin != JvmLoweredDeclarationOrigin.GENERATED_MEMBER_IN_CALLABLE_REFERENCE &&
+        origin != IrDeclarationOrigin.BRIDGE &&
+        origin != IrDeclarationOrigin.BRIDGE_SPECIAL &&
+        origin != IrDeclarationOrigin.DELEGATED_MEMBER &&
+        !isInvokeOfSuspendCallableReference() &&
+        !isBridgeToSuspendImplMethod()

-internal fun IrFunction.shouldNotContainSuspendMarkers(): Boolean =
-    isInvokeSuspendOfContinuation() || isKnownToBeTailCall()
+internal fun IrFunction.hasContinuation(): Boolean = isSuspend && shouldContainSuspendMarkers() &&
+        // These are templates for the inliner; the continuation will be generated after it runs.
+        origin != IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA &&
+        origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE &&
+        origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE

 internal fun IrExpression?.isReadOfCrossinline(): Boolean = when (this) {
    is IrGetValue -> (symbol.owner as? IrValueParameter)?.isCrossinline == true
@@ -156,31 +155,3 @@ internal fun createFakeContinuation(context: JvmBackendContext): IrExpression =
    context.ir.symbols.continuationClass.createType(true, listOf(makeTypeProjection(context.irBuiltIns.anyNType, Variance.INVARIANT))),
    "FAKE_CONTINUATION"
 )
-
-internal fun generateFakeContinuationConstructorCall(
-    v: InstructionAdapter,
-    classCodegen: ClassCodegen,
-    continuationClass: IrClass,
-    irFunction: IrFunction
-) {
-    val continuationType = classCodegen.typeMapper.mapClass(continuationClass)
-    // TODO: This is different in case of DefaultImpls
-    val thisNameType = Type.getObjectType(classCodegen.visitor.thisName.replace(".", "/"))
-    val continuationIndex = listOfNotNull(irFunction.dispatchReceiverParameter, irFunction.extensionReceiverParameter).size +
-            irFunction.valueParameters.size - 1
-    with(v) {
-        addFakeContinuationConstructorCallMarker(this, true)
-        anew(continuationType)
-        dup()
-        if (irFunction.dispatchReceiverParameter != null) {
-            load(0, AsmTypes.OBJECT_TYPE)
-            load(continuationIndex, Type.getObjectType("kotlin/coroutines/Continuation"))
-            invokespecial(continuationType.internalName, "<init>", "(${thisNameType}Lkotlin/coroutines/Continuation;)V", false)
-        } else {
-            load(continuationIndex, Type.getObjectType("kotlin/coroutines/Continuation"))
-            invokespecial(continuationType.internalName, "<init>", "(Lkotlin/coroutines/Continuation;)V", false)
-        }
-        addFakeContinuationConstructorCallMarker(this, false)
-        pop()
-    }
-}
@@ -226,7 +226,23 @@ 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, continuationClass, irFunction)
+        val continuationType = typeMapper.mapClass(continuationClass)
+        val continuationIndex = frameMap.getIndex(irFunction.continuationParameter()!!.symbol)
+        with(mv) {
+            addFakeContinuationConstructorCallMarker(this, true)
+            anew(continuationType)
+            dup()
+            if (irFunction.dispatchReceiverParameter != null) {
+                load(0, OBJECT_TYPE)
+                load(continuationIndex, Type.getObjectType("kotlin/coroutines/Continuation"))
+                invokespecial(continuationType.internalName, "<init>", "(${classCodegen.type}Lkotlin/coroutines/Continuation;)V", false)
+            } else {
+                load(continuationIndex, Type.getObjectType("kotlin/coroutines/Continuation"))
+                invokespecial(continuationType.internalName, "<init>", "(Lkotlin/coroutines/Continuation;)V", false)
+            }
+            addFakeContinuationConstructorCallMarker(this, false)
+            pop()
+        }
    }

    private fun generateNonNullAssertions() {
@@ -444,7 +460,7 @@ class ExpressionCodegen(
            }
            expression is IrConstructorCall ->
                MaterialValue(this, asmType, expression.type)
-            !irFunction.shouldNotContainSuspendMarkers() && expression.type.isUnit() -> {
+            expression.type.isUnit() && irFunction.shouldContainSuspendMarkers() -> {
                // NewInference allows casting `() -> T` to `() -> Unit`. A CHECKCAST here will fail.
                // Also, if the callee is a suspend function with a suspending tail call, the next `resumeWith`
                // will continue from here, but the value passed to it might not have been `Unit`. An exception
@@ -469,7 +485,7 @@ class ExpressionCodegen(
    }

    private fun IrFunctionAccessExpression.isSuspensionPoint(): Boolean = when {
-        !symbol.owner.isSuspend || irFunction.shouldNotContainSuspendMarkers() -> false
+        !symbol.owner.isSuspend || !irFunction.shouldContainSuspendMarkers() -> false
        // Copy-pasted bytecode blocks are not suspension points.
        symbol.owner.isInline ->
            symbol.owner.fqNameForIrSerialization == FqName("kotlin.coroutines.intrinsics.IntrinsicsKt.suspendCoroutineUninterceptedOrReturn")
@@ -143,16 +143,6 @@ open class FunctionCodegen(
            context.suspendLambdaToOriginalFunctionMap[irFunction.parentAsClass.attributeOwnerId]!!.symbol.descriptor.psiElement)
                as KtElement

-    private fun IrFunction.hasContinuation(): Boolean = isSuspend &&
-            // We do not generate continuation and state-machine for synthetic accessors, bridges, and delegated members,
-            // in a sense, they are tail-call
-            !isKnownToBeTailCall() &&
-            // This is suspend lambda parameter of inline function
-            origin != IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA &&
-            // This is just a template for inliner
-            origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE &&
-            origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
-
    private fun IrFunction.continuationClass(): IrClass =
            (body as IrBlockBody).statements.first { it is IrClass && it.origin == JvmLoweredDeclarationOrigin.CONTINUATION_CLASS } as IrClass

@@ -12,9 +12,9 @@ import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
-import org.jetbrains.kotlin.backend.jvm.codegen.isInvokeSuspendForInlineOfLambda
+import org.jetbrains.kotlin.backend.jvm.codegen.continuationParameter
+import org.jetbrains.kotlin.backend.jvm.codegen.hasContinuation
 import org.jetbrains.kotlin.backend.jvm.codegen.isInvokeSuspendOfContinuation
-import org.jetbrains.kotlin.backend.jvm.codegen.isInvokeSuspendOfLambda
 import org.jetbrains.kotlin.backend.jvm.codegen.isReadOfCrossinline
 import org.jetbrains.kotlin.backend.jvm.ir.IrInlineReferenceLocator
 import org.jetbrains.kotlin.backend.jvm.ir.defaultValue
@@ -33,7 +33,6 @@ import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrGetFieldImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
-import org.jetbrains.kotlin.ir.expressions.impl.IrTypeOperatorCallImpl
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
@@ -616,11 +615,11 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                functionsStack.push(declaration)
                val function = super.visitFunction(declaration) as IrFunction
                functionsStack.pop()
-                if (skip(function)) return function
+                if (!function.isSuspend || function.origin == IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA) return function

                val view = function.getOrCreateSuspendFunctionViewIfNeeded(context) as IrSimpleFunction

-                if (withoutContinuationClass(function)) return view
+                if (function.body == null || !function.hasContinuation()) return view

                if (function in suspendFunctionsCapturingCrossinline || function.isInline) {
                    val newFunction = buildFunWithDescriptorForInlining(view.descriptor) {
@@ -661,19 +660,6 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                return newFunction
            }

-            // 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.origin == JvmLoweredDeclarationOrigin.GENERATED_MEMBER_IN_CALLABLE_REFERENCE ||
-                        function.origin == JvmLoweredDeclarationOrigin.DEFAULT_IMPLS_BRIDGE ||
-                        function.origin == JvmLoweredDeclarationOrigin.JVM_OVERLOADS_WRAPPER
-
-            private fun skip(function: IrFunction) =
-                !function.isSuspend ||
-                        function.origin == IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA
-
            private fun registerNewFunction(function: IrFunction) {
                functionsToAdd.peek()!!.add(function)
            }
@@ -731,19 +717,22 @@ private fun IrFunction.suspendFunctionStub(context: JvmBackendContext): IrFuncti

        function.copyAttributes(this)
        function.copyTypeParametersFrom(this)
+        function.copyReceiverParametersFrom(this)

        if (origin != JvmLoweredDeclarationOrigin.DEFAULT_IMPLS_BRIDGE) {
            function.overriddenSymbols +=
                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.
+        // 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))
-        }
+        val index = valueParameters.firstOrNull { it.origin == IrDeclarationOrigin.MASK_FOR_DEFAULT_FUNCTION }?.index
+            ?: valueParameters.size
+        function.valueParameters += valueParameters.take(index).map { it.copyTo(function, index = it.index) }
+        function.addValueParameter(
+            SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME, continuationType(context), JvmLoweredDeclarationOrigin.CONTINUATION_CLASS
+        )
+        function.valueParameters += valueParameters.drop(index).map { it.copyTo(function, index = it.index + 1) }

        context.recordSuspendFunctionViewStub(this, function)
    }
@@ -792,44 +781,20 @@ private fun IrCall.createSuspendFunctionCallViewIfNeeded(context: JvmBackendCont

    // While the new callee technically returns `<original type> | COROUTINE_SUSPENDED`, the latter case is handled
    // by a method visitor so at an IR overview we don't need to consider it.
-    return IrCallImpl(startOffset, endOffset, type, view.symbol, superQualifierSymbol = superQualifierSymbol).also {
+    return IrCallImpl(startOffset, endOffset, type, view.symbol, origin, superQualifierSymbol).also {
        it.copyTypeArgumentsFrom(this)
        it.dispatchReceiver = dispatchReceiver
        it.extensionReceiver = extensionReceiver
-        val valueArguments = computeSuspendFunctionCallViewValueArguments(caller, context, view)
-        valueArguments.forEachIndexed { index, argument -> it.putValueArgument(index, argument) }
+        val continuationIndex = view.continuationParameter()!!.index
+        for (i in 0 until valueArgumentsCount) {
+            it.putValueArgument(i + if (i >= continuationIndex) 1 else 0, getValueArgument(i)!!)
+        }
+        // At this point the only LOCAL_FUNCTION_FOR_LAMBDAs are inline and crossinline lambdas
+        val continuation = if (caller.origin == IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA)
+            context.fakeContinuation
+        else
+            IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, caller.continuationParameter()?.symbol
+                ?: throw AssertionError("${caller.render()} has no continuation; can't call ${symbol.owner.render()}"))
+        it.putValueArgument(continuationIndex, continuation)
    }
 }
-
-private fun IrCall.computeSuspendFunctionCallViewValueArguments(
-    caller: IrFunction,
-    context: JvmBackendContext,
-    view: IrFunction
-): List<IrExpression> {
-    // 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.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)
-    val valueArguments = mutableListOf<IrExpression>()
-    for (i in 0 until continuationIndex) {
-        valueArguments.add(getValueArgument(i)!!)
-    }
-    valueArguments.add(continuationParameter)
-    if (numberOfMasks != 0) {
-        for (i in 0 until numberOfMasks + 1) {
-            valueArguments.add(getValueArgument(continuationIndex + i)!!)
-        }
-    }
-    return valueArguments
-}
@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
-import org.jetbrains.kotlin.backend.jvm.ir.createJvmIrBuilder
 import org.jetbrains.kotlin.backend.jvm.ir.getJvmNameFromAnnotation
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.languageVersionSettings
@@ -21,17 +20,11 @@ import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.builders.declarations.addFunction
 import org.jetbrains.kotlin.ir.builders.declarations.addValueParameter
-import org.jetbrains.kotlin.ir.builders.declarations.buildFun
 import org.jetbrains.kotlin.ir.declarations.IrClass
-import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
 import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.declarations.IrValueParameter
-import org.jetbrains.kotlin.ir.declarations.impl.IrFunctionImpl
-import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
 import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionReferenceImpl
-import org.jetbrains.kotlin.ir.expressions.impl.IrReturnImpl
 import org.jetbrains.kotlin.ir.types.*
-import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.util.functions
 import org.jetbrains.kotlin.ir.util.isFileClass
 import org.jetbrains.kotlin.name.Name
@@ -81,27 +74,19 @@ private class MainMethodGenerationLowering(private val context: JvmBackendContex

        irClass.functions.find { it.isMainMethod() }?.let { mainMethod ->
            if (mainMethod.isSuspend) {
-                generateMainMethod(irClass) { args ->
-                    body = buildRunSuspendWrapper {
-                        +IrReturnImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, context.irBuiltIns.nothingType, it, irCall(mainMethod).apply {
-                            putValueArgument(0, irGet(args))
-                        })
-                    }
+                irClass.generateMainMethod { args ->
+                    +irRunSuspend(mainMethod, args)
                }
            }
            return
        }

        irClass.functions.find { it.isParameterlessMainMethod() }?.let { parameterlessMainMethod ->
-            generateMainMethod(irClass) {
-                body = if (parameterlessMainMethod.isSuspend) {
-                    buildRunSuspendWrapper {
-                        +IrReturnImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, context.irBuiltIns.nothingType, it, irCall(parameterlessMainMethod))
-                    }
+            irClass.generateMainMethod {
+                if (parameterlessMainMethod.isSuspend) {
+                    +irRunSuspend(parameterlessMainMethod, null)
                } else {
-                    context.createIrBuilder(symbol).irBlockBody {
-                        +irCall(parameterlessMainMethod)
-                    }
+                    +irCall(parameterlessMainMethod)
                }
            }
        }
@@ -130,11 +115,8 @@ private class MainMethodGenerationLowering(private val context: JvmBackendContex
        }
    }

-    private fun generateMainMethod(
-        irClass: IrClass,
-        makeBody: IrFunctionImpl.(IrValueParameter) -> Unit
-    ) {
-        irClass.addFunction {
+    private fun IrClass.generateMainMethod(makeBody: IrBlockBodyBuilder.(IrValueParameter) -> Unit) =
+        addFunction {
            name = Name.identifier("main")
            visibility = Visibilities.PUBLIC
            returnType = context.irBuiltIns.unitType
@@ -145,41 +127,22 @@ private class MainMethodGenerationLowering(private val context: JvmBackendContex
                name = Name.identifier("args")
                type = context.irBuiltIns.arrayClass.typeWith(context.irBuiltIns.stringType)
            }
-            makeBody(args)
+            body = context.createIrBuilder(symbol).irBlockBody { makeBody(args) }
        }
-    }

-    private fun IrFunctionImpl.buildRunSuspendWrapper(buildMainInvocation: IrBlockBodyBuilder.(IrSimpleFunctionSymbol) -> Unit) =
-        context.createJvmIrBuilder(symbol).irBlockBody {
-            val lambda = irBlock(UNDEFINED_OFFSET, UNDEFINED_OFFSET, IrStatementOrigin.LAMBDA) {
-                val function = buildFun {
-                    origin = IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA
-                    name = Name.identifier("wrapper_for_suspend_main")
-                    visibility = Visibilities.LOCAL
-                    returnType = context.irBuiltIns.anyNType
-                }.apply {
-                    addValueParameter("k", context.irBuiltIns.anyNType)
-                    parent = this
-                    body = irBlockBody {
-                        buildMainInvocation(symbol)
-                    }
-                }
-                +function
-                +IrFunctionReferenceImpl(
-                    UNDEFINED_OFFSET,
-                    UNDEFINED_OFFSET,
-                    context.irBuiltIns.function1.typeWith(
-                        context.irBuiltIns.anyNType,
-                        context.irBuiltIns.anyNType
-                    ),
-                    function.symbol,
-                    typeArgumentsCount = 0,
-                    origin = IrStatementOrigin.LAMBDA,
-                    reflectionTarget = null
-                )
-            }
-            +irCall(this@MainMethodGenerationLowering.context.ir.symbols.runSuspendFunction).apply {
-                putValueArgument(0, lambda)
+    private fun IrBuilderWithScope.irRunSuspend(target: IrSimpleFunction, args: IrValueParameter?) =
+        irCall(this@MainMethodGenerationLowering.context.ir.symbols.runSuspendFunction).apply {
+            val reference = IrFunctionReferenceImpl(
+                UNDEFINED_OFFSET,
+                UNDEFINED_OFFSET,
+                this@MainMethodGenerationLowering.context.ir.symbols.suspendFunctionN(0).typeWith(context.irBuiltIns.anyNType),
+                target.symbol,
+                typeArgumentsCount = 0,
+                reflectionTarget = null
+            )
+            if (args != null) {
+                reference.putValueArgument(0, irGet(args))
            }
+            putValueArgument(0, reference)
        }
 }
@@ -9,10 +9,8 @@ import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.ir.isSuspend
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
-import org.jetbrains.kotlin.backend.jvm.codegen.anyOfOverriddenFunctionsReturnsNonUnit
-import org.jetbrains.kotlin.backend.jvm.codegen.isKnownToBeTailCall
+import org.jetbrains.kotlin.backend.jvm.codegen.hasContinuation
 import org.jetbrains.kotlin.ir.IrStatement
-import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
 import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.expressions.*
@@ -20,7 +18,6 @@ import org.jetbrains.kotlin.ir.expressions.impl.IrReturnImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrTypeOperatorCallImpl
 import org.jetbrains.kotlin.ir.types.isUnit
 import org.jetbrains.kotlin.ir.util.fqNameWhenAvailable
-import org.jetbrains.kotlin.ir.util.isSuspend
 import org.jetbrains.kotlin.ir.util.parentAsClass
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformer
 import org.jetbrains.kotlin.name.FqName
@@ -77,10 +74,7 @@ private class TailCallOptimizationData(val function: IrFunction) {
    }

    init {
-        if (function.isSuspend && function.origin != IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA && !function.isKnownToBeTailCall() &&
-            // See `disableTailCallOptimizationForFunctionReturningUnit` in `generateStateMachineForNamedFunction`:
-            !(returnsUnit && function.anyOfOverriddenFunctionsReturnsNonUnit())
-        ) {
+        if (function.hasContinuation()) {
            when (val body = function.body) {
                is IrBlockBody -> body.statements.findTailCall(returnsUnit)?.let(::findCallsOnTailPositionWithoutImmediateReturn)
                is IrExpressionBody -> findCallsOnTailPositionWithoutImmediateReturn(body.expression)
@@ -318,6 +318,4 @@ class IrBuiltIns(
        const val ANDAND = "ANDAND"
        const val OROR = "OROR"
    }
-
-    val function1: IrClassSymbol = builtIns.getFunction(1).toIrSymbol()
 }
@@ -0,0 +1,22 @@
+// WITH_RUNTIME
+// WITH_COROUTINES
+// TARGET_BACKEND: JVM
+// FILE: a.kt
+
+inline suspend fun f(crossinline lambda: suspend (Double) -> Double): Double {
+    val obj = object {
+        suspend fun g(x: Double): Double = lambda(x)
+    }
+    return obj.g(1.0)
+}
+
+// FILE: b.kt
+import helpers.*
+import kotlin.coroutines.*
+import kotlin.coroutines.intrinsics.*
+
+fun box(): String {
+    var result: Double = 0.0
+    suspend { result = f { it } }.startCoroutine(EmptyContinuation)
+    return if (result == 1.0) "OK" else "fail: $result"
+}
@@ -3813,6 +3813,11 @@ public class BlackBoxInlineCodegenTestGenerated extends AbstractBlackBoxInlineCo
            runTestWithPackageReplacement("compiler/testData/codegen/boxInline/suspend/multipleSuspensionPoints.kt", "kotlin.coroutines");
        }

+        @TestMetadata("nestedMethodWith2XParameter.kt")
+        public void testNestedMethodWith2XParameter() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/suspend/nestedMethodWith2XParameter.kt");
+        }
+
        @TestMetadata("nonSuspendCrossinline.kt")
        public void testNonSuspendCrossinline_1_2() throws Exception {
            runTestWithPackageReplacement("compiler/testData/codegen/boxInline/suspend/nonSuspendCrossinline.kt", "kotlin.coroutines.experimental");
@@ -3813,6 +3813,11 @@ public class CompileKotlinAgainstInlineKotlinTestGenerated extends AbstractCompi
            runTestWithPackageReplacement("compiler/testData/codegen/boxInline/suspend/multipleSuspensionPoints.kt", "kotlin.coroutines");
        }

+        @TestMetadata("nestedMethodWith2XParameter.kt")
+        public void testNestedMethodWith2XParameter() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/suspend/nestedMethodWith2XParameter.kt");
+        }
+
        @TestMetadata("nonSuspendCrossinline.kt")
        public void testNonSuspendCrossinline_1_2() throws Exception {
            runTestWithPackageReplacement("compiler/testData/codegen/boxInline/suspend/nonSuspendCrossinline.kt", "kotlin.coroutines.experimental");
@@ -3728,6 +3728,11 @@ public class IrBlackBoxInlineCodegenTestGenerated extends AbstractIrBlackBoxInli
            runTestWithPackageReplacement("compiler/testData/codegen/boxInline/suspend/multipleSuspensionPoints.kt", "kotlin.coroutines");
        }

+        @TestMetadata("nestedMethodWith2XParameter.kt")
+        public void testNestedMethodWith2XParameter() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/suspend/nestedMethodWith2XParameter.kt");
+        }
+
        @TestMetadata("nonSuspendCrossinline.kt")
        public void testNonSuspendCrossinline_1_3() throws Exception {
            runTestWithPackageReplacement("compiler/testData/codegen/boxInline/suspend/nonSuspendCrossinline.kt", "kotlin.coroutines");
@@ -3728,6 +3728,11 @@ public class IrCompileKotlinAgainstInlineKotlinTestGenerated extends AbstractIrC
            runTestWithPackageReplacement("compiler/testData/codegen/boxInline/suspend/multipleSuspensionPoints.kt", "kotlin.coroutines");
        }

+        @TestMetadata("nestedMethodWith2XParameter.kt")
+        public void testNestedMethodWith2XParameter() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/suspend/nestedMethodWith2XParameter.kt");
+        }
+
        @TestMetadata("nonSuspendCrossinline.kt")
        public void testNonSuspendCrossinline_1_3() throws Exception {
            runTestWithPackageReplacement("compiler/testData/codegen/boxInline/suspend/nonSuspendCrossinline.kt", "kotlin.coroutines");