[IR] SuspendFunctionLowering refactoring

* distinguish common part which generates successor of `CoroutineImpl`
   into separate common lowering
 * merge it with K/N

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

@@ -0,0 +1,675 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. 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.common.lower
+
+import org.jetbrains.kotlin.backend.common.*
+import org.jetbrains.kotlin.backend.common.descriptors.*
+import org.jetbrains.kotlin.backend.common.ir.*
+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.builders.*
+import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.declarations.impl.*
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.expressions.impl.*
+import org.jetbrains.kotlin.ir.symbols.*
+import org.jetbrains.kotlin.ir.symbols.impl.*
+import org.jetbrains.kotlin.ir.types.*
+import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
+import org.jetbrains.kotlin.ir.util.*
+import org.jetbrains.kotlin.ir.util.defaultType
+import org.jetbrains.kotlin.ir.visitors.*
+import org.jetbrains.kotlin.name.Name
+
+abstract class AbstractSuspendFunctionsLowering<C: CommonBackendContext>(val context: C, val symbolTable: SymbolTable) : FileLoweringPass {
+
+    protected object STATEMENT_ORIGIN_COROUTINE_IMPL : IrStatementOriginImpl("COROUTINE_IMPL")
+    protected object DECLARATION_ORIGIN_COROUTINE_IMPL : IrDeclarationOriginImpl("COROUTINE_IMPL")
+
+    protected abstract val stateMachineMethodName: Name
+    protected abstract fun getCoroutineBaseClass(function: IrFunction): IrClassSymbol
+
+
+    protected abstract fun buildStateMachine(
+        originalBody: IrBody, stateMachineFunction: IrFunction,
+        transformingFunction: IrFunction,
+        argumentToPropertiesMap: Map<IrValueParameter, IrField>
+    )
+
+    protected abstract fun IrBlockBodyBuilder.generateCoroutineStart(invokeSuspendFunction: IrFunction, receiver: IrExpression)
+
+    protected abstract fun initializeStateMachine(coroutineConstructors: List<IrConstructor>, coroutineClassThis: IrValueDeclaration)
+
+
+    private val builtCoroutines = mutableMapOf<IrFunction, BuiltCoroutine>()
+    private val suspendLambdas = mutableMapOf<IrFunction, IrFunctionReference>()
+
+    override fun lower(irFile: IrFile) {
+        markSuspendLambdas(irFile)
+        buildCoroutines(irFile)
+        transformCallableReferencesToSuspendLambdas(irFile)
+    }
+
+    private fun buildCoroutines(irFile: IrFile) {
+        irFile.transformDeclarationsFlat(::tryTransformSuspendFunction)
+        irFile.acceptVoid(object : IrElementVisitorVoid {
+            override fun visitElement(element: IrElement) {
+                element.acceptChildrenVoid(this)
+            }
+
+            override fun visitClass(declaration: IrClass) {
+                declaration.acceptChildrenVoid(this)
+                declaration.transformDeclarationsFlat(::tryTransformSuspendFunction)
+            }
+        })
+    }
+
+
+    // Suppress since it is used in native
+    @Suppress("MemberVisibilityCanBePrivate")
+    protected fun IrCall.isReturnIfSuspendedCall() =
+        symbol.owner.run { fqNameSafe == context.internalPackageFqn.child(Name.identifier("returnIfSuspended")) }
+
+    private fun tryTransformSuspendFunction(element: IrElement) =
+        if (element is IrSimpleFunction && element.isSuspend && element.modality != Modality.ABSTRACT)
+            transformSuspendFunction(element, suspendLambdas[element])
+        else null
+
+    private fun markSuspendLambdas(irElement: IrElement) {
+        irElement.acceptChildrenVoid(object : IrElementVisitorVoid {
+            override fun visitElement(element: IrElement) {
+                element.acceptChildrenVoid(this)
+            }
+
+            override fun visitFunctionReference(expression: IrFunctionReference) {
+                expression.acceptChildrenVoid(this)
+
+                if (expression.isSuspend) {
+                    suspendLambdas[expression.symbol.owner] = expression
+                }
+            }
+        })
+    }
+
+    private fun transformCallableReferencesToSuspendLambdas(irElement: IrElement) {
+        irElement.transformChildrenVoid(object : IrElementTransformerVoid() {
+
+            override fun visitFunctionReference(expression: IrFunctionReference): IrExpression {
+                expression.transformChildrenVoid(this)
+
+                if (!expression.isSuspend)
+                    return expression
+                val coroutine = builtCoroutines[expression.symbol.owner]
+                    ?: throw Error("Non-local callable reference to suspend lambda: $expression")
+                val constructorParameters = coroutine.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(coroutine.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)) {
+            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 in suspendLambdas)             // 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): SuspendFunctionKind {
+        if (irFunction in suspendLambdas)
+            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) {
+                element.acceptChildrenVoid(this)
+            }
+
+            override fun visitCall(expression: IrCall) {
+                expression.acceptChildrenVoid(this)
+                if (expression.isSuspend)
+                    ++numberOfSuspendCalls
+            }
+        })
+        // It is important to optimize the case where there is only one suspend call and it is the last statement
+        // because we don't need to build a fat coroutine class in that case.
+        // This happens a lot in practise because of suspend functions with default arguments.
+        // TODO: use TailRecursionCallsCollector.
+        val lastCall = when (val lastStatement = (body as IrBlockBody).statements.lastOrNull()) {
+            is IrCall -> lastStatement
+            is IrReturn -> {
+                var value: IrElement = lastStatement
+                /*
+                 * Check if matches this pattern:
+                 * block/return {
+                 *     block/return {
+                 *         .. suspendCall()
+                 *     }
+                 * }
+                 */
+                loop@ while (true) {
+                    value = when {
+                        value is IrBlock && value.statements.size == 1 -> value.statements.first()
+                        value is IrReturn -> value.value
+                        else -> break@loop
+                    }
+                }
+                value as? IrCall
+            }
+            else -> null
+        }
+        val suspendCallAtEnd = lastCall != null && lastCall.isSuspend    // Suspend call.
+        return when {
+            numberOfSuspendCalls == 0 -> SuspendFunctionKind.NO_SUSPEND_CALLS
+            numberOfSuspendCalls == 1
+                    && suspendCallAtEnd -> SuspendFunctionKind.DELEGATING(lastCall!!)
+            else -> SuspendFunctionKind.NEEDS_STATE_MACHINE
+        }
+    }
+
+    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(returnValue)
+        }
+    }
+
+    private fun buildCoroutine(irFunction: IrSimpleFunction, functionReference: IrFunctionReference?): IrClass {
+        val coroutine = CoroutineBuilder(irFunction, functionReference).build()
+        builtCoroutines[irFunction] = coroutine
+
+        if (functionReference == null) {
+            // It is not a lambda - replace original function with a call to constructor of the built coroutine.
+            val irBuilder = context.createIrBuilder(irFunction.symbol, irFunction.startOffset, irFunction.endOffset)
+            irFunction.body = 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)
+                                               )
+                                           )
+                                       })
+            }
+        }
+
+        return coroutine.coroutineClass
+    }
+
+    private class BuiltCoroutine(
+        val coroutineClass: IrClass,
+        val coroutineConstructor: IrConstructor,
+        val stateMachineFunction: IrFunction
+    )
+
+    private var coroutineId = 0
+
+    private inner class CoroutineBuilder(val irFunction: IrFunction, val functionReference: IrFunctionReference?) {
+
+        private val startOffset = irFunction.startOffset
+        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 ->
+            IrClassImpl(
+                startOffset, endOffset,
+                DECLARATION_ORIGIN_COROUTINE_IMPL,
+                symbolTable.referenceClass(d),
+                "${irFunction.name}COROUTINE\$${coroutineId++}".synthesizedName,
+                ClassKind.CLASS,
+                irFunction.visibility,
+                Modality.FINAL,
+                isCompanion = false,
+                isInner = false,
+                isData = false,
+                isExternal = false,
+                isInline = false
+            ).apply {
+                d.bind(this)
+                parent = irFunction.parent
+                createParameterDeclarations()
+                irFunction.typeParameters.mapTo(typeParameters) { typeParam ->
+                    typeParam.copyToWithoutSuperTypes(this).apply { superTypes += typeParam.superTypes }
+                }
+            }
+        }
+        private val coroutineClassThis = coroutineClass.thisReceiver!!
+
+        private val continuationType = continuationClassSymbol.typeWith(irFunction.returnType)
+
+        // Save all arguments to fields.
+        private val argumentToPropertiesMap = functionParameters.associate {
+            it to 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.owner.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()
+
+            val superInvokeSuspendFunction = coroutineBaseClass.owner.simpleFunctions().single { it.name == stateMachineMethodName }
+            val invokeSuspendMethod = buildInvokeSuspendMethod(superInvokeSuspendFunction, coroutineClass)
+
+            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,
+                symbolTable.referenceConstructor(d),
+                coroutineBaseClassConstructor.name,
+                irFunction.visibility,
+                coroutineClass.defaultType,
+                isInline = false,
+                isExternal = false,
+                isPrimary = false
+            ).apply {
+                d.bind(this)
+                parent = coroutineClass
+                coroutineClass.declarations += this
+                coroutineConstructors += this
+
+                functionParameters.mapIndexedTo(valueParameters) { index, parameter ->
+                    parameter.copyTo(this, DECLARATION_ORIGIN_COROUTINE_IMPL, index)
+                }
+                val continuationParameter = coroutineBaseClassConstructor.valueParameters[0]
+                valueParameters += continuationParameter.copyTo(
+                    this, DECLARATION_ORIGIN_COROUTINE_IMPL,
+                    index = valueParameters.size, type = continuationType
+                )
+
+                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)
+
+                    functionParameters.forEachIndexed { index, parameter ->
+                        +irSetField(
+                            irGet(coroutineClassThis),
+                            argumentToPropertiesMap.getValue(parameter),
+                            irGet(valueParameters[index])
+                        )
+                    }
+                }
+            }
+        }
+
+        private fun buildFactoryConstructor(boundParams: List<IrValueParameter>) = WrappedClassConstructorDescriptor().let { d ->
+            IrConstructorImpl(
+                startOffset, endOffset,
+                DECLARATION_ORIGIN_COROUTINE_IMPL,
+                symbolTable.referenceConstructor(d),
+                coroutineBaseClassConstructor.name,
+                irFunction.visibility,
+                coroutineClass.defaultType,
+                isInline = false,
+                isExternal = false,
+                isPrimary = false
+            ).apply {
+                d.bind(this)
+                parent = coroutineClass
+                coroutineClass.declarations += this
+                coroutineConstructors += this
+
+                boundParams.mapIndexedTo(valueParameters) { 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(
+            unboundArgs: List<IrValueParameter>,
+            superFunctionSymbol: IrSimpleFunctionSymbol?,
+            coroutineConstructor: IrConstructor
+        ) = WrappedSimpleFunctionDescriptor().let { d ->
+            IrFunctionImpl(
+                startOffset, endOffset,
+                DECLARATION_ORIGIN_COROUTINE_IMPL,
+                symbolTable.referenceSimpleFunction(d),
+                Name.identifier("create"),
+                Visibilities.PROTECTED,
+                Modality.FINAL,
+                coroutineClass.defaultType,
+                isInline = false,
+                isExternal = false,
+                isTailrec = false,
+                isSuspend = false
+            ).apply {
+                d.bind(this)
+                parent = coroutineClass
+                coroutineClass.declarations += this
+
+                irFunction.typeParameters.mapTo(typeParameters) { parameter ->
+                    parameter.copyToWithoutSuperTypes(this, 0, DECLARATION_ORIGIN_COROUTINE_IMPL)
+                        .apply { superTypes += parameter.superTypes }
+                }
+
+                (unboundArgs + create1CompletionParameter)
+                    .mapIndexedTo(valueParameters) { 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,
+                symbolTable.referenceSimpleFunction(d),
+                Name.identifier("invoke"),
+                Visibilities.PROTECTED,
+                Modality.FINAL,
+                irFunction.returnType,
+                isInline = false,
+                isExternal = false,
+                isTailrec = false,
+                isSuspend = true
+            ).apply {
+                d.bind(this)
+                parent = coroutineClass
+                coroutineClass.declarations += this
+
+                irFunction.typeParameters.mapTo(typeParameters) { parameter ->
+                    parameter.copyToWithoutSuperTypes(this, 0, DECLARATION_ORIGIN_COROUTINE_IMPL)
+                        .apply { superTypes += parameter.superTypes }
+                }
+
+                createFunction.valueParameters
+                    // Skip completion - invoke() already has it implicitly as a suspend function.
+                    .take(createFunction.valueParameters.size - 1)
+                    .mapIndexedTo(valueParameters) { 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(
+                            valueParameters.size,
+                            irCall(getContinuationSymbol, getContinuationSymbol.owner.returnType, listOf(returnType))
+                        )
+                    })
+                }
+            }
+        }
+
+        private fun buildInvokeSuspendMethod(
+            stateMachineFunction: IrSimpleFunction,
+            coroutineClass: IrClass
+        ): IrSimpleFunction {
+            val originalBody = irFunction.body!!
+            val function = WrappedSimpleFunctionDescriptor().let { d ->
+                IrFunctionImpl(
+                    startOffset, endOffset,
+                    DECLARATION_ORIGIN_COROUTINE_IMPL,
+                    symbolTable.referenceSimpleFunction(d),
+                    stateMachineFunction.name,
+                    stateMachineFunction.visibility,
+                    Modality.FINAL,
+                    context.irBuiltIns.anyNType,
+                    stateMachineFunction.isInline,
+                    stateMachineFunction.isExternal,
+                    stateMachineFunction.isTailrec,
+                    stateMachineFunction.isSuspend
+                ).apply {
+                    d.bind(this)
+                    parent = coroutineClass
+                    coroutineClass.declarations += this
+
+                    stateMachineFunction.typeParameters.mapTo(typeParameters) { parameter ->
+                        parameter.copyToWithoutSuperTypes(this, 0, DECLARATION_ORIGIN_COROUTINE_IMPL)
+                            .apply { superTypes += parameter.superTypes }
+                    }
+
+                    stateMachineFunction.valueParameters.mapIndexedTo(valueParameters) { index, parameter ->
+                        parameter.copyTo(this, DECLARATION_ORIGIN_COROUTINE_IMPL, index)
+                    }
+
+                    this.createDispatchReceiverParameter()
+
+                    overriddenSymbols += stateMachineFunction.symbol
+                }
+            }
+
+            buildStateMachine(originalBody, function, irFunction, argumentToPropertiesMap)
+            return function
+        }
+    }
+
+    protected open class VariablesScopeTracker : IrElementVisitorVoid {
+
+        protected val scopeStack = mutableListOf<MutableSet<IrVariable>>(mutableSetOf())
+
+        override fun visitElement(element: IrElement) {
+            element.acceptChildrenVoid(this)
+        }
+
+        override fun visitContainerExpression(expression: IrContainerExpression) {
+            if (!expression.isTransparentScope)
+                scopeStack.push(mutableSetOf())
+            super.visitContainerExpression(expression)
+            if (!expression.isTransparentScope)
+                scopeStack.pop()
+        }
+
+        override fun visitCatch(aCatch: IrCatch) {
+            scopeStack.push(mutableSetOf())
+            super.visitCatch(aCatch)
+            scopeStack.pop()
+        }
+
+        override fun visitVariable(declaration: IrVariable) {
+            super.visitVariable(declaration)
+            scopeStack.peek()!!.add(declaration)
+        }
+    }
+
+    fun IrClass.addField(name: Name, type: IrType, isMutable: Boolean): IrField {
+        val descriptor = WrappedFieldDescriptor()
+        val symbol = IrFieldSymbolImpl(descriptor)
+        return IrFieldImpl(
+            startOffset,
+            endOffset,
+            DECLARATION_ORIGIN_COROUTINE_IMPL,
+            symbol,
+            name,
+            type,
+            Visibilities.PRIVATE,
+            !isMutable,
+            isExternal = false,
+            isStatic = false
+        ).also {
+            descriptor.bind(it)
+            it.parent = this
+            addChild(it)
+        }
+    }
+}

compiler/ir/backend.common/src/org/jetbrains/kotlin/ir/builders/IrBuilders.kt

@@ -74,8 +74,8 @@ fun Scope.createTmpVariable(
         Name.identifier(nameHint ?: "tmp"),
         varType,
         isMutable,
-        false,
+        isConst = false,
-        false
+        isLateinit = false
     ).apply {
         initializer = irExpression
         parent = getLocalDeclarationParent()

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

@@ -30,6 +30,7 @@ import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.declarations.impl.IrFileImpl
 import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
 import org.jetbrains.kotlin.ir.symbols.*
+import org.jetbrains.kotlin.ir.types.IrDynamicType
 import org.jetbrains.kotlin.ir.types.IrSimpleType
 import org.jetbrains.kotlin.ir.types.classifierOrFail
 import org.jetbrains.kotlin.ir.types.impl.IrDynamicTypeImpl
@@ -147,7 +148,7 @@ class JsIrBackendContext(
         return numbers + listOf(Name.identifier("String"))
     }
 
-    val dynamicType = IrDynamicTypeImpl(null, emptyList(), Variance.INVARIANT)
+    val dynamicType: IrDynamicType = IrDynamicTypeImpl(null, emptyList(), Variance.INVARIANT)
 
     fun getOperatorByName(name: Name, type: IrSimpleType) = operatorMap[name]?.get(type.classifier)
 

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

@@ -10,7 +10,7 @@ import org.jetbrains.kotlin.backend.common.lower.*
 import org.jetbrains.kotlin.backend.common.phaser.*
 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.coroutines.SuspendFunctionsLowering
+import org.jetbrains.kotlin.ir.backend.js.lower.coroutines.JsSuspendFunctionsLowering
 import org.jetbrains.kotlin.ir.backend.js.lower.inline.FunctionInlining
 import org.jetbrains.kotlin.ir.backend.js.lower.inline.RemoveInlineFunctionsWithReifiedTypeParametersLowering
 import org.jetbrains.kotlin.ir.backend.js.lower.inline.ReturnableBlockLowering
@@ -179,7 +179,7 @@ private val innerClassConstructorCallsLoweringPhase = makeJsModulePhase(
 )
 
 private val suspendFunctionsLoweringPhase = makeJsModulePhase(
-    ::SuspendFunctionsLowering,
+    ::JsSuspendFunctionsLowering,
     name = "SuspendFunctionsLowering",
     description = "Transform suspend functions into CoroutineImpl instance and build state machine",
     prerequisite = setOf(unitMaterializationLoweringPhase)

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

@@ -0,0 +1,212 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. 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.backend.js.lower.coroutines
+
+import org.jetbrains.kotlin.backend.common.ir.isSuspend
+import org.jetbrains.kotlin.backend.common.lower.AbstractSuspendFunctionsLowering
+import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
+import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
+import org.jetbrains.kotlin.ir.builders.*
+import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.expressions.IrBlockBody
+import org.jetbrains.kotlin.ir.expressions.IrBody
+import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.impl.*
+import org.jetbrains.kotlin.ir.symbols.IrFieldSymbol
+import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
+import org.jetbrains.kotlin.ir.util.explicitParameters
+import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
+import org.jetbrains.kotlin.ir.visitors.acceptVoid
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.utils.DFS
+
+class JsSuspendFunctionsLowering(ctx: JsIrBackendContext) : AbstractSuspendFunctionsLowering<JsIrBackendContext>(ctx, ctx.symbolTable) {
+
+    private val coroutineImplExceptionPropertyGetter = ctx.coroutineImplExceptionPropertyGetter
+    private val coroutineImplExceptionPropertySetter = ctx.coroutineImplExceptionPropertySetter
+    private val coroutineImplExceptionStatePropertyGetter = ctx.coroutineImplExceptionStatePropertyGetter
+    private val coroutineImplExceptionStatePropertySetter = ctx.coroutineImplExceptionStatePropertySetter
+    private val coroutineImplLabelPropertySetter = ctx.coroutineImplLabelPropertySetter
+    private val coroutineImplLabelPropertyGetter = ctx.coroutineImplLabelPropertyGetter
+    private val coroutineImplResultSymbolGetter = ctx.coroutineImplResultSymbolGetter
+    private val coroutineImplResultSymbolSetter = ctx.coroutineImplResultSymbolSetter
+
+    private var exceptionTrapId = -1
+
+    override val stateMachineMethodName = Name.identifier("doResume")
+    override fun getCoroutineBaseClass(function: IrFunction) = context.ir.symbols.coroutineImpl
+
+    override fun buildStateMachine(
+        originalBody: IrBody,
+        stateMachineFunction: IrFunction,
+        transformingFunction: IrFunction,
+        argumentToPropertiesMap: Map<IrValueParameter, IrField>
+    ) {
+        val body =
+            (originalBody as IrBlockBody).run {
+                IrBlockImpl(
+                    transformingFunction.startOffset,
+                    transformingFunction.endOffset,
+                    context.irBuiltIns.unitType,
+                    STATEMENT_ORIGIN_COROUTINE_IMPL,
+                    statements
+                )
+            }
+
+        val coroutineClass = stateMachineFunction.parent as IrClass
+        val suspendResult = JsIrBuilder.buildVar(
+            context.irBuiltIns.anyNType,
+            stateMachineFunction,
+            "suspendResult",
+            true,
+            initializer = JsIrBuilder.buildCall(coroutineImplResultSymbolGetter.symbol).apply {
+                dispatchReceiver = JsIrBuilder.buildGetValue(stateMachineFunction.dispatchReceiverParameter!!.symbol)
+            }
+        )
+
+        val suspendState = JsIrBuilder.buildVar(coroutineImplLabelPropertyGetter.returnType, stateMachineFunction, "suspendState", true)
+
+        val unit = context.irBuiltIns.unitType
+
+        val switch = IrWhenImpl(body.startOffset, body.endOffset, unit, COROUTINE_SWITCH)
+        val rootTry = IrTryImpl(body.startOffset, body.endOffset, unit).apply { tryResult = switch }
+        val rootLoop = IrDoWhileLoopImpl(
+            body.startOffset,
+            body.endOffset,
+            unit,
+            COROUTINE_ROOT_LOOP,
+            rootTry,
+            JsIrBuilder.buildBoolean(context.irBuiltIns.booleanType, true)
+        )
+
+        val suspendableNodes = mutableSetOf<IrElement>()
+        val loweredBody =
+            collectSuspendableNodes(body, suspendableNodes, context, stateMachineFunction, context.dynamicType)
+        val thisReceiver = (stateMachineFunction.dispatchReceiverParameter as IrValueParameter).symbol
+
+        val stateMachineBuilder = StateMachineBuilder(
+            suspendableNodes,
+            context,
+            stateMachineFunction.symbol,
+            rootLoop,
+            coroutineImplExceptionPropertyGetter,
+            coroutineImplExceptionPropertySetter,
+            coroutineImplExceptionStatePropertyGetter,
+            coroutineImplExceptionStatePropertySetter,
+            coroutineImplLabelPropertySetter,
+            thisReceiver,
+            suspendResult.symbol
+        )
+
+        loweredBody.acceptVoid(stateMachineBuilder)
+
+        stateMachineBuilder.finalizeStateMachine()
+
+        rootTry.catches += stateMachineBuilder.globalCatch
+
+        val visited = mutableSetOf<SuspendState>()
+
+        val sortedStates = DFS.topologicalOrder(listOf(stateMachineBuilder.entryState), { it.successors }, { visited.add(it) })
+        sortedStates.withIndex().forEach { it.value.id = it.index }
+
+        fun buildDispatch(target: SuspendState) = target.run {
+            assert(id >= 0)
+            JsIrBuilder.buildInt(context.irBuiltIns.intType, id)
+        }
+
+        val eqeqeqInt = context.irBuiltIns.eqeqeqSymbol
+
+        for (state in sortedStates) {
+            val condition = JsIrBuilder.buildCall(eqeqeqInt).apply {
+                putValueArgument(0, JsIrBuilder.buildCall(coroutineImplLabelPropertyGetter.symbol).also {
+                    it.dispatchReceiver = JsIrBuilder.buildGetValue(thisReceiver)
+                })
+                putValueArgument(1, JsIrBuilder.buildInt(context.irBuiltIns.intType, state.id))
+            }
+
+            switch.branches += IrBranchImpl(state.entryBlock.startOffset, state.entryBlock.endOffset, condition, state.entryBlock)
+        }
+
+        rootLoop.transform(DispatchPointTransformer(::buildDispatch), null)
+
+        exceptionTrapId = stateMachineBuilder.rootExceptionTrap.id
+
+        val functionBody =
+            IrBlockBodyImpl(stateMachineFunction.startOffset, stateMachineFunction.endOffset, listOf(suspendResult, rootLoop))
+
+        stateMachineFunction.body = functionBody
+
+        val liveLocals = computeLivenessAtSuspensionPoints(functionBody).values.flatten().toSet()
+
+        val localToPropertyMap = mutableMapOf<IrValueSymbol, IrFieldSymbol>()
+        var localCounter = 0
+        // TODO: optimize by using the same property for different locals.
+        liveLocals.forEach {
+            if (it != suspendState && it != suspendResult) {
+                localToPropertyMap.getOrPut(it.symbol) {
+                    coroutineClass.addField(Name.identifier("${it.name}${localCounter++}"), it.type, (it as? IrVariable)?.isVar ?: false)
+                        .symbol
+                }
+            }
+        }
+        transformingFunction.explicitParameters.forEach {
+            localToPropertyMap.getOrPut(it.symbol) {
+                argumentToPropertiesMap.getValue(it).symbol
+            }
+        }
+
+        stateMachineFunction.transform(LiveLocalsTransformer(localToPropertyMap, { JsIrBuilder.buildGetValue(thisReceiver) }, unit), null)
+    }
+
+    private fun computeLivenessAtSuspensionPoints(body: IrBody): Map<IrCall, List<IrValueDeclaration>> {
+        // TODO: data flow analysis.
+        // Just save all visible for now.
+        val result = mutableMapOf<IrCall, List<IrValueDeclaration>>()
+        body.acceptChildrenVoid(object : VariablesScopeTracker() {
+            override fun visitCall(expression: IrCall) {
+                if (!expression.isSuspend) return super.visitCall(expression)
+
+                expression.acceptChildrenVoid(this)
+                val visibleVariables = mutableListOf<IrValueDeclaration>()
+                scopeStack.forEach { visibleVariables += it }
+                result[expression] = visibleVariables
+            }
+        })
+
+        return result
+    }
+
+
+    override fun initializeStateMachine(coroutineConstructors: List<IrConstructor>, coroutineClassThis: IrValueDeclaration) {
+        for (it in coroutineConstructors) {
+            (it.body as? IrBlockBody)?.run {
+                val receiver = JsIrBuilder.buildGetValue(coroutineClassThis.symbol)
+                val id = JsIrBuilder.buildInt(context.irBuiltIns.intType, exceptionTrapId)
+                statements += JsIrBuilder.buildCall(coroutineImplExceptionStatePropertySetter.symbol).also { call ->
+                    call.dispatchReceiver = receiver
+                    call.putValueArgument(0, id)
+                }
+            }
+        }
+    }
+
+    override fun IrBlockBodyBuilder.generateCoroutineStart(invokeSuspendFunction: IrFunction, receiver: IrExpression) {
+        val dispatchReceiverVar = createTmpVariable(receiver, irType = receiver.type)
+        +irCall(coroutineImplResultSymbolSetter).apply {
+            dispatchReceiver = irGet(dispatchReceiverVar)
+            putValueArgument(0, irGetObject(context.irBuiltIns.unitClass))
+        }
+        +irCall(coroutineImplExceptionPropertySetter).apply {
+            dispatchReceiver = irGet(dispatchReceiverVar)
+            putValueArgument(0, irNull())
+        }
+        +irReturn(irCall(invokeSuspendFunction.symbol).apply {
+            dispatchReceiver = irGet(dispatchReceiverVar)
+        })
+    }
+}

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.ir.backend.js.lower.coroutines
 
+import org.jetbrains.kotlin.backend.common.CommonBackendContext
 import org.jetbrains.kotlin.backend.common.ir.isElseBranch
 import org.jetbrains.kotlin.backend.common.ir.isSuspend
 import org.jetbrains.kotlin.backend.common.peek
@@ -13,7 +14,6 @@ import org.jetbrains.kotlin.backend.common.push
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
-import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
 import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
 import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.declarations.IrVariable
@@ -61,7 +61,7 @@ class DispatchPointTransformer(val action: (SuspendState) -> IrExpression) : IrE
 
 class StateMachineBuilder(
     private val suspendableNodes: MutableSet<IrElement>,
-    val context: JsIrBackendContext,
+    val context: CommonBackendContext,
     val function: IrFunctionSymbol,
     private val rootLoop: IrLoop,
     private val exceptionSymbolGetter: IrSimpleFunction,
@@ -90,10 +90,6 @@ class StateMachineBuilder(
     lateinit var globalCatch: IrCatch
 
     fun finalizeStateMachine() {
-        val unitValue = JsIrBuilder.buildGetObjectValue(
-            unit,
-            context.symbolTable.referenceClass(context.builtIns.unit)
-        )
         globalCatch = buildGlobalCatch()
         if (currentBlock.statements.lastOrNull() !is IrReturn) {
             addStatement(JsIrBuilder.buildReturn(function, unitValue, nothing))
@@ -537,10 +533,7 @@ class StateMachineBuilder(
         })
     }
 
-    private val unitValue = JsIrBuilder.buildGetObjectValue(
+    private val unitValue get() = JsIrBuilder.buildGetObjectValue(unit, context.irBuiltIns.unitClass)
-        unit,
-        context.symbolTable.referenceClass(context.builtIns.unit)
-    )
 
     override fun visitReturn(expression: IrReturn) {
         expression.acceptChildrenVoid(this)

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

@@ -5,13 +5,13 @@
 
 package org.jetbrains.kotlin.ir.backend.js.lower.coroutines
 
+import org.jetbrains.kotlin.backend.common.CommonBackendContext
 import org.jetbrains.kotlin.backend.common.ir.isSuspend
 import org.jetbrains.kotlin.backend.common.lower.FinallyBlocksLowering
 import org.jetbrains.kotlin.backend.common.pop
 import org.jetbrains.kotlin.backend.common.push
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
-import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
 import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
 import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.declarations.IrVariable
@@ -25,7 +25,6 @@ import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.visitors.*
 
-
 object COROUTINE_ROOT_LOOP : IrStatementOriginImpl("COROUTINE_ROOT_LOOP")
 object COROUTINE_SWITCH : IrStatementOriginImpl("COROUTINE_SWITCH")
 
@@ -55,8 +54,9 @@ open class SuspendableNodesCollector(protected val suspendableNodes: MutableSet<
 fun collectSuspendableNodes(
     body: IrBlock,
     suspendableNodes: MutableSet<IrElement>,
-    context: JsIrBackendContext,
+    context: CommonBackendContext,
-    function: IrFunction
+    function: IrFunction,
+    throwableType: IrType
 ): IrBlock {
 
     // 1st: mark suspendable loops and tries
@@ -66,7 +66,7 @@ fun collectSuspendableNodes(
     body.acceptVoid(terminatorsCollector)
 
     if (terminatorsCollector.shouldFinalliesBeLowered) {
-        val finallyLower = FinallyBlocksLowering(context, context.dynamicType)
+        val finallyLower = FinallyBlocksLowering(context, throwableType)
 
         function.body = IrBlockBodyImpl(body.startOffset, body.endOffset, body.statements)
         function.transform(finallyLower, null)
@@ -76,7 +76,7 @@ fun collectSuspendableNodes(
         suspendableNodes.clear()
         val newBlock = JsIrBuilder.buildBlock(body.type, newBody.statements)
 
-        return collectSuspendableNodes(newBlock, suspendableNodes, context, function)
+        return collectSuspendableNodes(newBlock, suspendableNodes, context, function, throwableType)
     }
 
     return body

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/builders/ExpressionHelpers.kt

@@ -229,6 +229,13 @@ fun IrBuilderWithScope.irCall(callee: IrFunctionSymbol, type: IrType, typeArgume
         }
     }
 
+fun IrBuilderWithScope.irCallConstructor(callee: IrConstructorSymbol, typeArguments: List<IrType>): IrCall =
+    IrCallImpl(startOffset, endOffset, callee.owner.returnType, callee, callee.descriptor, typeArguments.size, callee.owner.valueParameters.size).apply {
+        typeArguments.forEachIndexed { index, irType ->
+            this.putTypeArgument(index, irType)
+        }
+    }
+
 fun IrBuilderWithScope.irCall(callee: IrFunctionSymbol): IrCall =
     IrCallImpl(startOffset, endOffset, callee.owner.returnType, callee, callee.descriptor)