[JS IR] Make coroutines work with new lambda scheme

- Do not create extra class for suspend lambda

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/ir/IrUtils.kt

@@ -526,7 +526,7 @@ fun IrClass.addFakeOverrides() {
             ).apply {
                 descriptor.bind(this)
                 parent = this@addFakeOverrides
-                overriddenSymbols += overriddenFunctions.map { it.symbol }
+                overriddenSymbols = overriddenFunctions.map { it.symbol }
                 copyParameterDeclarationsFrom(irFunction)
                 copyAttributes(irFunction)
             }
@@ -537,7 +537,9 @@ fun IrClass.addFakeOverrides() {
         .associate { it.value.first() to createFakeOverride(it.value) }
         .toMutableMap()
 
-    declarations += fakeOverriddenFunctions.values
+    for (fo in fakeOverriddenFunctions.values) {
+        addChild(fo)
+    }
 }
 
 fun createStaticFunctionWithReceivers(

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/JsLoweringPhases.kt

@@ -16,7 +16,6 @@ import org.jetbrains.kotlin.ir.backend.js.lower.*
 import org.jetbrains.kotlin.ir.backend.js.lower.calls.CallsLowering
 import org.jetbrains.kotlin.ir.backend.js.lower.cleanup.CleanupLowering
 import org.jetbrains.kotlin.ir.backend.js.lower.coroutines.JsSuspendFunctionsLowering
-import org.jetbrains.kotlin.ir.backend.js.lower.coroutines.RemoveSuspendLambdas
 import org.jetbrains.kotlin.ir.backend.js.lower.inline.CopyInlineFunctionBodyLowering
 import org.jetbrains.kotlin.ir.backend.js.lower.inline.RemoveInlineFunctionsWithReifiedTypeParametersLowering
 import org.jetbrains.kotlin.ir.declarations.IrFile
@@ -360,12 +359,6 @@ private val suspendFunctionsLoweringPhase = makeBodyLoweringPhase(
     description = "Transform suspend functions into CoroutineImpl instance and build state machine"
 )
 
-private val suspendLambdasRemovalLoweringPhase = makeDeclarationTransformerPhase(
-    { RemoveSuspendLambdas() },
-    name = "RemoveSuspendLambdas",
-    description = "Remove suspend lambdas"
-)
-
 private val privateMembersLoweringPhase = makeDeclarationTransformerPhase(
     ::PrivateMembersLowering,
     name = "PrivateMembersLowering",
@@ -632,7 +625,6 @@ val loweringList = listOf<Lowering>(
     enumUsageLoweringPhase,
     enumEntryRemovalLoweringPhase,
     suspendFunctionsLoweringPhase,
-    suspendLambdasRemovalLoweringPhase,
     interopCallableReferenceLoweringPhase,
     returnableBlockLoweringPhase,
     forLoopsLoweringPhase,

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/CallableReferenceLowering.kt

@@ -177,6 +177,7 @@ class CallableReferenceLowering(private val context: CommonBackendContext) : Bod
                 name = superMethod.name
                 returnType = function.returnType
                 isSuspend = superMethod.isSuspend
+                isOperator = superMethod.isOperator
             }.apply {
                 overriddenSymbols = listOf(superMethod.symbol)
                 dispatchReceiverParameter = clazz.createDispatchReceiver().also { it.parent = this }

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/coroutines/AbstractSuspendFunctionsLowering.kt

@@ -7,12 +7,13 @@ package org.jetbrains.kotlin.ir.backend.js.lower.coroutines
 
 import org.jetbrains.kotlin.backend.common.*
 import org.jetbrains.kotlin.backend.common.ir.*
+import org.jetbrains.kotlin.backend.common.lower.LocalDeclarationsLowering
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
-import org.jetbrains.kotlin.backend.common.lower.irBlockBody
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.backend.js.lower.CallableReferenceLowering
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.declarations.impl.IrClassImpl
@@ -67,114 +68,27 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
     private val builtCoroutines = mutableMapOf<IrFunction, BuiltCoroutine>()
 
     override fun lower(irBody: IrBody, container: IrDeclaration) {
-        transformCallableReferencesToSuspendLambdas(irBody)
+        transformToCoroutine(irBody, container)
-
-        if (container.isLambda) return
-
-        tryTransformSuspendFunction(container, null)
     }
 
-    // Suppress since it is used in native
-    @Suppress("MemberVisibilityCanBePrivate")
-    protected fun IrCall.isReturnIfSuspendedCall() =
-        symbol.owner.run { fqNameWhenAvailable == context.internalPackageFqn.child(Name.identifier("returnIfSuspended")) }
 
-    private fun tryTransformSuspendFunction(element: IrElement, functionReference: IrFunctionReference?) {
+    private fun transformToCoroutine(irBody: IrBody, container: IrDeclaration) {
-        if (element is IrSimpleFunction && element.isSuspend && element.modality != Modality.ABSTRACT) {
+        if (container !is IrSimpleFunction || !container.isSuspend) return
 
-            transformCallableReferencesToSuspendLambdas(element)
+        transformSuspendFunction2(container, irBody)?.let {
-
+            val dc = container.parent as IrDeclarationContainer
-            transformSuspendFunction(element, functionReference)?.let { result ->
+            dc.addChild(it)
-                result.forEach { declaration ->
-                    if (declaration !== element) {
-                        // TODO Use proper means to emerge declarations
-                        element.file.declarations += declaration
-                        declaration.parent = element.file
-
-                        // TODO investigate IrJsCodegenBoxTestGenerated$Coroutines$ControlFlow.testBreakFinally_1_3
-                        declaration.patchDeclarationParents(declaration.parent)
-                    }
-                }
-            }
         }
     }
 
-    private fun transformCallableReferencesToSuspendLambdas(irElement: IrElement) {
+    private fun getSuspendFunctionKind2(function: IrSimpleFunction, body: IrBody): SuspendFunctionKind {
-        irElement.transformChildrenVoid(object : IrElementTransformerVoid() {
 
-            override fun visitFunctionReference(expression: IrFunctionReference): IrExpression {
+        fun IrSimpleFunction.isSuspendLambda() =
-                expression.transformChildrenVoid(this)
+            name.asString() == "invoke" && parentClassOrNull?.let { it.origin === CallableReferenceLowering.Companion.LAMBDA_IMPL } == true
 
-                if (!expression.isSuspend)
+        if (function.isSuspendLambda())
-                    return expression
-
-                val coroutineConstructor = if (expression.symbol.owner.isLambda) {
-                    if (expression.symbol.owner in builtCoroutines) {
-                        error("Lambda revisiting?")
-                    }
-
-                    tryTransformSuspendFunction(expression.symbol.owner, expression)
-
-                    builtCoroutines[expression.symbol.owner]?.coroutineConstructor
-                        ?: throw Error("Non-local callable reference to suspend lambda: $expression")
-                } else {
-                    expression.symbol.owner.coroutineConstructor ?: run {
-
-                        tryTransformSuspendFunction(expression.symbol.owner, expression)
-
-                        builtCoroutines[expression.symbol.owner]!!.coroutineConstructor
-                    }
-                }
-
-                val constructorParameters = coroutineConstructor.valueParameters
-                val expressionArguments = expression.getArguments().map { it.second }
-                assert(constructorParameters.size == expressionArguments.size) {
-                    "Inconsistency between callable reference to suspend lambda and the corresponding coroutine"
-                }
-                val irBuilder = context.createIrBuilder(expression.symbol, expression.startOffset, expression.endOffset)
-                irBuilder.run {
-                    return irCall(coroutineConstructor.symbol).apply {
-                        expressionArguments.forEachIndexed { index, argument ->
-                            putValueArgument(index, argument)
-                        }
-                    }
-                }
-            }
-        })
-    }
-
-    private sealed class SuspendFunctionKind {
-        object NO_SUSPEND_CALLS : SuspendFunctionKind()
-        class DELEGATING(val delegatingCall: IrCall) : SuspendFunctionKind()
-        object NEEDS_STATE_MACHINE : SuspendFunctionKind()
-    }
-
-    private fun transformSuspendFunction(irFunction: IrSimpleFunction, functionReference: IrFunctionReference?) =
-        when (val suspendFunctionKind = getSuspendFunctionKind(irFunction, functionReference)) {
-            is SuspendFunctionKind.NO_SUSPEND_CALLS -> {
-                null                                                            // No suspend function calls - just an ordinary function.
-            }
-
-            is SuspendFunctionKind.DELEGATING -> {                              // Calls another suspend function at the end.
-                removeReturnIfSuspendedCallAndSimplifyDelegatingCall(irFunction, suspendFunctionKind.delegatingCall)
-                null                                                            // No need in state machine.
-            }
-
-            is SuspendFunctionKind.NEEDS_STATE_MACHINE -> {
-                val coroutine = buildCoroutine(irFunction, functionReference)   // Coroutine implementation.
-                if (irFunction.isLambda)             // Suspend lambdas are called through factory method <create>,
-                    listOf(coroutine)                                           // thus we can eliminate original body.
-                else
-                    listOf<IrDeclaration>(coroutine, irFunction)
-            }
-        }
-
-    private fun getSuspendFunctionKind(irFunction: IrSimpleFunction, functionReference: IrFunctionReference?): SuspendFunctionKind {
-        if (irFunction.isLambda || functionReference != null)
             return SuspendFunctionKind.NEEDS_STATE_MACHINE            // Suspend lambdas always need coroutine implementation.
 
-        val body = irFunction.body ?: return SuspendFunctionKind.NO_SUSPEND_CALLS
-
         var numberOfSuspendCalls = 0
         body.acceptVoid(object : IrElementVisitorVoid {
             override fun visitElement(element: IrElement) {
@@ -223,81 +137,106 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
         }
     }
 
-    private val symbols = context.ir.symbols
+    private fun transformSuspendFunction2(function: IrSimpleFunction, body: IrBody): IrClass? {
-    private val getContinuationSymbol = symbols.getContinuation
+        assert(function.isSuspend)
-    private val continuationClassSymbol = getContinuationSymbol.owner.returnType.classifierOrFail as IrClassSymbol
 
-    private fun removeReturnIfSuspendedCallAndSimplifyDelegatingCall(irFunction: IrFunction, delegatingCall: IrCall) {
+        return when (val functionKind = getSuspendFunctionKind2(function, body)) {
-        val returnValue =
+            is SuspendFunctionKind.NO_SUSPEND_CALLS -> {
-            if (delegatingCall.isReturnIfSuspendedCall())
+                null                                                            // No suspend function calls - just an ordinary function.
-                delegatingCall.getValueArgument(0)!!
+            }
-            else delegatingCall
+
-        context.createIrBuilder(irFunction.symbol).run {
+            is SuspendFunctionKind.DELEGATING -> {                              // Calls another suspend function at the end.
-            val statements = (irFunction.body as IrBlockBody).statements
+                removeReturnIfSuspendedCallAndSimplifyDelegatingCall(function, functionKind.delegatingCall)
-            val lastStatement = statements.last()
+                null                                                            // No need in state machine.
-            assert(lastStatement == delegatingCall || lastStatement is IrReturn) { "Unexpected statement $lastStatement" }
+            }
-            statements[statements.lastIndex] = irReturn(generateDelegatedCall(irFunction.returnType, returnValue))
+
+            is SuspendFunctionKind.NEEDS_STATE_MACHINE -> {
+                val coroutine = buildCoroutine2(function)      // Coroutine implementation.
+                if (coroutine === function.parent)             // Suspend lambdas are called through factory method <create>,
+                    null
+                else
+                    coroutine
+            }
         }
     }
 
-    private fun buildCoroutine(irFunction: IrSimpleFunction, functionReference: IrFunctionReference?): IrClass {
+    private fun IrBlockBodyBuilder.createCoroutineInstance(function: IrSimpleFunction, parameters: Collection<IrValueParameter>, coroutine: BuiltCoroutine): IrExpression {
-        val coroutine = CoroutineBuilder(irFunction, functionReference).build()
+        val constructor = coroutine.coroutineConstructor
-        builtCoroutines[irFunction] = coroutine
+        val coroutineTypeArgs = function.typeParameters.map {
-        irFunction.coroutineConstructor = coroutine.coroutineConstructor
+            IrSimpleTypeImpl(it.symbol, true, emptyList(), emptyList())
+        }
 
-        if (!irFunction.isLambda && functionReference == null) {
+        return irCallConstructor(constructor.symbol, coroutineTypeArgs).apply {
-            // It is not a lambda - replace original function with a call to constructor of the built coroutine.
+            parameters.forEachIndexed { index, argument ->
-            val irBuilder = context.createIrBuilder(irFunction.symbol, irFunction.startOffset, irFunction.endOffset)
+                putValueArgument(index, irGet(argument))
-            (irFunction.body as IrBlockBody).statements.let {
-                it.clear()
-                it += irBuilder.irBlockBody(irFunction) {
-                    val constructor = coroutine.coroutineConstructor
-                    generateCoroutineStart(coroutine.stateMachineFunction,
-                                           irCallConstructor(constructor.symbol, irFunction.typeParameters.map {
-                                               IrSimpleTypeImpl(it.symbol, true, emptyList(), emptyList())
-                                           }).apply {
-                                               val functionParameters = irFunction.explicitParameters
-                                               functionParameters.forEachIndexed { index, argument ->
-                                                   putValueArgument(index, irGet(argument))
-                                               }
-                                               putValueArgument(
-                                                   functionParameters.size,
-                                                   irCall(
-                                                       getContinuationSymbol,
-                                                       getContinuationSymbol.owner.returnType,
-                                                       listOf(irFunction.returnType)
-                                                   )
-                                               )
-                                           })
-                }.statements
             }
+            putValueArgument(
+                parameters.size,
+                irCall(
+                    getContinuationSymbol,
+                    getContinuationSymbol.owner.returnType,
+                    listOf(function.returnType)
+                )
+            )
+        }
+    }
+
+    private fun buildCoroutine2(function: IrSimpleFunction): IrClass {
+        val coroutine = CoroutineBuilder2(function).build()
+
+        val isSuspendLambda = coroutine.coroutineClass === function.parent
+
+        if (isSuspendLambda) return coroutine.coroutineClass
+
+        // It is not a lambda - replace original function with a call to constructor of the built coroutine.
+
+        with(function) {
+            val irBuilder = context.createIrBuilder(symbol, startOffset, endOffset)
+            val functionBody = body as IrBlockBody
+            functionBody.statements.clear()
+            functionBody.statements.addAll(irBuilder.irBlockBody {
+                generateCoroutineStart(coroutine.stateMachineFunction, createCoroutineInstance(this@with, explicitParameters, coroutine))
+            }.statements)
         }
 
         return coroutine.coroutineClass
     }
 
-    private class BuiltCoroutine(
+    private inner class CoroutineBuilder2(private val function: IrSimpleFunction) {
-        val coroutineClass: IrClass,
+        private val startOffset = function.startOffset
-        val coroutineConstructor: IrConstructor,
+        private val endOffset = function.endOffset
-        val stateMachineFunction: IrFunction
+        private val isSuspendLambda = function.isOperator && function.name.asString() == "invoke" && function.parentClassOrNull
-    )
+            ?.let { it.origin === CallableReferenceLowering.Companion.LAMBDA_IMPL } == true
+        private val functionParameters = if (isSuspendLambda) function.valueParameters else function.explicitParameters
 
-    private inner class CoroutineBuilder(val irFunction: IrFunction, val functionReference: IrFunctionReference?) {
+        private val coroutineClass: IrClass = getCoroutineClass(function)
 
-        private val startOffset = irFunction.startOffset
+        private val coroutineClassThis = coroutineClass.thisReceiver!!
-        private val endOffset = irFunction.endOffset
-        private val functionParameters = irFunction.explicitParameters
-        private val boundFunctionParameters = functionReference?.getArgumentsWithIr()?.map { it.first }
-        private val unboundFunctionParameters = boundFunctionParameters?.let { functionParameters - it }
 
-        private val coroutineClass: IrClass = WrappedClassDescriptor().let { d ->
+        private val continuationType = continuationClassSymbol.typeWith(function.returnType)
+
+        // Save all arguments to fields.
+        private val argumentToPropertiesMap = functionParameters.associateWith { coroutineClass.addField(it.name, it.type, false) }
+
+        private val coroutineBaseClass = getCoroutineBaseClass(function)
+        private val coroutineBaseClassConstructor = coroutineBaseClass.owner.constructors.single { it.valueParameters.size == 1 }
+        private val create1Function = coroutineBaseClass.owner.simpleFunctions()
+            .single { it.name.asString() == "create" && it.valueParameters.size == 1 }
+        private val create1CompletionParameter = create1Function.valueParameters[0]
+
+        private fun getCoroutineClass(function: IrSimpleFunction): IrClass {
+            return if (isSuspendLambda) function.parentAsClass
+            else buildNewCoroutineClass(function)
+        }
+
+        private fun buildNewCoroutineClass(function: IrSimpleFunction): IrClass = WrappedClassDescriptor().let { d ->
             IrClassImpl(
                 startOffset, endOffset,
                 DECLARATION_ORIGIN_COROUTINE_IMPL,
                 IrClassSymbolImpl(d),
-                nameForCoroutineClass(irFunction),
+                nameForCoroutineClass(function),
                 ClassKind.CLASS,
-                irFunction.visibility,
+                function.visibility,
                 Modality.FINAL,
                 isCompanion = false,
                 isInner = false,
@@ -308,315 +247,70 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
                 isFun = false
             ).apply {
                 d.bind(this)
-                parent = irFunction.parent
+                parent = function.parent
                 createParameterDeclarations()
-                typeParameters = irFunction.typeParameters.map { typeParam ->
+                typeParameters = function.typeParameters.map { typeParam ->
+                    // TODO: remap types
                     typeParam.copyToWithoutSuperTypes(this).apply { superTypes += typeParam.superTypes }
                 }
             }
         }
-        private val coroutineClassThis = coroutineClass.thisReceiver!!
 
-        private val continuationType = continuationClassSymbol.typeWith(irFunction.returnType)
+        private fun buildConstructor(): IrConstructor {
-
+            if (isSuspendLambda) {
-        // Save all arguments to fields.
+                return coroutineClass.declarations.filterIsInstance<IrConstructor>().single()
-        private val argumentToPropertiesMap = functionParameters.associateWith { coroutineClass.addField(it.name, it.type, false) }
-
-        private val coroutineBaseClass = getCoroutineBaseClass(irFunction)
-        private val coroutineBaseClassConstructor = coroutineBaseClass.owner.constructors.single { it.valueParameters.size == 1 }
-        private val create1Function = coroutineBaseClass.owner.simpleFunctions()
-            .single { it.name.asString() == "create" && it.valueParameters.size == 1 }
-        private val create1CompletionParameter = create1Function.valueParameters[0]
-
-        private val coroutineConstructors = mutableListOf<IrConstructor>()
-
-        fun build(): BuiltCoroutine {
-            val superTypes = mutableListOf(coroutineBaseClass.defaultType)
-            var suspendFunctionClass: IrClass? = null
-            var functionClass: IrClass? = null
-            val suspendFunctionClassTypeArguments: List<IrType>?
-            val functionClassTypeArguments: List<IrType>?
-            if (unboundFunctionParameters != null) {
-                // Suspend lambda inherits SuspendFunction.
-                val numberOfParameters = unboundFunctionParameters.size
-                suspendFunctionClass = symbols.suspendFunctionN(numberOfParameters).owner
-                val unboundParameterTypes = unboundFunctionParameters.map { it.type }
-                suspendFunctionClassTypeArguments = unboundParameterTypes + irFunction.returnType
-                superTypes += suspendFunctionClass.typeWith(suspendFunctionClassTypeArguments)
-
-                functionClass = symbols.functionN(numberOfParameters + 1).owner
-                functionClassTypeArguments = unboundParameterTypes + continuationType + context.irBuiltIns.anyNType
-                superTypes += functionClass.typeWith(functionClassTypeArguments)
             }
 
-            val coroutineConstructor = buildConstructor()
+            return WrappedClassConstructorDescriptor().let { d ->
+                IrConstructorImpl(
+                    startOffset, endOffset,
+                    DECLARATION_ORIGIN_COROUTINE_IMPL,
+                    IrConstructorSymbolImpl(d),
+                    coroutineBaseClassConstructor.name,
+                    function.visibility,
+                    coroutineClass.defaultType,
+                    isInline = false,
+                    isExternal = false,
+                    isPrimary = true,
+                    isExpect = false
+                ).apply {
+                    d.bind(this)
+                    parent = coroutineClass
+                    coroutineClass.addChild(this)
 
-            val superInvokeSuspendFunction = coroutineBaseClass.owner.simpleFunctions().single { it.name == stateMachineMethodName }
+                    valueParameters = functionParameters.mapIndexed { index, parameter ->
-            val invokeSuspendMethod = buildInvokeSuspendMethod(superInvokeSuspendFunction, coroutineClass)
+                        parameter.copyTo(this, DECLARATION_ORIGIN_COROUTINE_IMPL, index, defaultValue = null)
-
-            var coroutineFactoryConstructor: IrConstructor? = null
-            val createMethod: IrSimpleFunction?
-            if (functionReference != null) {
-                // Suspend lambda - create factory methods.
-                coroutineFactoryConstructor = buildFactoryConstructor(boundFunctionParameters!!)
-
-                val createFunctionSymbol = coroutineBaseClass.owner.simpleFunctions()
-                    .atMostOne { it.name.asString() == "create" && it.valueParameters.size == unboundFunctionParameters!!.size + 1 }
-                    ?.symbol
-
-                createMethod = buildCreateMethod(
-                    unboundArgs = unboundFunctionParameters!!,
-                    superFunctionSymbol = createFunctionSymbol,
-                    coroutineConstructor = coroutineConstructor
-                )
-
-                val invokeFunctionSymbol =
-                    functionClass!!.simpleFunctions().single { it.name.asString() == "invoke" }.symbol
-                val suspendInvokeFunctionSymbol =
-                    suspendFunctionClass!!.simpleFunctions().single { it.name.asString() == "invoke" }.symbol
-
-                buildInvokeMethod(
-                    suspendFunctionInvokeFunctionSymbol = suspendInvokeFunctionSymbol,
-                    functionInvokeFunctionSymbol = invokeFunctionSymbol,
-                    createFunction = createMethod,
-                    stateMachineFunction = invokeSuspendMethod
-                )
-            }
-
-            coroutineClass.superTypes += superTypes
-            coroutineClass.addFakeOverrides()
-
-            initializeStateMachine(coroutineConstructors, coroutineClassThis)
-
-            return BuiltCoroutine(
-                coroutineClass = coroutineClass,
-                coroutineConstructor = coroutineFactoryConstructor ?: coroutineConstructor,
-                stateMachineFunction = invokeSuspendMethod
-            )
-        }
-
-        fun buildConstructor(): IrConstructor = WrappedClassConstructorDescriptor().let { d ->
-            IrConstructorImpl(
-                startOffset, endOffset,
-                DECLARATION_ORIGIN_COROUTINE_IMPL,
-                IrConstructorSymbolImpl(d),
-                coroutineBaseClassConstructor.name,
-                irFunction.visibility,
-                coroutineClass.defaultType,
-                isInline = false,
-                isExternal = false,
-                isPrimary = true,
-                isExpect = false
-            ).apply {
-                d.bind(this)
-                parent = coroutineClass
-                coroutineClass.declarations += this
-                coroutineConstructors += this
-
-                valueParameters = functionParameters.mapIndexed { index, parameter ->
-                    parameter.copyTo(this, DECLARATION_ORIGIN_COROUTINE_IMPL, index, defaultValue = null)
-                }
-                val continuationParameter = coroutineBaseClassConstructor.valueParameters[0]
-                valueParameters += continuationParameter.copyTo(
-                    this, DECLARATION_ORIGIN_COROUTINE_IMPL,
-                    index = valueParameters.size,
-                    type = continuationType,
-                    defaultValue = null
-                )
-
-                val irBuilder = context.createIrBuilder(symbol, startOffset, endOffset)
-                body = irBuilder.irBlockBody {
-                    val completionParameter = valueParameters.last()
-                    +irDelegatingConstructorCall(coroutineBaseClassConstructor).apply {
-                        putValueArgument(0, irGet(completionParameter))
                     }
-                    +IrInstanceInitializerCallImpl(startOffset, endOffset, coroutineClass.symbol, context.irBuiltIns.unitType)
+                    val continuationParameter = coroutineBaseClassConstructor.valueParameters[0]
-
+                    valueParameters += continuationParameter.copyTo(
-                    functionParameters.forEachIndexed { index, parameter ->
+                        this, DECLARATION_ORIGIN_COROUTINE_IMPL,
-                        +irSetField(
+                        index = valueParameters.size,
-                            irGet(coroutineClassThis),
+                        type = continuationType,
-                            argumentToPropertiesMap.getValue(parameter),
+                        defaultValue = null
-                            irGet(valueParameters[index])
-                        )
-                    }
-                }
-            }
-        }
-
-        private fun buildFactoryConstructor(boundParams: List<IrValueParameter>) = WrappedClassConstructorDescriptor().let { d ->
-            IrConstructorImpl(
-                startOffset, endOffset,
-                DECLARATION_ORIGIN_COROUTINE_IMPL,
-                IrConstructorSymbolImpl(d),
-                coroutineBaseClassConstructor.name,
-                irFunction.visibility,
-                coroutineClass.defaultType,
-                isInline = false,
-                isExternal = false,
-                isPrimary = false,
-                isExpect = false
-            ).apply {
-                d.bind(this)
-                parent = coroutineClass
-                coroutineClass.declarations += this
-                coroutineConstructors += this
-
-                valueParameters = boundParams.mapIndexed { index, parameter ->
-                    parameter.copyTo(this, DECLARATION_ORIGIN_COROUTINE_IMPL, index)
-                }
-
-                val irBuilder = context.createIrBuilder(symbol, startOffset, endOffset)
-                body = irBuilder.irBlockBody {
-                    +irDelegatingConstructorCall(coroutineBaseClassConstructor).apply {
-                        putValueArgument(0, irNull()) // Completion.
-                    }
-                    +IrInstanceInitializerCallImpl(
-                        startOffset, endOffset, coroutineClass.symbol,
-                        context.irBuiltIns.unitType
                     )
-                    // Save all arguments to fields.
-                    boundParams.forEachIndexed { index, parameter ->
-                        +irSetField(
-                            irGet(coroutineClassThis), argumentToPropertiesMap.getValue(parameter),
-                            irGet(valueParameters[index])
-                        )
-                    }
-                }
-            }
-        }
 
-        private fun buildCreateMethod(
+                    val irBuilder = context.createIrBuilder(symbol, startOffset, endOffset)
-            unboundArgs: List<IrValueParameter>,
+                    body = irBuilder.irBlockBody {
-            superFunctionSymbol: IrSimpleFunctionSymbol?,
+                        val completionParameter = valueParameters.last()
-            coroutineConstructor: IrConstructor
+                        +irDelegatingConstructorCall(coroutineBaseClassConstructor).apply {
-        ) = WrappedSimpleFunctionDescriptor().let { d ->
+                            putValueArgument(0, irGet(completionParameter))
-            IrFunctionImpl(
-                startOffset, endOffset,
-                DECLARATION_ORIGIN_COROUTINE_IMPL,
-                IrSimpleFunctionSymbolImpl(d),
-                Name.identifier("create"),
-                Visibilities.PROTECTED,
-                Modality.FINAL,
-                coroutineClass.defaultType,
-                isInline = false,
-                isExternal = false,
-                isTailrec = false,
-                isSuspend = false,
-                isExpect = false,
-                isFakeOverride = false,
-                isOperator = false
-            ).apply {
-                d.bind(this)
-                parent = coroutineClass
-                coroutineClass.declarations += this
-
-                typeParameters = irFunction.typeParameters.map { parameter ->
-                    parameter.copyToWithoutSuperTypes(this, origin = DECLARATION_ORIGIN_COROUTINE_IMPL)
-                        .apply { superTypes += parameter.superTypes }
-                }
-
-                valueParameters = (unboundArgs + create1CompletionParameter).mapIndexed { index, parameter ->
-                    parameter.copyTo(this, DECLARATION_ORIGIN_COROUTINE_IMPL, index)
-                }
-
-                this.createDispatchReceiverParameter()
-
-                superFunctionSymbol?.let {
-                    overriddenSymbols += it.owner.overriddenSymbols
-                    overriddenSymbols += it
-                }
-
-                val thisReceiver = this.dispatchReceiverParameter!!
-
-                val irBuilder = context.createIrBuilder(symbol, startOffset, endOffset)
-                body = irBuilder.irBlockBody(startOffset, endOffset) {
-                    +irReturn(
-                        irCall(coroutineConstructor).apply {
-                            var unboundIndex = 0
-                            val unboundArgsSet = unboundArgs.toSet()
-                            functionParameters.map {
-                                if (unboundArgsSet.contains(it))
-                                    irGet(valueParameters[unboundIndex++])
-                                else
-                                    irGetField(irGet(thisReceiver), argumentToPropertiesMap.getValue(it))
-                            }.forEachIndexed { index, argument ->
-                                putValueArgument(index, argument)
-                            }
-                            putValueArgument(functionParameters.size, irGet(valueParameters[unboundIndex]))
-                            assert(unboundIndex == valueParameters.size - 1) {
-                                "Not all arguments of <create> are used"
-                            }
-                        })
-                }
-            }
-        }
-
-        private fun buildInvokeMethod(
-            suspendFunctionInvokeFunctionSymbol: IrSimpleFunctionSymbol,
-            functionInvokeFunctionSymbol: IrSimpleFunctionSymbol,
-            createFunction: IrFunction,
-            stateMachineFunction: IrFunction
-        ) = WrappedSimpleFunctionDescriptor().let { d ->
-            IrFunctionImpl(
-                startOffset, endOffset,
-                DECLARATION_ORIGIN_COROUTINE_IMPL,
-                IrSimpleFunctionSymbolImpl(d),
-                Name.identifier("invoke"),
-                Visibilities.PROTECTED,
-                Modality.FINAL,
-                context.irBuiltIns.anyNType,
-                isInline = false,
-                isExternal = false,
-                isTailrec = false,
-                isSuspend = true,
-                isExpect = false,
-                isFakeOverride = false,
-                isOperator = false
-            ).apply {
-                d.bind(this)
-                parent = coroutineClass
-                coroutineClass.declarations += this
-
-                typeParameters = irFunction.typeParameters.map { parameter ->
-                    parameter.copyToWithoutSuperTypes(this, origin = DECLARATION_ORIGIN_COROUTINE_IMPL)
-                        .apply { superTypes += parameter.superTypes }
-                }
-
-                valueParameters = createFunction.valueParameters
-                    // Skip completion - invoke() already has it implicitly as a suspend function.
-                    .take(createFunction.valueParameters.size - 1)
-                    .mapIndexed { index, parameter ->
-                        parameter.copyTo(this, DECLARATION_ORIGIN_COROUTINE_IMPL, index)
-                    }
-
-                this.createDispatchReceiverParameter()
-
-                overriddenSymbols += functionInvokeFunctionSymbol
-                overriddenSymbols += suspendFunctionInvokeFunctionSymbol
-
-                val thisReceiver = dispatchReceiverParameter!!
-
-                val irBuilder = context.createIrBuilder(symbol, startOffset, endOffset)
-                body = irBuilder.irBlockBody(startOffset, endOffset) {
-                    generateCoroutineStart(stateMachineFunction, irCall(createFunction).apply {
-                        dispatchReceiver = irGet(thisReceiver)
-                        valueParameters.forEachIndexed { index, parameter ->
-                            putValueArgument(index, irGet(parameter))
                         }
-                        putValueArgument(
+                        +IrInstanceInitializerCallImpl(startOffset, endOffset, coroutineClass.symbol, context.irBuiltIns.unitType)
-                            valueParameters.size,
+
-                            irCall(getContinuationSymbol, getContinuationSymbol.owner.returnType, listOf(returnType))
+                        functionParameters.forEachIndexed { index, parameter ->
-                        )
+                            +irSetField(
-                    })
+                                irGet(coroutineClassThis),
+                                argumentToPropertiesMap.getValue(parameter),
+                                irGet(valueParameters[index])
+                            )
+                        }
+                    }
                 }
             }
         }
 
-        private fun buildInvokeSuspendMethod(
+        private fun buildInvokeSuspendMethod(stateMachineFunction: IrSimpleFunction): IrSimpleFunction {
-            stateMachineFunction: IrSimpleFunction,
+            val smFunction = WrappedSimpleFunctionDescriptor().let { d ->
-            coroutineClass: IrClass
-        ): IrSimpleFunction {
-            val function = WrappedSimpleFunctionDescriptor().let { d ->
                 IrFunctionImpl(
                     startOffset, endOffset,
                     DECLARATION_ORIGIN_COROUTINE_IMPL,
@@ -635,7 +329,7 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
                 ).apply {
                     d.bind(this)
                     parent = coroutineClass
-                    coroutineClass.declarations += this
+                    coroutineClass.addChild(this)
 
                     typeParameters = stateMachineFunction.typeParameters.map { parameter ->
                         parameter.copyToWithoutSuperTypes(this, origin = DECLARATION_ORIGIN_COROUTINE_IMPL)
@@ -648,15 +342,188 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
 
                     this.createDispatchReceiverParameter()
 
-                    overriddenSymbols += stateMachineFunction.symbol
+                    overriddenSymbols = listOf(stateMachineFunction.symbol)
                 }
             }
 
-            buildStateMachine(function, irFunction, argumentToPropertiesMap)
+            buildStateMachine(smFunction, function, argumentToPropertiesMap)
-            return function
+            return smFunction
+        }
+
+        private fun buildCreateMethod(superCreateFunction: IrSimpleFunction?, constructor: IrConstructor): IrSimpleFunction {
+            // val i = $lambdaN(this.f1, this.f2, ..., this.fn, continuation) // bound
+            // i.s1 = p1 // unbound
+            // ...
+            // i.sn = pn
+            // return i
+            val createFunction = WrappedSimpleFunctionDescriptor().let { d ->
+                IrFunctionImpl(
+                    startOffset, endOffset,
+                    DECLARATION_ORIGIN_COROUTINE_IMPL,
+                    IrSimpleFunctionSymbolImpl(d),
+                    Name.identifier("create"),
+                    Visibilities.PROTECTED,
+                    Modality.FINAL,
+                    coroutineClass.defaultType,
+                    isInline = false,
+                    isExternal = false,
+                    isTailrec = false,
+                    isSuspend = false,
+                    isExpect = false,
+                    isFakeOverride = false,
+                    isOperator = false
+                ).apply {
+                    d.bind(this)
+                    parent = coroutineClass
+                    coroutineClass.addChild(this)
+
+                    typeParameters = function.typeParameters.map { parameter ->
+                        parameter.copyToWithoutSuperTypes(this, origin = DECLARATION_ORIGIN_COROUTINE_IMPL)
+                            .apply { superTypes += parameter.superTypes }
+                    }
+
+                    val unboundArgs = function.valueParameters
+
+                    val createValueParameters = (unboundArgs + create1CompletionParameter).mapIndexed { index, parameter ->
+                        parameter.copyTo(this, DECLARATION_ORIGIN_COROUTINE_IMPL, index)
+                    }
+
+                    valueParameters = createValueParameters
+
+                    this.createDispatchReceiverParameter()
+
+                    superCreateFunction?.let {
+                        overriddenSymbols = ArrayList<IrSimpleFunctionSymbol>(it.overriddenSymbols.size + 1).apply {
+                            addAll(it.overriddenSymbols)
+                            add(it.symbol)
+                        }
+                    }
+
+                    val thisReceiver = this.dispatchReceiverParameter!!
+
+                    val boundFields = coroutineClass.declarations.filterIsInstance<IrField>().filter { f ->
+                        f.origin == LocalDeclarationsLowering.DECLARATION_ORIGIN_FIELD_FOR_CAPTURED_VALUE || f.origin == LocalDeclarationsLowering.DECLARATION_ORIGIN_FIELD_FOR_CROSSINLINE_CAPTURED_VALUE
+                    }
+
+                    val irBuilder = context.createIrBuilder(symbol, startOffset, endOffset)
+                    body = irBuilder.irBlockBody(startOffset, endOffset) {
+                        val instanceCreate = irCall(constructor).apply {
+                            var unboundIndex = 0
+                            for (f in boundFields) {
+                                putValueArgument(unboundIndex++, irGetField(irGet(thisReceiver), f))
+                            }
+                            putValueArgument(unboundIndex++, irGet(createValueParameters.last()))
+                            assert(unboundIndex == constructor.valueParameters.size) {
+                                "Not all arguments of <create> are used"
+                            }
+                        }
+                        val instanceVal = scope.createTmpVariable(instanceCreate, "i")
+                        +instanceVal
+
+                        assert(createValueParameters.size - 1 == argumentToPropertiesMap.size)
+
+                        val ptofMap = createValueParameters.zip(argumentToPropertiesMap.values)
+
+                        for ((p, f) in ptofMap) {
+                            +irSetField(irGet(instanceVal), f, irGet(p))
+                        }
+
+                        +irReturn(irGet(instanceVal))
+                    }
+                }
+            }
+
+            return createFunction
+        }
+
+        private fun transformInvokeMethod(createFunction: IrSimpleFunction, stateMachineFunction: IrSimpleFunction) {
+            val irBuilder = context.createIrBuilder(function.symbol, startOffset, endOffset)
+            val thisReceiver = function.dispatchReceiverParameter ?: error("Expected dispatch receiver for invoke")
+            val functionBody = function.body as IrBlockBody
+            functionBody.statements.clear()
+            functionBody.statements.addAll(irBuilder.irBlockBody(startOffset, endOffset) {
+                generateCoroutineStart(stateMachineFunction, irCall(createFunction).apply {
+                    dispatchReceiver = irGet(thisReceiver)
+                    var index = 0
+                    for (parameter in function.valueParameters) {
+                        putValueArgument(index++, irGet(parameter))
+                    }
+                    putValueArgument(
+                        index++,
+                        irCall(getContinuationSymbol, getContinuationSymbol.owner.returnType, listOf(function.returnType))
+                    )
+                    assert(index == createFunction.valueParameters.size)
+                })
+            }.statements)
+        }
+
+        fun build(): BuiltCoroutine {
+            val coroutineConstructor = buildConstructor()
+
+            val superInvokeSuspendFunction = coroutineBaseClass.owner.simpleFunctions().single { it.name == stateMachineMethodName }
+            val invokeSuspendMethod = buildInvokeSuspendMethod(superInvokeSuspendFunction)
+
+            if (isSuspendLambda) {
+                // Suspend lambda - create factory methods.
+                val createFunction = coroutineBaseClass.owner.simpleFunctions()
+                    .atMostOne {
+                        it.name.asString() == "create" && it.valueParameters.size == function.valueParameters.size + 1
+                    }
+
+                val createMethod = buildCreateMethod(createFunction, coroutineConstructor)
+
+                transformInvokeMethod(createMethod, invokeSuspendMethod)
+            } else {
+                coroutineClass.superTypes += coroutineBaseClass.defaultType
+            }
+
+            coroutineClass.addFakeOverrides()
+
+            initializeStateMachine(listOf(coroutineConstructor), coroutineClassThis)
+
+            return BuiltCoroutine(
+                coroutineClass = coroutineClass,
+                coroutineConstructor = coroutineConstructor,
+                stateMachineFunction = invokeSuspendMethod
+            )
+
         }
     }
 
+    // Suppress since it is used in native
+    @Suppress("MemberVisibilityCanBePrivate")
+    protected fun IrCall.isReturnIfSuspendedCall() =
+        symbol.owner.run { fqNameWhenAvailable == context.internalPackageFqn.child(Name.identifier("returnIfSuspended")) }
+
+    private sealed class SuspendFunctionKind {
+        object NO_SUSPEND_CALLS : SuspendFunctionKind()
+        class DELEGATING(val delegatingCall: IrCall) : SuspendFunctionKind()
+        object NEEDS_STATE_MACHINE : SuspendFunctionKind()
+    }
+
+    private val symbols = context.ir.symbols
+    private val getContinuationSymbol = symbols.getContinuation
+    private val continuationClassSymbol = getContinuationSymbol.owner.returnType.classifierOrFail as IrClassSymbol
+
+    private fun removeReturnIfSuspendedCallAndSimplifyDelegatingCall(irFunction: IrFunction, delegatingCall: IrCall) {
+        val returnValue =
+            if (delegatingCall.isReturnIfSuspendedCall())
+                delegatingCall.getValueArgument(0)!!
+            else delegatingCall
+        context.createIrBuilder(irFunction.symbol).run {
+            val statements = (irFunction.body as IrBlockBody).statements
+            val lastStatement = statements.last()
+            assert(lastStatement == delegatingCall || lastStatement is IrReturn) { "Unexpected statement $lastStatement" }
+            statements[statements.lastIndex] = irReturn(generateDelegatedCall(irFunction.returnType, returnValue))
+        }
+    }
+
+    private class BuiltCoroutine(
+        val coroutineClass: IrClass,
+        val coroutineConstructor: IrConstructor,
+        val stateMachineFunction: IrFunction
+    )
+
     protected open class VariablesScopeTracker : IrElementVisitorVoid {
 
         protected val scopeStack = mutableListOf<MutableSet<IrVariable>>(mutableSetOf())
@@ -708,12 +575,12 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
     }
 }
 
-private val IrDeclaration.isLambda
+//private val IrDeclaration.isLambda
-    get() = origin == IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA
+//    get() = origin == IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA
 
-class RemoveSuspendLambdas() : DeclarationTransformer {
+//class RemoveSuspendLambdas() : DeclarationTransformer {
-
+//
-    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
+//    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
-        return if (declaration.isLambda && declaration is IrFunction && declaration.isSuspend) emptyList() else null
+//        return if (declaration.isLambda && declaration is IrFunction && declaration.isSuspend) emptyList() else null
-    }
+//    }
-}
+//}

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/coroutines/JsSuspendFunctionsLowering.kt

@@ -158,7 +158,9 @@ class JsSuspendFunctionsLowering(ctx: JsIrBackendContext) : AbstractSuspendFunct
                 }
             }
         }
-        simplifiedFunction.explicitParameters.forEach {
+        val isSuspendLambda = transformingFunction.parent === coroutineClass
+        val parameters = if (isSuspendLambda) simplifiedFunction.valueParameters else simplifiedFunction.explicitParameters
+        parameters.forEach {
             localToPropertyMap.getOrPut(it.symbol) {
                 argumentToPropertiesMap.getValue(it).symbol
             }