Use WrappedSimpleFunctionDescriptor instead of real one for

suspend function views

compiler/backend/src/org/jetbrains/kotlin/codegen/state/KotlinTypeMapper.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2018 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -55,10 +55,7 @@ import org.jetbrains.kotlin.resolve.DescriptorUtils.*
 import org.jetbrains.kotlin.resolve.calls.model.DefaultValueArgument
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.VarargValueArgument
-import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
-import org.jetbrains.kotlin.resolve.descriptorUtil.classId
-import org.jetbrains.kotlin.resolve.descriptorUtil.isPublishedApi
-import org.jetbrains.kotlin.resolve.descriptorUtil.module
+import org.jetbrains.kotlin.resolve.descriptorUtil.*
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes.DEFAULT_CONSTRUCTOR_MARKER
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes.OBJECT_TYPE
 import org.jetbrains.kotlin.resolve.jvm.JvmClassName
@@ -1318,6 +1315,10 @@ class KotlinTypeMapper @JvmOverloads constructor(
                 return descriptor.containingDeclaration.name.asString()
             }
 
+            if (descriptor.containingDeclaration.fqNameSafe.asString() == "kotlin.coroutines.intrinsics") {
+                return "kotlin.coroutines.IntrinsicsKt" // TODO: why this change is needed in JVM_IR?
+            }
+
             if (directMember is FictitiousArrayConstructor) {
                 return "kotlin.Array"
             }

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/descriptors/WrappedDescriptors.kt

@@ -10,7 +10,6 @@ import org.jetbrains.kotlin.descriptors.annotations.AnnotationDescriptor
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationDescriptorImpl
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
 import org.jetbrains.kotlin.descriptors.impl.ReceiverParameterDescriptorImpl
-import org.jetbrains.kotlin.descriptors.impl.SimpleFunctionDescriptorImpl
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.descriptors.IrBasedDeclarationDescriptor
@@ -375,6 +374,7 @@ open class WrappedSimpleFunctionDescriptor(
     annotations: Annotations = Annotations.EMPTY,
     sourceElement: SourceElement = SourceElement.NO_SOURCE
 ) : SimpleFunctionDescriptor, WrappedCallableDescriptor<IrSimpleFunction>(annotations, sourceElement) {
+    private val userDataMap = mutableMapOf<CallableDescriptor.UserDataKey<*>, Any?>()
 
     // TODO: Remove as soon as all IR declarations have their originalDescriptor.
     constructor(originalDescriptor: FunctionDescriptor) : this(originalDescriptor.annotations, originalDescriptor.source) {
@@ -423,7 +423,7 @@ open class WrappedSimpleFunctionDescriptor(
     }
 
     override fun setOverriddenDescriptors(overriddenDescriptors: MutableCollection<out CallableMemberDescriptor>) {
-        TODO("not implemented")
+        if (!isSuspend) TODO("not implemented")
     }
 
     override fun getKind() =
@@ -450,21 +450,15 @@ open class WrappedSimpleFunctionDescriptor(
 
     override fun getInitialSignatureDescriptor() = null
 
-    override fun <V : Any?> getUserData(key: CallableDescriptor.UserDataKey<V>?): V? = null
+    override fun <V : Any?> getUserData(key: CallableDescriptor.UserDataKey<V>?): V? =
+        userDataMap[key as CallableDescriptor.UserDataKey<*>] as V?
 
     override fun newCopyBuilder(): FunctionDescriptor.CopyBuilder<out SimpleFunctionDescriptor> {
-        return SimpleFunctionDescriptorImpl
-            .create(containingDeclaration, annotations, name, kind, source)
-            .initialize(
-                extensionReceiverParameter,
-                dispatchReceiverParameter,
-                typeParameters,
-                valueParameters,
-                returnType,
-                modality,
-                visibility,
-                null
-            ).newCopyBuilder()
+        // KotlinTypeMapper uses getOrCreateJvmSuspendFunctionView extensively, but since
+        // after AddContinuationLowering all suspend descriptors are views anyway, return original
+        // one as a copy (with user data changed)
+        if (isSuspend) return ChangeUserDataCopyBuilder()
+        else TODO("not implemented")
     }
 
     override fun <R, D> accept(visitor: DeclarationDescriptorVisitor<R, D>?, data: D) =
@@ -473,6 +467,96 @@ open class WrappedSimpleFunctionDescriptor(
     override fun acceptVoid(visitor: DeclarationDescriptorVisitor<Void, Void>?) {
         visitor!!.visitFunctionDescriptor(this, null)
     }
+
+    inner class ChangeUserDataCopyBuilder : FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+        override fun setOwner(owner: DeclarationDescriptor): FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+            return this
+        }
+
+        override fun setModality(modality: Modality): FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+            return this
+        }
+
+        override fun setVisibility(visibility: Visibility): FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+            return this
+        }
+
+        override fun setKind(kind: CallableMemberDescriptor.Kind): FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+            return this
+        }
+
+        override fun setCopyOverrides(copyOverrides: Boolean): FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+            return this
+        }
+
+        override fun setName(name: Name): FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+            return this
+        }
+
+        override fun setValueParameters(parameters: MutableList<ValueParameterDescriptor>): FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+            return this
+        }
+
+        override fun setTypeParameters(parameters: MutableList<TypeParameterDescriptor>): FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+            return this
+        }
+
+        override fun setReturnType(type: KotlinType): FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+            return this
+        }
+
+        override fun setExtensionReceiverParameter(extensionReceiverParameter: ReceiverParameterDescriptor?): FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+            return this
+        }
+
+        override fun setDispatchReceiverParameter(dispatchReceiverParameter: ReceiverParameterDescriptor?): FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+            return this
+        }
+
+        override fun setOriginal(original: CallableMemberDescriptor?): FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+            return this
+        }
+
+        override fun setSignatureChange(): FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+            return this
+        }
+
+        override fun setPreserveSourceElement(): FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+            return this
+        }
+
+        override fun setDropOriginalInContainingParts(): FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+            return this
+        }
+
+        override fun setHiddenToOvercomeSignatureClash(): FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+            return this
+        }
+
+        override fun setHiddenForResolutionEverywhereBesideSupercalls(): FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+            return this
+        }
+
+        override fun setAdditionalAnnotations(additionalAnnotations: Annotations): FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+            return this
+        }
+
+        override fun setSubstitution(substitution: TypeSubstitution): FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+            return this
+        }
+
+        override fun <V : Any?> putUserData(
+            userDataKey: CallableDescriptor.UserDataKey<V>,
+            value: V
+        ): FunctionDescriptor.CopyBuilder<WrappedSimpleFunctionDescriptor> {
+            this@WrappedSimpleFunctionDescriptor.userDataMap[userDataKey] = value
+            return this
+        }
+
+        override fun build(): WrappedSimpleFunctionDescriptor? {
+            return this@WrappedSimpleFunctionDescriptor
+        }
+    }
 }
 
 class WrappedFunctionDescriptorWithContainerSource(

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

@@ -167,7 +167,7 @@ fun IrDeclarationContainer.addFunction(
 fun IrDeclarationContainer.addFunctionOverride(function: IrSimpleFunction, modality: Modality = Modality.FINAL): IrSimpleFunction =
     addFunction(function.name.asString(), function.returnType, modality).apply {
         overriddenSymbols.add(function.symbol)
-        valueParameters.addAll(function.valueParameters.map { it.copyTo(this) })
+        function.valueParameters.mapTo(valueParameters) { it.copyTo(this) }
     }
 
 inline fun buildConstructor(b: IrFunctionBuilder.() -> Unit): IrConstructor =

compiler/ir/backend.common/src/org/jetbrains/kotlin/ir/util/IrTypeUtils.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2018 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
 import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.utils.DFS
 
 val kotlinPackageFqn = FqName.fromSegments(listOf("kotlin"))
@@ -31,6 +32,10 @@ fun IrType.isKClass() = this.isNameInPackage("KClass", kotlinReflectionPackageFq
 fun IrType.isKFunction() = this.isNameInPackage("KFunction", kotlinReflectionPackageFqn)
 fun IrType.isSuspendFunction() = this.isNameInPackage("SuspendFunction", kotlinCoroutinesPackageFqn)
 
+fun IrType.isKotlinResult(): Boolean = isNameInPackage("Result", kotlinPackageFqn)
+fun IrType.isContinuation(): Boolean = isNameInPackage("Continuation", kotlinCoroutinesPackageFqn)
+fun IrType.isNullableContinuation(): Boolean = isContinuation() && this is IrSimpleType && hasQuestionMark
+
 fun IrType.isNameInPackage(prefix: String, packageFqName: FqName): Boolean {
     val classifier = classifierOrNull ?: return false
     val name = classifier.descriptor.name.asString()

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmBackend.kt

@@ -23,15 +23,10 @@ class JvmBackend(val context: JvmBackendContext) {
     }
 
     fun generateLoweredFile(irFile: IrFile) {
-        for (loweredClass in context.suspendFunctionContinuations.values) {
-            codegen.generateClass(loweredClass)
-        }
-
         for (loweredClass in irFile.declarations) {
             if (loweredClass !is IrClass) {
                 throw AssertionError("File-level declaration should be IrClass after JvmLower, got: " + loweredClass.render())
             }
-            if (loweredClass in context.suspendFunctionContinuations.values) continue
 
             codegen.generateClass(loweredClass)
         }

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmBackendContext.kt

@@ -53,9 +53,8 @@ class JvmBackendContext(
     override val internalPackageFqn = FqName("kotlin.jvm")
 
     val suspendFunctionContinuations = mutableMapOf<IrFunction, IrClass>()
-    val suspendLambdaClasses = mutableMapOf<IrClass, KtElement>()
+    val suspendLambdaToOriginalFunctionMap = mutableMapOf<IrClass, IrFunction>()
     val continuationClassBuilders = mutableMapOf<IrClass, ClassBuilder>()
-    var functionReferenceAndContinuationsCount = 0
 
     internal fun getTopLevelClass(fqName: FqName): IrClassSymbol {
         val descriptor = state.module.getPackage(fqName.parent()).memberScope.getContributedClassifier(

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ClassCodegen.kt

@@ -95,16 +95,23 @@ open class ClassCodegen protected constructor(
             signature.interfaces.toTypedArray()
         )
         AnnotationCodegen(this, context.state, visitor.visitor::visitAnnotation).genAnnotations(irClass, null)
-        /* TODO: Temporary workaround: ClassBuilder needs a pathless name. */
-        val shortName = File(fileEntry.name).name
-        visitor.visitSource(shortName, null)
 
         val nestedClasses = irClass.declarations.mapNotNull { declaration ->
-            if (declaration is IrClass && declaration !in context.suspendFunctionContinuations.values) {
+            if (declaration is IrClass) {
                 ClassCodegen(declaration, context, this)
             } else null
         }
 
+        // Suspend function state-machine builder requires half-built continuation class
+        val continuationCodegens = nestedClasses.filter { it.irClass in context.suspendFunctionContinuations.values }
+        for (continuationCodegen in continuationCodegens) {
+            continuationCodegen.generate()
+        }
+
+        /* TODO: Temporary workaround: ClassBuilder needs a pathless name. */
+        val shortName = File(fileEntry.name).name
+        visitor.visitSource(shortName, null)
+
         val companionObjectCodegen = nestedClasses.firstOrNull { it.irClass.isCompanion }
 
         for (declaration in irClass.declarations) {
@@ -113,7 +120,7 @@ open class ClassCodegen protected constructor(
 
         // Generate nested classes at the end, to ensure that codegen for companion object will have the necessary JVM signatures in its
         // trace for properties moved to the outer class
-        for (codegen in nestedClasses) {
+        for (codegen in (nestedClasses - continuationCodegens)) {
             codegen.generate()
         }
 

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/CoroutineCodegen.kt

@@ -14,7 +14,6 @@ import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.util.isSuspend
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodGenericSignature
-import org.jetbrains.kotlin.utils.sure
 import org.jetbrains.org.objectweb.asm.MethodVisitor
 
 fun generateStateMachineForNamedFunction(
@@ -32,11 +31,7 @@ fun generateStateMachineForNamedFunction(
     assert(languageVersionSettings.isReleaseCoroutines()) { "Experimental coroutines are unsupported in JVM_IR backend" }
     return CoroutineTransformerMethodVisitor(
         methodVisitor, access, signature.asmMethod.name, signature.asmMethod.descriptor, null, null,
-        obtainClassBuilderForCoroutineState = {
-            continuationClassBuilder.sure {
-                "continuationClassBuilder is null"
-            }
-        },
+        obtainClassBuilderForCoroutineState = { continuationClassBuilder!! },
         element = element,
         diagnostics = state.diagnostics,
         languageVersionSettings = languageVersionSettings,
@@ -72,4 +67,4 @@ fun generateStateMachineForLambda(
 }
 
 fun IrFunction.isInvokeSuspendOfLambda(context: JvmBackendContext): Boolean =
-    name.asString() == INVOKE_SUSPEND_METHOD_NAME && parent in context.suspendLambdaClasses
+    name.asString() == INVOKE_SUSPEND_METHOD_NAME && parent in context.suspendLambdaToOriginalFunctionMap

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ExpressionCodegen.kt

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.backend.jvm.lower.constantValue
 import org.jetbrains.kotlin.backend.jvm.lower.inlineclasses.unboxInlineClass
 import org.jetbrains.kotlin.codegen.*
 import org.jetbrains.kotlin.codegen.AsmUtil.*
+import org.jetbrains.kotlin.codegen.coroutines.INVOKE_SUSPEND_METHOD_NAME
 import org.jetbrains.kotlin.codegen.inline.*
 import org.jetbrains.kotlin.codegen.inline.ReifiedTypeInliner.OperationKind.AS
 import org.jetbrains.kotlin.codegen.inline.ReifiedTypeInliner.OperationKind.SAFE_AS
@@ -314,7 +315,7 @@ class ExpressionCodegen(
                 mv.load(0, OBJECT_TYPE)
             expression.descriptor is ConstructorDescriptor ->
                 throw AssertionError("IrCall with ConstructorDescriptor: ${expression.javaClass.simpleName}")
-            callee.isSuspend -> addInlineMarker(mv, isStartNotEnd = true)
+            callee.isSuspend && !irFunction.isInvokeSuspendInContinuation() -> addInlineMarker(mv, isStartNotEnd = true)
         }
 
         val receiver = expression.dispatchReceiver
@@ -389,7 +390,8 @@ class ExpressionCodegen(
 
         expression.markLineNumber(true)
 
-        if (callee.isSuspend) {
+        // Do not generate redundant markers in continuation, otherwise Roman will demand to remove them
+        if (callee.isSuspend && !irFunction.isInvokeSuspendInContinuation()) {
             addSuspendMarker(mv, isStartNotEnd = true)
         }
 
@@ -402,7 +404,7 @@ class ExpressionCodegen(
 
         val returnType = callee.returnType.substitute(typeSubstitutionMap)
 
-        if (callee.isSuspend) {
+        if (callee.isSuspend && !irFunction.isInvokeSuspendInContinuation()) {
             addSuspendMarker(mv, isStartNotEnd = false)
             addInlineMarker(mv, isStartNotEnd = false)
         }
@@ -425,6 +427,9 @@ class ExpressionCodegen(
         }
     }
 
+    private fun IrFunction.isInvokeSuspendInContinuation(): Boolean =
+        name.asString() == INVOKE_SUSPEND_METHOD_NAME && parentAsClass in classCodegen.context.suspendFunctionContinuations.values
+
     override fun visitVariable(declaration: IrVariable, data: BlockInfo): PromisedValue {
         val varType = typeMapper.mapType(declaration)
         val index = frameMap.enter(declaration.symbol, varType)

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/FunctionCodegen.kt

@@ -65,8 +65,8 @@ open class FunctionCodegen(
         } else {
             val frameMap = createFrameMapWithReceivers(signature)
             val irClass = classCodegen.context.suspendFunctionContinuations[irFunction]
-            val element = irFunction.symbol.descriptor.psiElement as? KtElement
-                ?: classCodegen.context.suspendLambdaClasses[classCodegen.context.suspendLambdaClasses.keys.find { it == irFunction.parent }]
+            val element = (irFunction.symbol.descriptor.psiElement
+                ?: classCodegen.context.suspendLambdaToOriginalFunctionMap[irFunction.parent]?.symbol?.descriptor?.psiElement) as? KtElement
             val continuationClassBuilder = classCodegen.context.continuationClassBuilders[irClass]
             methodVisitor = when {
                 irFunction.isSuspend -> generateStateMachineForNamedFunction(
@@ -206,7 +206,9 @@ fun generateParameterAnnotations(
     state: GenerationState
 ) {
     val iterator = irFunction.valueParameters.iterator()
-    val kotlinParameterTypes = jvmSignature.valueParameters
+    val kotlinParameterTypes =
+        if (irFunction.isSuspend) jvmSignature.valueParameters.dropLast(1) // do not generate annotation for continuation parameter
+        else jvmSignature.valueParameters
     var syntheticParameterCount = 0
     kotlinParameterTypes.forEachIndexed { i, parameterSignature ->
         val kind = parameterSignature.kind

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/AddContinuationLowering.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.backend.jvm.lower
 
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.descriptors.WrappedSimpleFunctionDescriptor
 import org.jetbrains.kotlin.backend.common.descriptors.synthesizedName
 import org.jetbrains.kotlin.backend.common.ir.copyTo
 import org.jetbrains.kotlin.backend.common.ir.copyTypeParametersFrom
@@ -13,9 +14,11 @@ import org.jetbrains.kotlin.backend.common.ir.createImplicitParameterDeclaration
 import org.jetbrains.kotlin.backend.common.ir.isSuspend
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.common.lower.parents
+import org.jetbrains.kotlin.backend.common.peek
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
+import org.jetbrains.kotlin.backend.common.pop
+import org.jetbrains.kotlin.backend.common.push
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
-import org.jetbrains.kotlin.backend.jvm.codegen.psiElement
 import org.jetbrains.kotlin.codegen.coroutines.*
 import org.jetbrains.kotlin.config.coroutinesPackageFqName
 import org.jetbrains.kotlin.ir.IrElement
@@ -28,19 +31,20 @@ import org.jetbrains.kotlin.ir.declarations.impl.IrFunctionImpl
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.*
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
+import org.jetbrains.kotlin.ir.symbols.impl.IrSimpleFunctionSymbolImpl
 import org.jetbrains.kotlin.ir.types.IrSimpleType
 import org.jetbrains.kotlin.ir.types.createType
 import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
 import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
-import org.jetbrains.kotlin.ir.types.impl.originalKotlinType
+import org.jetbrains.kotlin.ir.types.isInt
+import org.jetbrains.kotlin.ir.types.isNullableAny
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.*
 import org.jetbrains.kotlin.load.java.JavaVisibilities
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.util.OperatorNameConventions
-import java.util.*
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
 
 // TODO: Split into several lowerings, merge with ones in JS and Native
 internal val addContinuationPhase = makeIrFilePhase(
@@ -66,8 +70,10 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
 
     override fun lower(irFile: IrFile) {
         val suspendLambdas = markSuspendLambdas(irFile)
-        val suspendFunctions = addContinuationParameter(irFile, suspendLambdas)
-        suspendLambdas.forEach { generateContinuationClassForLambda(it) }
+        val suspendFunctions = addContinuationParameter(irFile, suspendLambdas.map { it.function }.toSet())
+        for (lambda in suspendLambdas) {
+            generateContinuationClassForLambda(lambda)
+        }
         transformReferencesToSuspendLambdas(irFile, suspendLambdas)
         transformSuspendCalls(irFile)
         for (suspendFunction in suspendFunctions) {
@@ -77,15 +83,6 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
 
     private fun transformReferencesToSuspendLambdas(irFile: IrFile, suspendLambdas: List<SuspendLambdaInfo>) {
         irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
-            private val functionStack = Stack<IrFunction>()
-
-            override fun visitFunction(declaration: IrFunction): IrStatement {
-                functionStack.push(declaration)
-                val res = super.visitFunction(declaration)
-                functionStack.pop()
-                return res
-            }
-
             override fun visitFunctionReference(expression: IrFunctionReference): IrExpression {
                 expression.transformChildrenVoid(this)
 
@@ -129,6 +126,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             val constructor = addPrimaryConstructorForLambda(info.arity, info.function, parametersFields)
             val secondaryConstructor = addSecondaryConstructorForLambda(constructor)
             val invokeToOverride = functionNClass.functions.single { it.owner.valueParameters.size == info.arity + 1 }
+            assert(invokeToOverride.owner.name.asString() == "invoke")
             val invokeSuspend = addInvokeSuspendForLambda(info.function, parametersFields)
             if (info.arity <= 1) {
                 val create = addCreate(constructor, suspendLambda, info.arity, parametersFields)
@@ -137,14 +135,16 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                 addInvoke(constructor, invokeSuspend, invokeToOverride)
             }
 
-            context.suspendLambdaClasses[this] = info.function.symbol.descriptor.psiElement as KtElement
+            context.suspendLambdaToOriginalFunctionMap[this] = info.function
 
             info.constructor = secondaryConstructor
         }
     }
 
     private fun IrClass.addInvokeSuspendForLambda(irFunction: IrFunction, fields: List<IrField>): IrFunction {
-        return addFunctionOverride(suspendLambda.functions.single { it.owner.name.asString() == INVOKE_SUSPEND_METHOD_NAME }.owner)
+        val superMethod = suspendLambda.functions.single { it.owner.name.asString() == INVOKE_SUSPEND_METHOD_NAME }.owner
+        assert(superMethod.valueParameters.size == 1 && superMethod.valueParameters[0].type.isKotlinResult())
+        return addFunctionOverride(superMethod)
             .also { function ->
                 function.body = irFunction.body?.deepCopyWithSymbols()
                 function.body?.transformChildrenVoid(object : IrElementTransformerVoid() {
@@ -193,6 +193,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         parametersFields: List<IrField>
     ): IrFunction {
         val create = superType.functions.single { it.name == Name.identifier("create") && it.valueParameters.size == arity + 1 }
+        assert(create.valueParameters.last().type.isContinuation() && if (arity == 1) create.valueParameters.first().type.isNullableAny() else true)
         return addFunctionOverride(create).also { function ->
             function.body = context.createIrBuilder(function.symbol).irBlockBody {
                 +irReturn(irCall(constructor).also {
@@ -208,14 +209,14 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
     }
 
     private fun IrClass.createContinuationClassFor(irFunction: IrFunction): IrClass = buildClass {
-        name = "${irFunction.name}\$${context.functionReferenceAndContinuationsCount++}".synthesizedName
+        name = "${irFunction.name}\$Continuation".synthesizedName
         origin = CONTINUATION_CLASS
         visibility = JavaVisibilities.PACKAGE_VISIBILITY
     }.also { irClass ->
         irClass.createImplicitParameterDeclarationWithWrappedDescriptor()
         irClass.superTypes.add(defaultType)
 
-        val parent = irFunction.parents.find { it is IrDeclarationContainer } as IrDeclarationContainer
+        val parent = irFunction.parents.firstIsInstance<IrDeclarationContainer>()
         irClass.parent = parent
         parent.declarations.add(irClass)
     }
@@ -229,10 +230,11 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             isPrimary = true
             returnType = defaultType
         }.also { constructor ->
-            constructor.valueParameters.addAll(irFunction.valueParameters.map { it.copyTo(constructor) })
+            irFunction.valueParameters.mapTo(constructor.valueParameters) { it.copyTo(constructor) }
             val completionParameterSymbol = constructor.addCompletionValueParameter()
 
             val superClassConstructor = suspendLambda.owner.constructors.single { it.valueParameters.size == 2 }
+            assert(superClassConstructor.valueParameters[0].type.isInt() && superClassConstructor.valueParameters[1].type.isNullableContinuation())
             constructor.body = context.createIrBuilder(constructor.symbol).irBlockBody {
                 +irDelegatingConstructorCall(superClassConstructor).also {
                     it.putValueArgument(0, irInt(arity + 1))
@@ -265,9 +267,6 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
     private fun IrFunction.addCompletionValueParameter(): IrValueParameter =
         addValueParameter(SUSPEND_FUNCTION_CONTINUATION_PARAMETER, continuationType())
 
-    private fun IrFunction.addObjectParameter(name: String): IrValueParameter =
-        addValueParameter(name, context.irBuiltIns.anyNType)
-
     private fun IrFunction.continuationType(): IrSimpleType =
         continuation.createType(true, listOf(makeTypeProjection(returnType, Variance.INVARIANT)))
 
@@ -319,15 +318,14 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         }
     }
 
-    private fun addContinuationParameter(irFile: IrFile, suspendLambdas: List<SuspendLambdaInfo>): Set<IrFunction> {
+    private fun addContinuationParameter(irFile: IrFile, suspendLambdas: Set<IrFunction>): Set<IrFunction> {
         val views = hashSetOf<IrFunction>()
         fun tryAddContinuationParameter(element: IrElement): List<IrDeclaration>? {
-            return if (element is IrSimpleFunction && element.isSuspend && suspendLambdas.none { it.function == element }) {
+            return if (element is IrSimpleFunction && element.isSuspend && element !in suspendLambdas) {
                 listOf(element.getOrCreateView().also { views.add(it) })
             } else null
         }
 
-        irFile.transformDeclarationsFlat(::tryAddContinuationParameter)
         irFile.acceptVoid(object : IrElementVisitorVoid {
             override fun visitElement(element: IrElement) {
                 element.acceptChildrenVoid(this)
@@ -352,7 +350,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                 expression.acceptChildrenVoid(this)
 
                 if (expression.isSuspend) {
-                    suspendLambdas += SuspendLambdaInfo(expression.symbol.owner, expression.type.originalKotlinType!!.arguments.size - 1)
+                    suspendLambdas += SuspendLambdaInfo(expression.symbol.owner, (expression.type as IrSimpleType).arguments.size - 1)
                 }
             }
         })
@@ -361,7 +359,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
 
     private fun transformSuspendCalls(irFile: IrFile) {
         irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
-            private val functionStack = Stack<IrFunction>()
+            private val functionStack = mutableListOf<IrFunction>()
             override fun visitElement(element: IrElement): IrElement {
                 element.transformChildrenVoid(this)
                 return element
@@ -375,10 +373,10 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             }
 
             private fun getContinuationFromCaller(): IrGetValue {
-                val caller = functionStack.peek()
+                val caller = functionStack.peek()!!
                 val continuationParameter =
                     if (caller.isSuspend) caller.valueParameters.last().symbol
-                    else if (caller.name.asString() == INVOKE_SUSPEND_METHOD_NAME && caller.parent in context.suspendLambdaClasses)
+                    else if (caller.name.asString() == INVOKE_SUSPEND_METHOD_NAME && caller.parent in context.suspendLambdaToOriginalFunctionMap)
                         caller.dispatchReceiverParameter!!.symbol
                     else error("suspend call outside suspend context")
                 return IrGetValueImpl(
@@ -427,20 +425,25 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
 
     private fun IrSimpleFunction.getOrCreateView(): IrSimpleFunction =
         transformedSuspendFunctionsCache.getOrPut(this) {
+            // Copy source element, so we can check for suspend calls in monitor during state-machine generation
+            val descriptor = WrappedSimpleFunctionDescriptor(sourceElement = this.descriptor.source)
             IrFunctionImpl(
-                startOffset, endOffset, origin, getOrCreateJvmSuspendFunctionView(descriptor, context.state), context.irBuiltIns.anyType
+                startOffset, endOffset, origin, IrSimpleFunctionSymbolImpl(descriptor),
+                name, visibility, modality, context.irBuiltIns.anyType,
+                isInline, isExternal, isTailrec, isSuspend
             ).also {
+                descriptor.bind(it)
                 it.parent = parent
                 it.copyTypeParametersFrom(this)
 
                 it.dispatchReceiverParameter = dispatchReceiverParameter?.copyTo(it)
+                it.extensionReceiverParameter = extensionReceiverParameter?.copyTo(it)
 
                 valueParameters.mapTo(it.valueParameters) { p -> p.copyTo(it) }
                 it.addCompletionValueParameter()
 
                 // Fix links to parameters
-                val valueParametersMapping: Map<IrValueDeclaration?, IrValueParameter?> = (valueParameters + dispatchReceiverParameter)
-                    .zip(it.valueParameters.dropLast(1) + it.dispatchReceiverParameter).toMap()
+                val valueParametersMapping = explicitParameters.zip(it.explicitParameters).toMap()
                 it.body = body?.deepCopyWithSymbols(this)
                 it.body?.transformChildrenVoid(object : IrElementTransformerVoid() {
                     override fun visitGetValue(expression: IrGetValue) =