Generate continuation class for suspend lambdas

generate state-machine for coroutines.

compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/CoroutineTransformerMethodVisitor.kt

@@ -69,7 +69,9 @@ class CoroutineTransformerMethodVisitor(
     // It's only matters for named functions, may differ from '!isStatic(access)' in case of DefaultImpls
     private val needDispatchReceiver: Boolean = false,
     // May differ from containingClassInternalName in case of DefaultImpls
-    private val internalNameForDispatchReceiver: String? = null
+    private val internalNameForDispatchReceiver: String? = null,
+    // JVM_IR backend generates $completion, while old backend does not
+    private val putContinuationParameterToLvt: Boolean = true
 ) : TransformationMethodVisitor(delegate, access, name, desc, signature, exceptions) {
 
     private val classBuilderForCoroutineState: ClassBuilder by lazy(obtainClassBuilderForCoroutineState)
@@ -105,7 +107,9 @@ class CoroutineTransformerMethodVisitor(
         if (isForNamedFunction) {
             ReturnUnitMethodTransformer.transform(containingClassInternalName, methodNode)
 
-            addCompletionParameterToLVT(methodNode)
+            if (putContinuationParameterToLvt) {
+                addCompletionParameterToLVT(methodNode)
+            }
 
             if (allSuspensionPointsAreTailCalls(containingClassInternalName, methodNode, suspensionPoints)) {
                 dropSuspensionMarkers(methodNode, suspensionPoints)

compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/coroutineCodegenUtil.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.
  */
 

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/inlineCodegenUtils.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-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.
  */
 
@@ -395,7 +395,7 @@ internal fun removeFinallyMarkers(intoNode: MethodNode) {
     }
 }
 
-internal fun addInlineMarker(v: InstructionAdapter, isStartNotEnd: Boolean) {
+fun addInlineMarker(v: InstructionAdapter, isStartNotEnd: Boolean) {
     v.visitMethodInsn(
         Opcodes.INVOKESTATIC, INLINE_MARKER_CLASS_NAME,
         if (isStartNotEnd) INLINE_MARKER_BEFORE_METHOD_NAME else INLINE_MARKER_AFTER_METHOD_NAME,
@@ -423,7 +423,7 @@ internal fun addReturnsUnitMarkerIfNecessary(v: InstructionAdapter, resolvedCall
     }
 }
 
-internal fun addSuspendMarker(v: InstructionAdapter, isStartNotEnd: Boolean) {
+fun addSuspendMarker(v: InstructionAdapter, isStartNotEnd: Boolean) {
     v.emitInlineMarker(if (isStartNotEnd) INLINE_MARKER_BEFORE_SUSPEND_ID else INLINE_MARKER_AFTER_SUSPEND_ID)
 }
 

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

@@ -8,11 +8,14 @@ package org.jetbrains.kotlin.ir.builders.declarations
 import org.jetbrains.kotlin.backend.common.descriptors.*
 import org.jetbrains.kotlin.backend.common.ir.copyTo
 import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.descriptors.Visibility
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.declarations.impl.*
 import org.jetbrains.kotlin.ir.symbols.impl.*
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.util.defaultType
+import org.jetbrains.kotlin.ir.util.overrides
 import org.jetbrains.kotlin.ir.util.parentAsClass
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.types.Variance
@@ -58,6 +61,13 @@ inline fun IrDeclarationContainer.addField(b: IrFieldBuilder.() -> Unit) =
         declarations.add(field)
     }
 
+fun IrClass.addField(fieldName: String, fieldType: IrType, fieldVisibility: Visibility = Visibilities.PRIVATE): IrField =
+    addField {
+        name = Name.identifier(fieldName)
+        type = fieldType
+        visibility = fieldVisibility
+    }
+
 fun IrPropertyBuilder.buildProperty(): IrProperty {
     val wrappedDescriptor = WrappedPropertyDescriptor()
     return IrPropertyImpl(
@@ -154,6 +164,12 @@ 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) })
+    }
+
 inline fun buildConstructor(b: IrFunctionBuilder.() -> Unit): IrConstructor =
     IrFunctionBuilder().run {
         b()

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

@@ -1,17 +1,6 @@
 /*
- * Copyright 2010-2016 JetBrains s.r.o.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
+ * Copyright 2010-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.
  */
 
 package org.jetbrains.kotlin.backend.jvm
@@ -34,10 +23,15 @@ 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

@@ -11,21 +11,21 @@ import org.jetbrains.kotlin.backend.common.phaser.PhaseConfig
 import org.jetbrains.kotlin.backend.jvm.descriptors.JvmDeclarationFactory
 import org.jetbrains.kotlin.backend.jvm.descriptors.JvmSharedVariablesManager
 import org.jetbrains.kotlin.backend.jvm.intrinsics.IrIntrinsicMethods
+import org.jetbrains.kotlin.codegen.ClassBuilder
 import org.jetbrains.kotlin.codegen.coroutines.coroutinesJvmInternalPackageFqName
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.config.coroutinesPackageFqName
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.ir.IrElement
-import org.jetbrains.kotlin.ir.declarations.IrFile
-import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
-import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
+import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.util.ReferenceSymbolTable
 import org.jetbrains.kotlin.ir.util.SymbolTable
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi2ir.PsiSourceManager
 
 class JvmBackendContext(
@@ -52,7 +52,12 @@ class JvmBackendContext(
 
     override val internalPackageFqn = FqName("kotlin.jvm")
 
+    // TODO: Can I make it local to AddContinuationLowering?
     val transformedSuspendFunctionsCache = mutableMapOf<IrSimpleFunction, IrSimpleFunction>()
+    val suspendFunctionContinuations = mutableMapOf<IrFunction, IrClass>()
+    val suspendLambdaClasses = mutableMapOf<IrClass, KtElement>()
+    val continuationClassBuilders = mutableMapOf<IrClass, ClassBuilder>()
+    var functionReferenceAndContinuationsCount = 0
 
     private val coroutinePackage = state.module.getPackage(state.languageVersionSettings.coroutinesPackageFqName())
     private val coroutinesJvmInternalPackage = state.module.getPackage(state.languageVersionSettings.coroutinesJvmInternalPackageFqName())
@@ -69,6 +74,12 @@ class JvmBackendContext(
         ) as ClassDescriptor
     )
 
+    val suspendLambda = symbolTable.referenceClass(
+        coroutinesJvmInternalPackage.memberScope.getContributedClassifier(
+            Name.identifier("SuspendLambda"), NoLookupLocation.FROM_BACKEND
+        ) as ClassDescriptor
+    )
+
     internal fun getTopLevelClass(fqName: FqName): IrClassSymbol {
         val descriptor = state.module.getPackage(fqName.parent()).memberScope.getContributedClassifier(
             fqName.shortName(), NoLookupLocation.FROM_BACKEND

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

@@ -115,10 +115,10 @@ val jvmPhases = namedIrFilePhase<JvmBackendContext>(
             makePatchParentsPhase(1) then
 
             enumWhenPhase then
-            singletonReferencesPhase then
             localDeclarationsPhase then
             singleAbstractMethodPhase then
             addContinuationPhase then
+            singletonReferencesPhase then
             callableReferencePhase then
             functionNVarargInvokePhase then
 

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

@@ -1,17 +1,6 @@
 /*
- * Copyright 2010-2016 JetBrains s.r.o.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
+ * Copyright 2010-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.
  */
 
 package org.jetbrains.kotlin.backend.jvm.codegen
@@ -111,7 +100,7 @@ open class ClassCodegen protected constructor(
         visitor.visitSource(shortName, null)
 
         val nestedClasses = irClass.declarations.mapNotNull { declaration ->
-            if (declaration is IrClass) {
+            if (declaration is IrClass && declaration !in context.suspendFunctionContinuations.values) {
                 ClassCodegen(declaration, context, this)
             } else null
         }
@@ -130,7 +119,11 @@ open class ClassCodegen protected constructor(
 
         generateKotlinMetadataAnnotation()
 
-        done()
+        if (irClass in context.suspendFunctionContinuations.values) {
+            context.continuationClassBuilders[irClass] = visitor
+        } else {
+            done()
+        }
     }
 
     private fun generateKotlinMetadataAnnotation() {

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

@@ -5,33 +5,71 @@
 
 package org.jetbrains.kotlin.backend.jvm.codegen
 
+import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.codegen.ClassBuilder
+import org.jetbrains.kotlin.codegen.coroutines.CoroutineTransformerMethodVisitor
+import org.jetbrains.kotlin.codegen.coroutines.INVOKE_SUSPEND_METHOD_NAME
+import org.jetbrains.kotlin.config.isReleaseCoroutines
 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 wrapWithCoroutineTransformer(
+fun generateStateMachineForNamedFunction(
     irFunction: IrFunction,
     classCodegen: ClassCodegen,
     methodVisitor: MethodVisitor,
     access: Int,
-    signature: JvmMethodGenericSignature
+    signature: JvmMethodGenericSignature,
+    continuationClassBuilder: ClassBuilder?,
+    element: KtElement
 ): MethodVisitor {
-    return methodVisitor
-//    assert(irFunction.isSuspend)
-//    val state = classCodegen.state
-//    val languageVersionSettings = state.languageVersionSettings
-//    assert(languageVersionSettings.isReleaseCoroutines()) { "Experimental coroutines are unsupported in JVM_IR backend" }
-//    val continuationClass = createContinuationClassForNamedFunction()
-//    CoroutineTransformerMethodVisitor(
-//        methodVisitor, access, signature.asmMethod.name, signature.asmMethod.descriptor, null, null,
-//        obtainClassBuilderForCoroutineState = { TODO() },
-//        element = irFunction.symbol.descriptor.psiElement as KtElement,
-//        diagnostics = state.diagnostics,
-//        languageVersionSettings = languageVersionSettings,
-//        shouldPreserveClassInitialization = state.constructorCallNormalizationMode.shouldPreserveClassInitialization,
-//        containingClassInternalName = classCodegen.visitor.thisName,
-//        isForNamedFunction = true,
-//        needDispatchReceiver = true,
-//        internalNameForDispatchReceiver = classCodegen.visitor.thisName
-//    )
+    assert(irFunction.isSuspend)
+    val state = classCodegen.state
+    val languageVersionSettings = state.languageVersionSettings
+    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"
+            }
+        },
+        element = element,
+        diagnostics = state.diagnostics,
+        languageVersionSettings = languageVersionSettings,
+        shouldPreserveClassInitialization = state.constructorCallNormalizationMode.shouldPreserveClassInitialization,
+        containingClassInternalName = classCodegen.visitor.thisName,
+        isForNamedFunction = true,
+        needDispatchReceiver = irFunction.dispatchReceiverParameter != null,
+        internalNameForDispatchReceiver = classCodegen.visitor.thisName,
+        putContinuationParameterToLvt = false
+    )
 }
+
+fun generateStateMachineForLambda(
+    classCodegen: ClassCodegen,
+    methodVisitor: MethodVisitor,
+    access: Int,
+    signature: JvmMethodGenericSignature,
+    element: KtElement
+): MethodVisitor {
+    val state = classCodegen.state
+    val languageVersionSettings = state.languageVersionSettings
+    assert(languageVersionSettings.isReleaseCoroutines()) { "Experimental coroutines are unsupported in JVM_IR backend" }
+    return CoroutineTransformerMethodVisitor(
+        methodVisitor, access, signature.asmMethod.name, signature.asmMethod.descriptor, null, null,
+        obtainClassBuilderForCoroutineState = { classCodegen.visitor },
+        element = element,
+        diagnostics = state.diagnostics,
+        languageVersionSettings = languageVersionSettings,
+        shouldPreserveClassInitialization = state.constructorCallNormalizationMode.shouldPreserveClassInitialization,
+        containingClassInternalName = classCodegen.visitor.thisName,
+        isForNamedFunction = false
+    )
+}
+
+fun IrFunction.isInvokeSuspendOfLambda(context: JvmBackendContext): Boolean =
+    name.asString() == INVOKE_SUSPEND_METHOD_NAME && parent in context.suspendLambdaClasses

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

@@ -314,6 +314,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)
         }
 
         val receiver = expression.dispatchReceiver
@@ -387,6 +388,11 @@ class ExpressionCodegen(
         }
 
         expression.markLineNumber(true)
+
+        if (callee.isSuspend) {
+            addSuspendMarker(mv, isStartNotEnd = true)
+        }
+
         callGenerator.genCall(
             callable,
             defaultMask.generateOnStackIfNeeded(callGenerator, callee is IrConstructor, this),
@@ -394,7 +400,13 @@ class ExpressionCodegen(
             expression
         )
 
-        val returnType = callee.returnType
+        val returnType = callee.returnType.substitute(typeSubstitutionMap)
+
+        if (callee.isSuspend) {
+            addSuspendMarker(mv, isStartNotEnd = false)
+            addInlineMarker(mv, isStartNotEnd = false)
+        }
+
         return when {
             returnType.substitute(typeSubstitutionMap).isNothing() -> {
                 mv.aconst(null)

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

@@ -5,23 +5,18 @@
 
 package org.jetbrains.kotlin.backend.jvm.codegen
 
-import org.jetbrains.kotlin.backend.common.descriptors.WrappedValueParameterDescriptor
-import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns.FQ_NAMES
-import org.jetbrains.kotlin.codegen.*
+import org.jetbrains.kotlin.codegen.AsmUtil
+import org.jetbrains.kotlin.codegen.ClassBuilderMode
+import org.jetbrains.kotlin.codegen.OwnerKind
 import org.jetbrains.kotlin.codegen.state.GenerationState
+import org.jetbrains.kotlin.codegen.visitAnnotableParameterCount
 import org.jetbrains.kotlin.descriptors.Modality
-import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.declarations.impl.IrValueParameterImpl
 import org.jetbrains.kotlin.ir.expressions.*
-import org.jetbrains.kotlin.ir.symbols.impl.IrValueParameterSymbolImpl
-import org.jetbrains.kotlin.ir.types.IrType
-import org.jetbrains.kotlin.ir.types.createType
-import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.resolve.jvm.annotations.JVM_SYNTHETIC_ANNOTATION_FQ_NAME
 import org.jetbrains.kotlin.resolve.jvm.annotations.STRICTFP_ANNOTATION_FQ_NAME
@@ -29,7 +24,6 @@ import org.jetbrains.kotlin.resolve.jvm.annotations.SYNCHRONIZED_ANNOTATION_FQ_N
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodGenericSignature
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodParameterKind
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature
-import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.org.objectweb.asm.MethodVisitor
 import org.jetbrains.org.objectweb.asm.Opcodes
@@ -70,17 +64,39 @@ open class FunctionCodegen(
             generateAnnotationDefaultValueIfNeeded(methodVisitor)
         } 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 continuationClassBuilder = classCodegen.context.continuationClassBuilders[irClass]
             ExpressionCodegen(
                 irFunction,
                 frameMap,
                 InstructionAdapter(
-                    if (irFunction.isSuspend) wrapWithCoroutineTransformer(irFunction, classCodegen, methodVisitor, flags, signature)
-                    else methodVisitor
+                    when {
+                        irFunction.isSuspend -> generateStateMachineForNamedFunction(
+                            irFunction,
+                            classCodegen,
+                            methodVisitor,
+                            flags,
+                            signature,
+                            continuationClassBuilder,
+                            element!!
+                        )
+                        irFunction.isInvokeSuspendOfLambda(classCodegen.context) -> generateStateMachineForLambda(
+                            classCodegen,
+                            methodVisitor,
+                            flags,
+                            signature,
+                            element!!
+                        )
+                        else -> methodVisitor
+                    }
                 ),
                 classCodegen,
                 isInlineLambda
             ).generate()
             methodVisitor.visitMaxs(-1, -1)
+            continuationClassBuilder?.done()
         }
         methodVisitor.visitEnd()
 

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

@@ -6,188 +6,309 @@
 package org.jetbrains.kotlin.backend.jvm.lower
 
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
-import org.jetbrains.kotlin.backend.common.descriptors.WrappedValueParameterDescriptor
 import org.jetbrains.kotlin.backend.common.descriptors.synthesizedName
 import org.jetbrains.kotlin.backend.common.ir.copyTo
 import org.jetbrains.kotlin.backend.common.ir.copyTypeParametersFrom
 import org.jetbrains.kotlin.backend.common.ir.createImplicitParameterDeclarationWithWrappedDescriptor
 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.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
-import org.jetbrains.kotlin.codegen.coroutines.COROUTINE_LABEL_FIELD_NAME
-import org.jetbrains.kotlin.codegen.coroutines.INVOKE_SUSPEND_METHOD_NAME
-import org.jetbrains.kotlin.codegen.coroutines.dataFieldName
-import org.jetbrains.kotlin.codegen.coroutines.getOrCreateJvmSuspendFunctionView
-import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.backend.jvm.codegen.psiElement
+import org.jetbrains.kotlin.codegen.coroutines.*
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.builders.*
-import org.jetbrains.kotlin.ir.builders.declarations.buildClass
-import org.jetbrains.kotlin.ir.builders.declarations.buildConstructor
-import org.jetbrains.kotlin.ir.builders.declarations.buildField
-import org.jetbrains.kotlin.ir.builders.declarations.buildFun
+import org.jetbrains.kotlin.ir.builders.declarations.*
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.declarations.impl.IrFunctionImpl
-import org.jetbrains.kotlin.ir.declarations.impl.IrValueParameterImpl
 import org.jetbrains.kotlin.ir.expressions.*
-import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
-import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionReferenceImpl
-import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
-import org.jetbrains.kotlin.ir.symbols.impl.IrValueParameterSymbolImpl
-import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.expressions.impl.*
+import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
+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.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.*
 
-// TODO: Move to common prefix
+// TODO: Split into several lowerings, merge with ones in JS and Native
 internal val addContinuationPhase = makeIrFilePhase(
     ::AddContinuationLowering,
     "AddContinuation",
     "Add continuation parameter to suspend functions and suspend calls"
 )
 
-private object CONTINUATION_PARAMETER : IrDeclarationOriginImpl("CONTINUATION_PARAMETER")
 private object CONTINUATION_CLASS : IrDeclarationOriginImpl("CONTINUATION_CLASS")
-private object CONTINUATION_CLASS_CONSTRUCTOR : IrDeclarationOriginImpl("CONTINUATION_CLASS_CONSTRUCTOR")
-private object CONTINUATION_CLASS_COMPLETION_PARAMETER : IrDeclarationOriginImpl("CONTINUATION_CLASS_COMPLETION_PARAMETER")
-private object CONTINUATION_CLASS_INVOKE_SUSPEND : IrDeclarationOriginImpl("CONTINUATION_CLASS_INVOKE_SUSPEND")
-private object CONTINUATION_CLASS_RESULT_FIELD : IrDeclarationOriginImpl("CONTINUATION_CLASS_RESULT_FIELD")
-private object CONTINUATION_CLASS_LABEL_FIELD : IrDeclarationOriginImpl("CONTINUATION_CLASS_LABEL_FIELD")
 
 private class AddContinuationLowering(private val context: JvmBackendContext) : FileLoweringPass {
     override fun lower(irFile: IrFile) {
         val suspendLambdas = markSuspendLambdas(irFile)
         val suspendFunctions = addContinuationParameter(irFile, suspendLambdas)
-        transformSuspendCalls(irFile, suspendLambdas)
-        suspendFunctions.forEach {
-            generateContinuationClass(it)
+        suspendLambdas.forEach { generateContinuationClassForLambda(it) }
+        transformReferencesToSuspendLambdas(irFile, suspendLambdas)
+        transformSuspendCalls(irFile)
+        for (suspendFunction in suspendFunctions) {
+            generateContinuationClassForNamedFunction(suspendFunction)
         }
-        // TODO: suspend lambdas are covered by CallableReferenceLowering, which is not ideal
     }
 
-    // TODO: move to separate lowering and make the lowering common
-    private fun generateContinuationClass(irFunction: IrFunction) {
-        val continuationImpl = context.continuationImpl.owner
-        val continuationClass = buildClass {
-            // TODO: Use the same counter as CallableReferenceLowering
-            name = "${irFunction.name}\$1".synthesizedName
-            origin = CONTINUATION_CLASS
-            visibility = JavaVisibilities.PACKAGE_VISIBILITY
-        }.apply {
-            createImplicitParameterDeclarationWithWrappedDescriptor()
-            parent = irFunction.parent
-            superTypes.add(continuationImpl.defaultType)
+    private fun transformReferencesToSuspendLambdas(irFile: IrFile, suspendLambdas: List<SuspendLambdaInfo>) {
+        irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
+            private val functionStack = Stack<IrFunction>()
 
-            (parent as IrDeclarationContainer).declarations.add(this)
-        }
+            override fun visitFunction(declaration: IrFunction): IrStatement {
+                functionStack.push(declaration)
+                val res = super.visitFunction(declaration)
+                functionStack.pop()
+                return res
+            }
 
-        val resultField = buildField {
-            origin = CONTINUATION_CLASS_RESULT_FIELD
-            name = Name.identifier(context.state.languageVersionSettings.dataFieldName())
-            type = context.irBuiltIns.anyType
-            isFinal = false
-            visibility = JavaVisibilities.PACKAGE_VISIBILITY
-        }.apply {
-            parent = continuationClass
-            continuationClass.declarations.add(this)
-        }
+            override fun visitFunctionReference(expression: IrFunctionReference): IrExpression {
+                expression.transformChildrenVoid(this)
 
-        val labelField = buildField {
-            origin = CONTINUATION_CLASS_LABEL_FIELD
-            name = Name.identifier(COROUTINE_LABEL_FIELD_NAME)
-            type = context.irBuiltIns.intType
-            isFinal = false
-            visibility = JavaVisibilities.PACKAGE_VISIBILITY
-        }.apply {
-            parent = continuationClass
-            continuationClass.declarations.add(this)
-        }
+                if (!expression.isSuspend)
+                    return expression
+                val constructor = suspendLambdas.single { it.function == expression.symbol.owner }.constructor
+                val expressionArguments = expression.getArguments().map { it.second }
+                assert(constructor.valueParameters.size == expressionArguments.size) {
+                    "Inconsistency between callable reference to suspend lambda and the corresponding continuation"
+                }
+                val irBuilder = context.createIrBuilder(expression.symbol, expression.startOffset, expression.endOffset)
+                irBuilder.run {
+                    return irCall(constructor.symbol).apply {
+                        expressionArguments.forEachIndexed { index, argument ->
+                            putValueArgument(index, argument)
+                        }
+                    }
+                }
+            }
+        })
+    }
 
-        // TODO: How the fuck do I disable parameter assertions generation????
-        buildConstructor {
-            isPrimary = true
-            origin = CONTINUATION_CLASS_CONSTRUCTOR
-            visibility = Visibilities.PUBLIC
-            returnType = continuationClass.defaultType
-        }.apply {
-            parent = continuationClass
-            continuationClass.declarations.add(this)
-
-            val completionParameterDescriptor = WrappedValueParameterDescriptor()
-            val completionParameterSymbol = IrValueParameterSymbolImpl(completionParameterDescriptor)
-            valueParameters.add(
-                0, IrValueParameterImpl(
-                    UNDEFINED_OFFSET, UNDEFINED_OFFSET,
-                    CONTINUATION_CLASS_COMPLETION_PARAMETER,
-                    completionParameterSymbol,
-                    Name.identifier("completion"),
-                    0,
-                    context.continuationClass.owner.defaultType,
-                    null,
-                    isCrossinline = false, isNoinline = false
-                ).also { completionParameterDescriptor.bind(it) }
+    private fun generateContinuationClassForLambda(info: SuspendLambdaInfo) {
+        val suspendLambda = context.suspendLambda.owner
+        suspendLambda.createContinuationClassFor(info.function).apply {
+            val functionNClass = context.ir.symbols.getJvmFunctionClass(info.arity + 1)
+            superTypes.add(
+                IrSimpleTypeImpl(
+                    functionNClass,
+                    hasQuestionMark = false,
+                    arguments = (info.function.explicitParameters.subList(0, info.arity).map { it.type }
+                            + info.function.continuationType() + info.function.returnType)
+                        .map { makeTypeProjection(it, Variance.INVARIANT) },
+                    annotations = emptyList()
+                )
             )
 
-            val superClassConstructor = continuationImpl.constructors.single { it.valueParameters.size == 1 }
-            body = context.createIrBuilder(symbol).irBlockBody {
+            addField(COROUTINE_LABEL_FIELD_NAME, context.irBuiltIns.intType)
+
+            val parametersFields = info.function.valueParameters.map { addField(it.name.asString(), it.type) }
+            val constructor = addPrimaryConstructorForLambda(info.arity, info.function, parametersFields)
+            val secondaryConstructor = addSecondaryConstructorForLambda(constructor)
+            val invokeToOverride = functionNClass.functions.single { it.owner.valueParameters.size == info.arity + 1 }
+            val invokeSuspend = addInvokeSuspendForLambda(info.function, parametersFields)
+            if (info.arity <= 1) {
+                val create = addCreate(constructor, suspendLambda, info.arity, parametersFields)
+                addInvoke(create, invokeSuspend, invokeToOverride)
+            } else {
+                addInvoke(constructor, invokeSuspend, invokeToOverride)
+            }
+
+            context.suspendLambdaClasses[this] = info.function.symbol.descriptor.psiElement as KtElement
+
+            info.constructor = secondaryConstructor
+        }
+    }
+
+    private fun IrClass.addInvokeSuspendForLambda(irFunction: IrFunction, fields: List<IrField>): IrFunction {
+        return addFunctionOverride(context.suspendLambda.functions.single { it.owner.name.asString() == INVOKE_SUSPEND_METHOD_NAME }.owner)
+            .also { function ->
+                function.body = irFunction.body?.deepCopyWithSymbols()
+                function.body?.transformChildrenVoid(object : IrElementTransformerVoid() {
+                    override fun visitGetValue(expression: IrGetValue): IrExpression {
+                        val field = fields.single { it.name == expression.symbol.owner.name }
+                        return IrGetFieldImpl(expression.startOffset, expression.endOffset, field.symbol, field.type).also {
+                            it.receiver = IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, function.dispatchReceiverParameter!!.symbol)
+                        }
+                    }
+
+                    override fun visitReturn(expression: IrReturn): IrExpression {
+                        val ret = super.visitReturn(expression) as IrReturn
+                        return IrReturnImpl(ret.startOffset, ret.endOffset, context.irBuiltIns.anyType, function.symbol, ret.value)
+                    }
+                })
+                (irFunction.parent as IrDeclarationContainer).declarations.remove(irFunction)
+            }
+    }
+
+    private fun IrClass.addInvoke(
+        create: IrFunction,
+        invokeSuspend: IrFunction,
+        invokeToOverride: IrSimpleFunctionSymbol
+    ) {
+        addFunctionOverride(invokeToOverride.owner).also { function ->
+            function.body = context.createIrBuilder(function.symbol).irBlockBody {
+                +irReturn(irCall(invokeSuspend).also { invokeSuspendCall ->
+                    invokeSuspendCall.dispatchReceiver = irCall(create).also {
+                        it.dispatchReceiver = irGet(function.dispatchReceiverParameter!!)
+                        for ((index, param) in function.valueParameters.withIndex()) {
+                            it.putValueArgument(index, irGet(param))
+                        }
+                    }
+                    invokeSuspendCall.putValueArgument(0, irUnit())
+                })
+            }
+        }
+    }
+
+    private fun IrClass.addCreate(
+        constructor: IrFunction,
+        superType: IrClass,
+        arity: Int,
+        parametersFields: List<IrField>
+    ): IrFunction {
+        val create = superType.functions.single { it.name == Name.identifier("create") && it.valueParameters.size == arity + 1 }
+        return addFunctionOverride(create).also { function ->
+            function.body = context.createIrBuilder(function.symbol).irBlockBody {
+                +irReturn(irCall(constructor).also {
+                    for ((i, field) in parametersFields.withIndex()) {
+                        it.putValueArgument(i, irGetField(irGet(function.dispatchReceiverParameter!!), field))
+                    }
+                    for ((i, parameter) in function.valueParameters.withIndex()) {
+                        it.putValueArgument(parametersFields.size + i, irGet(parameter))
+                    }
+                })
+            }
+        }
+    }
+
+    private fun IrClass.createContinuationClassFor(irFunction: IrFunction): IrClass = buildClass {
+        name = "${irFunction.name}\$${context.functionReferenceAndContinuationsCount++}".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
+        irClass.parent = parent
+        parent.declarations.add(irClass)
+    }
+
+    private fun IrClass.addPrimaryConstructorForLambda(
+        arity: Int,
+        irFunction: IrFunction,
+        fields: List<IrField>
+    ): IrConstructor =
+        addConstructor {
+            isPrimary = true
+            returnType = defaultType
+        }.also { constructor ->
+            constructor.valueParameters.addAll(irFunction.valueParameters.map { it.copyTo(constructor) })
+            val completionParameterSymbol = constructor.addCompletionValueParameter()
+
+            val superClassConstructor = context.suspendLambda.owner.constructors.single { it.valueParameters.size == 2 }
+            constructor.body = context.createIrBuilder(constructor.symbol).irBlockBody {
                 +irDelegatingConstructorCall(superClassConstructor).also {
-                    it.putValueArgument(0, irGet(completionParameterSymbol.owner))
+                    it.putValueArgument(0, irInt(arity + 1))
+                    it.putValueArgument(1, irGet(completionParameterSymbol))
+                }
+
+                for ((index, param) in constructor.valueParameters.dropLast(1).withIndex()) {
+                    +irSetField(irGet(thisReceiver!!), fields[index], irGet(param))
                 }
             }
         }
 
-        val thisReceiver = continuationClass.thisReceiver!!
+    private fun IrClass.addSecondaryConstructorForLambda(primary: IrConstructor): IrConstructor =
+        addConstructor {
+            isPrimary = false
+            returnType = defaultType
+        }.also { constructor ->
+            constructor.valueParameters.addAll(primary.valueParameters.dropLast(1).map { it.copyTo(constructor) })
 
-        val invokeSuspendName = Name.identifier(INVOKE_SUSPEND_METHOD_NAME)
-        val invokeSuspend = continuationImpl.functions.single { it.name == invokeSuspendName }.symbol
-        buildFun {
-            origin = CONTINUATION_CLASS_INVOKE_SUSPEND
-            name = invokeSuspendName
-            visibility = Visibilities.PUBLIC
-            isSuspend = false
-            returnType = context.irBuiltIns.anyType
-        }.apply {
-            parent = continuationClass
-            continuationClass.declarations.add(this)
-            overriddenSymbols.add(invokeSuspend)
-            dispatchReceiverParameter = thisReceiver.copyTo(this)
+            constructor.body = context.createIrBuilder(constructor.symbol).irBlockBody {
+                +irDelegatingConstructorCall(primary).also {
+                    for ((index, param) in constructor.valueParameters.withIndex()) {
+                        it.putValueArgument(index, irGet(param))
+                    }
+                    it.putValueArgument(constructor.valueParameters.size, irNull())
+                }
+            }
+        }
 
-            val result = invokeSuspend.owner.valueParameters[0].copyTo(this)
-            valueParameters.add(result)
-            body = context.createIrBuilder(symbol).irBlockBody {
-                +irSetField(irGet(dispatchReceiverParameter!!), resultField, irGet(result))
+    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 =
+        context.continuationClass.createType(true, listOf(makeTypeProjection(returnType, Variance.INVARIANT)))
+
+    private fun generateContinuationClassForNamedFunction(irFunction: IrFunction) {
+        context.continuationImpl.owner.createContinuationClassFor(irFunction).apply {
+            val resultField = addField(
+                context.state.languageVersionSettings.dataFieldName(),
+                context.irBuiltIns.anyType,
+                JavaVisibilities.PACKAGE_VISIBILITY
+            )
+            val labelField = addField(COROUTINE_LABEL_FIELD_NAME, context.irBuiltIns.intType, JavaVisibilities.PACKAGE_VISIBILITY)
+            addConstructorForNamedFunction()
+            addInvokeSuspendForNamedFunction(irFunction, resultField, labelField)
+
+            context.suspendFunctionContinuations[irFunction] = this
+        }
+    }
+
+    private fun IrClass.addConstructorForNamedFunction(): IrConstructor = addConstructor {
+        isPrimary = true
+        returnType = defaultType
+    }.also { constructor ->
+        val completionParameterSymbol = constructor.addCompletionValueParameter()
+
+        val superClassConstructor = context.continuationImpl.owner.constructors.single { it.valueParameters.size == 1 }
+        constructor.body = context.createIrBuilder(constructor.symbol).irBlockBody {
+            +irDelegatingConstructorCall(superClassConstructor).also {
+                it.putValueArgument(0, irGet(completionParameterSymbol))
+            }
+        }
+    }
+
+    private fun IrClass.addInvokeSuspendForNamedFunction(irFunction: IrFunction, resultField: IrField, labelField: IrField) {
+        val invokeSuspend = context.continuationImpl.owner.functions.single { it.name == Name.identifier(INVOKE_SUSPEND_METHOD_NAME) }
+        addFunctionOverride(invokeSuspend).also { function ->
+            function.body = context.createIrBuilder(function.symbol).irBlockBody {
+                +irSetField(irGet(function.dispatchReceiverParameter!!), resultField, irGet(function.valueParameters[0]))
                 +irSetField(
-                    irGet(dispatchReceiverParameter!!), labelField,
-                    // TODO: Why the fuck label is boxed
+                    irGet(function.dispatchReceiverParameter!!), labelField,
                     irCallOp(
                         context.irBuiltIns.intClass.functions.single { it.owner.name == OperatorNameConventions.OR },
                         context.irBuiltIns.intType,
-                        irGetField(irGet(dispatchReceiverParameter!!), labelField),
+                        irGetField(irGet(function.dispatchReceiverParameter!!), labelField),
                         irInt(1 shl 31)
                     )
                 )
-                +irReturn(irCall(irFunction).also { it.putValueArgument(0, irGet(dispatchReceiverParameter!!)) })
+                +irReturn(irCall(irFunction).also { it.putValueArgument(0, irGet(function.dispatchReceiverParameter!!)) })
             }
         }
     }
 
-    private fun addContinuationParameter(irFile: IrFile, suspendLambdas: List<IrFunction>): Set<IrFunction> {
+    private fun addContinuationParameter(irFile: IrFile, suspendLambdas: List<SuspendLambdaInfo>): Set<IrFunction> {
         val views = hashSetOf<IrFunction>()
         fun tryAddContinuationParameter(element: IrElement): List<IrDeclaration>? {
-            return if (element is IrSimpleFunction && element.isSuspend) {
-                val view = element.getOrCreateView()
-                if (element !in suspendLambdas) {
-                    views.add(view)
-                }
-                listOf(view)
+            return if (element is IrSimpleFunction && element.isSuspend && suspendLambdas.none { it.function == element }) {
+                listOf(element.getOrCreateView().also { views.add(it) })
             } else null
         }
 
@@ -205,8 +326,8 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         return views
     }
 
-    private fun markSuspendLambdas(irElement: IrElement): List<IrFunction> {
-        val suspendLambdas = arrayListOf<IrFunction>()
+    private fun markSuspendLambdas(irElement: IrElement): List<SuspendLambdaInfo> {
+        val suspendLambdas = arrayListOf<SuspendLambdaInfo>()
         irElement.acceptChildrenVoid(object : IrElementVisitorVoid {
             override fun visitElement(element: IrElement) {
                 element.acceptChildrenVoid(this)
@@ -216,17 +337,16 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                 expression.acceptChildrenVoid(this)
 
                 if (expression.isSuspend) {
-                    suspendLambdas += expression.symbol.owner
+                    suspendLambdas += SuspendLambdaInfo(expression.symbol.owner, expression.type.originalKotlinType!!.arguments.size - 1)
                 }
             }
         })
         return suspendLambdas
     }
 
-    private fun transformSuspendCalls(irFile: IrFile, suspendLambdas: List<IrFunction>) {
+    private fun transformSuspendCalls(irFile: IrFile) {
         irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
             private val functionStack = Stack<IrFunction>()
-
             override fun visitElement(element: IrElement): IrElement {
                 element.transformChildrenVoid(this)
                 return element
@@ -239,16 +359,22 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                 return res
             }
 
+            private fun getContinuationFromCaller(): IrGetValue {
+                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)
+                        caller.dispatchReceiverParameter!!.symbol
+                    else error("suspend call outside suspend context")
+                return IrGetValueImpl(
+                    UNDEFINED_OFFSET,
+                    UNDEFINED_OFFSET,
+                    continuationParameter
+                )
+            }
+
             override fun visitCall(expression: IrCall): IrExpression {
                 if (!expression.isSuspend) return super.visitCall(expression)
-                val caller = functionStack.peek()
-                val continuationArgument = IrGetValueImpl(
-                    UNDEFINED_OFFSET,
-                    UNDEFINED_OFFSET,
-                    continuationType(caller.returnType),
-                    if (caller in suspendLambdas) (caller.dispatchReceiverParameter as IrValueParameter).symbol
-                    else caller.valueParameters.last().symbol
-                )
                 val view = (expression.symbol.owner as IrSimpleFunction).getOrCreateView()
                 val res = IrCallImpl(expression.startOffset, expression.endOffset, expression.type, view.symbol).apply {
                     copyTypeArgumentsFrom(expression)
@@ -256,7 +382,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                     for (i in 0 until expression.valueArgumentsCount) {
                         putValueArgument(i, expression.getValueArgument(i))
                     }
-                    putValueArgument(expression.valueArgumentsCount, continuationArgument)
+                    putValueArgument(expression.valueArgumentsCount, getContinuationFromCaller())
                 }
                 return super.visitCall(res)
             }
@@ -295,10 +421,10 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                 it.dispatchReceiverParameter = dispatchReceiverParameter?.copyTo(it)
 
                 valueParameters.mapTo(it.valueParameters) { p -> p.copyTo(it) }
-                it.valueParameters += createContinuationValueParameter(it.valueParameters.size, returnType).apply { parent = it }
+                it.addCompletionValueParameter()
 
                 // Fix links to parameters
-                val valueParametersMapping = (valueParameters + dispatchReceiverParameter)
+                val valueParametersMapping: Map<IrValueDeclaration?, IrValueParameter?> = (valueParameters + dispatchReceiverParameter)
                     .zip(it.valueParameters.dropLast(1) + it.dispatchReceiverParameter).toMap()
                 it.body = body?.deepCopyWithSymbols(this)
                 it.body?.transformChildrenVoid(object : IrElementTransformerVoid() {
@@ -310,25 +436,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             }
         }
 
-    private fun createContinuationValueParameter(index: Int, returnType: IrType): IrValueParameter {
-        val descriptor = WrappedValueParameterDescriptor()
-        return IrValueParameterImpl(
-            UNDEFINED_OFFSET,
-            UNDEFINED_OFFSET,
-            CONTINUATION_PARAMETER,
-            IrValueParameterSymbolImpl(descriptor),
-            Name.identifier("\$completion"),
-            index,
-            continuationType(returnType),
-            null, isCrossinline = false, isNoinline = false
-        ).also {
-            descriptor.bind(it)
-        }
+    private class SuspendLambdaInfo(val function: IrFunction, val arity: Int) {
+        lateinit var constructor: IrConstructor
     }
-
-    private fun continuationType(returnType: IrType) = context.continuationClass.createType(
-        hasQuestionMark = false,
-        arguments = listOf(makeTypeProjection(returnType, Variance.INVARIANT))
-    )
-
 }