JVM_IR: Support crossinline suspend lambdas

The main idea is the following: since we need to generate
(fake)continuations before inlining, we move IrClasses of suspend
lambdas and continuation classes of named functions into the functions.
Thus, it allows the codegen to generate them prior to inlining and
the inliner will happily transform them for us.
Because of that, lowerings which transform call-site function are likely
to change reference to lowered suspend lambdas or functions.
Hence, do not rely on references to lowered suspend lambdas or
functions, instead, rely on attributes.

Do not generate continuation for inline suspend lambdas.
Previously, inline suspend lambdas were treated like suspend functions,
thus we generated continuations for them. Now we just do not treat them
as suspend functions or lambdas during AddContinuationLowering.
We should add continuation parameter to them, however.

Do not generate secondary constructor for suspend lambdas, otherwise,
the inliner is unable to transform them (it requires only one
constructor to be present).

Generate continuation classes for suspend functions as first statement
inside the function.
This enables suspend functions in local object inside inline functions.
Since we already have attributes inside suspend named functions, we
just reuse them to generate continuation class names. This allows us
to close the gap between code generated by old back-end and the new
one.

If a suspend named function captures crossinline lambda, we should
generate a template for inliner: a copy of the function without
state-machine and a continuation constructor call. The call is needed
so the inliner transforms the continuation as well.

Refactor CoroutineTransformerMethodVisitor, so it no longer depends on
PSI.

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

@@ -83,6 +83,7 @@ abstract class IrElementTransformerVoidWithContext : IrElementTransformerVoid()
     protected val currentScope get() = scopeStack.peek()
     protected val parentScope get() = if (scopeStack.size < 2) null else scopeStack[scopeStack.size - 2]
     protected val allScopes get() = scopeStack
+    protected val currentDeclarationParent get() = allScopes.last { it.irElement is IrDeclarationParent }.irElement as IrDeclarationParent
 
     fun printScopeStack() {
         scopeStack.forEach { println(it.scope.scopeOwner) }

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.backend.common.ir
 import org.jetbrains.kotlin.backend.common.CommonBackendContext
 import org.jetbrains.kotlin.backend.common.DumpIrTreeWithDescriptorsVisitor
 import org.jetbrains.kotlin.backend.common.deepCopyWithVariables
+import org.jetbrains.kotlin.backend.common.descriptors.*
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
@@ -566,9 +567,17 @@ fun copyBodyToStatic(oldFunction: IrFunction, staticFunction: IrFunction) {
     val mapping: Map<IrValueParameter, IrValueParameter> =
         (listOfNotNull(oldFunction.dispatchReceiverParameter, oldFunction.extensionReceiverParameter) + oldFunction.valueParameters)
             .zip(staticFunction.valueParameters).toMap()
-    staticFunction.body = oldFunction.body
+    copyBodyWithParametersMapping(staticFunction, oldFunction, mapping)
+}
+
+fun copyBodyWithParametersMapping(
+    newFunction: IrFunction,
+    oldFunction: IrFunction,
+    mapping: Map<IrValueParameter, IrValueParameter>
+) {
+    newFunction.body = oldFunction.body?.deepCopyWithSymbols(oldFunction)
         ?.transform(
-            object: IrElementTransformerVoid() {
+            object : IrElementTransformerVoid() {
                 // Remap return targets to the static method so they do not appear to be
                 // non-local returns.
                 override fun visitReturn(expression: IrReturn): IrExpression {
@@ -578,8 +587,9 @@ fun copyBodyToStatic(oldFunction: IrFunction, staticFunction: IrFunction) {
                             expression.startOffset,
                             expression.endOffset,
                             expression.type,
-                            staticFunction.symbol,
-                            expression.value)
+                            newFunction.symbol,
+                            expression.value
+                        )
                     } else expression
                 }
 
@@ -589,8 +599,9 @@ fun copyBodyToStatic(oldFunction: IrFunction, staticFunction: IrFunction) {
                         IrGetValueImpl(expression.startOffset, expression.endOffset, it.type, it.symbol, expression.origin)
                     } ?: expression
 
-            }, null)
-        ?.patchDeclarationParents(staticFunction)
+            }, null
+        )
+        ?.patchDeclarationParents(newFunction)
 }
 
 val IrSymbol.isSuspend: Boolean

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

@@ -87,6 +87,9 @@ class LocalDeclarationsLowering(
     object DECLARATION_ORIGIN_FIELD_FOR_CAPTURED_VALUE :
         IrDeclarationOriginImpl("FIELD_FOR_CAPTURED_VALUE", isSynthetic = true)
 
+    object DECLARATION_ORIGIN_FIELD_FOR_CROSSINLINE_CAPTURED_VALUE :
+        IrDeclarationOriginImpl("FIELD_FOR_CROSSINLINE_CAPTURED_VALUE", isSynthetic = true)
+
     private object STATEMENT_ORIGIN_INITIALIZER_OF_FIELD_FOR_CAPTURED_VALUE :
         IrStatementOriginImpl("INITIALIZER_OF_FIELD_FOR_CAPTURED_VALUE")
 
@@ -655,14 +658,15 @@ class LocalDeclarationsLowering(
             name: Name,
             visibility: Visibility,
             parent: IrClass,
-            fieldType: IrType
+            fieldType: IrType,
+            isCrossinline: Boolean
         ): IrField {
             val descriptor = WrappedFieldDescriptor()
             val symbol = IrFieldSymbolImpl(descriptor)
             return IrFieldImpl(
                 startOffset,
                 endOffset,
-                DECLARATION_ORIGIN_FIELD_FOR_CAPTURED_VALUE,
+                if (isCrossinline) DECLARATION_ORIGIN_FIELD_FOR_CROSSINLINE_CAPTURED_VALUE else DECLARATION_ORIGIN_FIELD_FOR_CAPTURED_VALUE,
                 symbol,
                 name,
                 fieldType,
@@ -682,16 +686,18 @@ class LocalDeclarationsLowering(
             val generatedNames = mutableSetOf<Name>()
             localClassContext.closure.capturedValues.forEach { capturedValue ->
 
+                val owner = capturedValue.owner
                 val irField = createFieldForCapturedValue(
                     classDeclaration.startOffset,
                     classDeclaration.endOffset,
-                    suggestNameForCapturedValue(capturedValue.owner, generatedNames),
+                    suggestNameForCapturedValue(owner, generatedNames),
                     Visibilities.PRIVATE,
                     classDeclaration,
-                    capturedValue.owner.type
+                    owner.type,
+                    owner is IrValueParameter && owner.isCrossinline
                 )
 
-                localClassContext.capturedValueToField[capturedValue.owner] = irField
+                localClassContext.capturedValueToField[owner] = irField
             }
         }
 

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

@@ -176,13 +176,15 @@ fun IrDeclarationContainer.addFunction(
     returnType: IrType,
     modality: Modality = Modality.FINAL,
     isStatic: Boolean = false,
-    isSuspend: Boolean = false
+    isSuspend: Boolean = false,
+    origin: IrDeclarationOrigin = IrDeclarationOrigin.DEFINED
 ): IrSimpleFunction =
     addFunction {
         this.name = Name.identifier(name)
         this.returnType = returnType
         this.modality = modality
         this.isSuspend = isSuspend
+        this.origin = origin
     }.apply {
         if (!isStatic) {
             dispatchReceiverParameter = parentAsClass.thisReceiver!!.copyTo(this)

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/DeclarationOrigins.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
@@ -51,4 +40,9 @@ interface JvmLoweredDeclarationOrigin : IrDeclarationOrigin {
     object POLYMORPHIC_SIGNATURE_INSTANTIATION : IrDeclarationOriginImpl("POLYMORPHIC_SIGNATURE_INSTANTIATION", isSynthetic = true)
     object ENUM_CONSTRUCTOR_SYNTHETIC_PARAMETER : IrDeclarationOriginImpl("ENUM_CONSTRUCTOR_SYNTHETIC_PARAMETER", isSynthetic = true)
     object OBJECT_SUPER_CONSTRUCTOR_PARAMETER : IrDeclarationOriginImpl("OBJECT_SUPER_CONSTURCTOR_PARAMETER", isSynthetic = true)
+    object CONTINUATION_CLASS : IrDeclarationOriginImpl("CONTINUATION_CLASS")
+    object SUSPEND_LAMBDA : IrDeclarationOriginImpl("SUSPEND_LAMBDA")
+    object FOR_INLINE_STATE_MACHINE_TEMPLATE : IrDeclarationOriginImpl("FOR_INLINE_TEMPLATE")
+    object FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE : IrDeclarationOriginImpl("FOR_INLINE_TEMPLATE_CROSSINLINE")
+    object FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE_VIEW : IrDeclarationOriginImpl("FOR_INLINE_TEMPLATE_CROSSINLINE_VIEW")
 }

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

@@ -31,6 +31,7 @@ import org.jetbrains.kotlin.ir.builders.irNull
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
 import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrFunctionReference
 import org.jetbrains.kotlin.ir.symbols.*
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.util.SymbolTable
@@ -97,9 +98,8 @@ class JvmBackendContext(
 
     override val internalPackageFqn = FqName("kotlin.jvm")
 
-    val suspendFunctionContinuations = mutableMapOf<IrFunction, IrClass>()
-    val suspendLambdaToOriginalFunctionMap = mutableMapOf<IrClass, IrFunction>()
-    val continuationClassBuilders = mutableMapOf<IrClass, ClassBuilder>()
+    val suspendLambdaToOriginalFunctionMap = mutableMapOf<IrFunctionReference, IrFunction>()
+    val continuationClassBuilders = mutableMapOf<IrSimpleFunction, ClassBuilder>()
     val suspendFunctionOriginalToView = mutableMapOf<IrFunction, IrFunction>()
     val suspendFunctionViewToOriginal = mutableMapOf<IrFunction, IrFunction>()
     val fakeContinuation: IrExpression = createFakeContinuation(this)

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.backend.jvm.codegen
 
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.backend.jvm.lower.MultifileFacadeFileEntry
 import org.jetbrains.kotlin.backend.jvm.lower.buildAssertionsDisabledField
 import org.jetbrains.kotlin.backend.jvm.lower.hasAssertionsDisabledField
@@ -126,12 +127,6 @@ open class ClassCodegen protected constructor(
             } 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()
-        }
-
         val fileEntry = context.psiSourceManager.getFileEntry(irClass.fileParent)
         if (fileEntry != null) {
             /* TODO: Temporary workaround: ClassBuilder needs a pathless name. */
@@ -153,14 +148,14 @@ 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 - continuationCodegens)) {
+        for (codegen in nestedClasses) {
             codegen.generate()
         }
 
         generateKotlinMetadataAnnotation()
 
-        if (irClass in context.suspendFunctionContinuations.values) {
-            context.continuationClassBuilders[irClass] = visitor
+        if (irClass.origin == JvmLoweredDeclarationOrigin.CONTINUATION_CLASS) {
+            context.continuationClassBuilders[irClass.attributeOwnerId as IrSimpleFunction] = visitor
         } else {
             done()
         }

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.backend.jvm.codegen
 
+import org.jetbrains.kotlin.backend.common.CodegenUtil
 import org.jetbrains.kotlin.backend.common.ir.copyTo
 import org.jetbrains.kotlin.backend.common.ir.copyTypeParametersFrom
 import org.jetbrains.kotlin.backend.common.ir.isSuspend
@@ -14,6 +15,8 @@ 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.codegen.coroutines.SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME
+import org.jetbrains.kotlin.codegen.coroutines.reportSuspensionPointInsideMonitor
+import org.jetbrains.kotlin.codegen.inline.addFakeContinuationConstructorCallMarker
 import org.jetbrains.kotlin.config.isReleaseCoroutines
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
@@ -37,11 +40,14 @@ import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodGenericSignature
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DescriptorWithContainerSource
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.org.objectweb.asm.MethodVisitor
 import org.jetbrains.org.objectweb.asm.Opcodes
+import org.jetbrains.org.objectweb.asm.Type
+import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter
 
 internal fun generateStateMachineForNamedFunction(
     irFunction: IrFunction,
@@ -49,21 +55,19 @@ internal fun generateStateMachineForNamedFunction(
     methodVisitor: MethodVisitor,
     access: Int,
     signature: JvmMethodGenericSignature,
-    continuationClassBuilder: ClassBuilder?,
+    obtainContinuationClassBuilder: () -> ClassBuilder,
     element: KtElement
 ): MethodVisitor {
     assert(irFunction.isSuspend)
-    assert(continuationClassBuilder != null) {
-        "Class builder for continuation is null"
-    }
     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, signature.genericsSignature, null,
-        obtainClassBuilderForCoroutineState = { continuationClassBuilder!! },
-        element = element,
-        diagnostics = state.diagnostics,
+        methodVisitor, access, signature.asmMethod.name, signature.asmMethod.descriptor, null, null,
+        obtainClassBuilderForCoroutineState = obtainContinuationClassBuilder,
+        reportSuspensionPointInsideMonitor = { reportSuspensionPointInsideMonitor(element, state, it) },
+        lineNumber = CodegenUtil.getLineNumberForElement(element, false) ?: 0,
+        sourceFile = classCodegen.irClass.file.name,
         languageVersionSettings = languageVersionSettings,
         shouldPreserveClassInitialization = state.constructorCallNormalizationMode.shouldPreserveClassInitialization,
         containingClassInternalName = classCodegen.visitor.thisName,
@@ -92,8 +96,9 @@ internal fun generateStateMachineForLambda(
     return CoroutineTransformerMethodVisitor(
         methodVisitor, access, signature.asmMethod.name, signature.asmMethod.descriptor, signature.genericsSignature, null,
         obtainClassBuilderForCoroutineState = { classCodegen.visitor },
-        element = element,
-        diagnostics = state.diagnostics,
+        reportSuspensionPointInsideMonitor = { reportSuspensionPointInsideMonitor(element, state, it) },
+        lineNumber = CodegenUtil.getLineNumberForElement(element, false) ?: 0,
+        sourceFile = classCodegen.irClass.file.name,
         languageVersionSettings = languageVersionSettings,
         shouldPreserveClassInitialization = state.constructorCallNormalizationMode.shouldPreserveClassInitialization,
         containingClassInternalName = classCodegen.visitor.thisName,
@@ -102,17 +107,18 @@ internal fun generateStateMachineForLambda(
     )
 }
 
-internal fun IrFunction.isInvokeSuspendOfLambda(context: JvmBackendContext): Boolean =
-    name.asString() == INVOKE_SUSPEND_METHOD_NAME && parent in context.suspendLambdaToOriginalFunctionMap
+internal fun IrFunction.isInvokeSuspendOfLambda(): Boolean =
+    name.asString() == INVOKE_SUSPEND_METHOD_NAME && parentAsClass.origin == JvmLoweredDeclarationOrigin.SUSPEND_LAMBDA
 
-internal fun IrFunction.isInvokeOfSuspendLambda(context: JvmBackendContext): Boolean =
-    name.asString() == "invoke" && parent in context.suspendLambdaToOriginalFunctionMap
+internal fun IrFunction.isInvokeSuspendForInlineOfLambda(): Boolean =
+    origin == JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE
+            && parentAsClass.origin == JvmLoweredDeclarationOrigin.SUSPEND_LAMBDA
 
-internal fun IrFunction.isInvokeSuspendOfContinuation(context: JvmBackendContext): Boolean =
-    name.asString() == INVOKE_SUSPEND_METHOD_NAME && parentAsClass in context.suspendFunctionContinuations.values
+internal fun IrFunction.isInvokeSuspendOfContinuation(): Boolean =
+    name.asString() == INVOKE_SUSPEND_METHOD_NAME && parentAsClass.origin == JvmLoweredDeclarationOrigin.CONTINUATION_CLASS
 
 internal fun IrFunction.isInvokeOfSuspendCallableReference(): Boolean = isSuspend && name.asString() == "invoke" &&
-        (parent as? IrClass)?.origin == JvmLoweredDeclarationOrigin.FUNCTION_REFERENCE_IMPL
+        parentAsClass.origin == JvmLoweredDeclarationOrigin.FUNCTION_REFERENCE_IMPL
 
 internal fun IrFunction.isKnownToBeTailCall(): Boolean =
     when (origin) {
@@ -126,13 +132,15 @@ internal fun IrFunction.isKnownToBeTailCall(): Boolean =
         else -> isInvokeOfSuspendCallableReference()
     }
 
-internal fun IrFunction.shouldNotContainSuspendMarkers(context: JvmBackendContext): Boolean =
-    isInvokeSuspendOfContinuation(context) || isKnownToBeTailCall()
+internal fun IrFunction.shouldNotContainSuspendMarkers(): Boolean =
+    isInvokeSuspendOfContinuation() || isKnownToBeTailCall()
 
 // Transform `suspend fun foo(params): RetType` into `fun foo(params, $completion: Continuation<RetType>): Any?`
 // the result is called 'view', just to be consistent with old backend.
 internal fun IrFunction.getOrCreateSuspendFunctionViewIfNeeded(context: JvmBackendContext): IrFunction {
-    if (!isSuspend || origin == JvmLoweredDeclarationOrigin.SUSPEND_FUNCTION_VIEW) return this
+    if (!isSuspend || origin == JvmLoweredDeclarationOrigin.SUSPEND_FUNCTION_VIEW ||
+        origin == JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE_VIEW
+    ) return this
     context.suspendFunctionOriginalToView[this]?.let { return it }
     return suspendFunctionView(context)
 }
@@ -152,7 +160,12 @@ private fun IrFunction.suspendFunctionView(context: JvmBackendContext): IrFuncti
         else
             WrappedSimpleFunctionDescriptor(sourceElement = originalDescriptor.source)
     return IrFunctionImpl(
-        startOffset, endOffset, JvmLoweredDeclarationOrigin.SUSPEND_FUNCTION_VIEW, IrSimpleFunctionSymbolImpl(descriptor),
+        startOffset, endOffset,
+        if (origin == JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE)
+            JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE_VIEW
+        else
+            JvmLoweredDeclarationOrigin.SUSPEND_FUNCTION_VIEW,
+        IrSimpleFunctionSymbolImpl(descriptor),
         name, visibility, modality, context.irBuiltIns.anyNType,
         isInline = isInline, isExternal = isExternal, isTailrec = isTailrec, isSuspend = isSuspend, isExpect = isExpect,
         isFakeOverride = false, isOperator = false
@@ -196,6 +209,7 @@ private fun IrFunction.suspendFunctionView(context: JvmBackendContext): IrFuncti
                 return super.visitCall(expression.createSuspendFunctionCallViewIfNeeded(context, it, callerIsInlineLambda = false))
             }
         })
+        it.copyAttributes(this)
     }
 }
 
@@ -216,7 +230,8 @@ internal fun IrCall.createSuspendFunctionCallViewIfNeeded(
         }
         val continuationParameter =
             when {
-                caller.isInvokeSuspendOfLambda(context) || caller.isInvokeSuspendOfContinuation(context) ->
+                caller.isInvokeSuspendOfLambda() || caller.isInvokeSuspendOfContinuation() ||
+                        caller.isInvokeSuspendForInlineOfLambda() ->
                     IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, caller.dispatchReceiverParameter!!.symbol)
                 callerIsInlineLambda -> context.fakeContinuation
                 else -> IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, caller.valueParameters.last().symbol)
@@ -231,3 +246,31 @@ internal fun createFakeContinuation(context: JvmBackendContext): IrExpression =
     context.ir.symbols.continuationClass.createType(true, listOf(makeTypeProjection(context.irBuiltIns.anyNType, Variance.INVARIANT))),
     "FAKE_CONTINUATION"
 )
+
+internal fun generateFakeContinuationConstructorCall(
+    v: InstructionAdapter,
+    containingClassBuilder: ClassBuilder,
+    classBuilder: ClassBuilder,
+    irFunction: IrFunction
+) {
+    val continuationType = Type.getObjectType(classBuilder.thisName)
+    // TODO: This is different in case of DefaultImpls
+    val thisNameType = Type.getObjectType(containingClassBuilder.thisName.replace(".", "/"))
+    val continuationIndex = listOfNotNull(irFunction.dispatchReceiverParameter, irFunction.extensionReceiverParameter).size +
+            irFunction.valueParameters.size - 1
+    with(v) {
+        addFakeContinuationConstructorCallMarker(this, true)
+        anew(continuationType)
+        dup()
+        if (irFunction.dispatchReceiverParameter != null) {
+            load(0, AsmTypes.OBJECT_TYPE)
+            load(continuationIndex, Type.getObjectType("kotlin/coroutines/Continuation"))
+            invokespecial(continuationType.internalName, "<init>", "(${thisNameType}Lkotlin/coroutines/Continuation;)V", false)
+        } else {
+            load(continuationIndex, Type.getObjectType("kotlin/coroutines/Continuation"))
+            invokespecial(continuationType.internalName, "<init>", "(Lkotlin/coroutines/Continuation;)V", false)
+        }
+        pop()
+        addFakeContinuationConstructorCallMarker(this, false)
+    }
+}

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

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.codegen.inline.*
 import org.jetbrains.kotlin.codegen.inline.ReifiedTypeInliner.Companion.putNeedClassReificationMarker
 import org.jetbrains.kotlin.codegen.inline.ReifiedTypeInliner.OperationKind.AS
 import org.jetbrains.kotlin.codegen.inline.ReifiedTypeInliner.OperationKind.SAFE_AS
+import org.jetbrains.kotlin.codegen.inline.coroutines.FOR_INLINE_SUFFIX
 import org.jetbrains.kotlin.codegen.intrinsics.TypeIntrinsics
 import org.jetbrains.kotlin.codegen.pseudoInsns.fakeAlwaysFalseIfeq
 import org.jetbrains.kotlin.codegen.pseudoInsns.fixStackAndJump
@@ -48,9 +49,11 @@ import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeSystemCommonBackendContext
 import org.jetbrains.kotlin.types.model.TypeParameterMarker
+import org.jetbrains.kotlin.utils.addToStdlib.cast
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.kotlin.utils.keysToMap
+import org.jetbrains.kotlin.utils.sure
 import org.jetbrains.org.objectweb.asm.Label
 import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter
@@ -193,6 +196,7 @@ class ExpressionCodegen(
         val info = BlockInfo()
         val body = irFunction.body!!
         generateNonNullAssertions()
+        generateFakeContinuationConstructorIfNeeded()
         val result = body.accept(this, info)
         // If this function has an expression body, return the result of that expression.
         // Otherwise, if it does not end in a return statement, it must be void-returning,
@@ -213,6 +217,19 @@ class ExpressionCodegen(
         writeParameterInLocalVariableTable(startLabel, endLabel)
     }
 
+    private fun generateFakeContinuationConstructorIfNeeded() {
+        if (irFunction.origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE_VIEW) return
+        val continuationClass = classCodegen.irClass.functions.find {
+            it.name.asString() == irFunction.name.asString().removeSuffix(FOR_INLINE_SUFFIX)
+        }?.body?.statements?.get(0) ?: error("could not find continuation for ${irFunction.render()}")
+        generateFakeContinuationConstructorCall(
+            mv,
+            classCodegen.visitor,
+            context.continuationClassBuilders[(continuationClass as IrClass).attributeOwnerId]!!,
+            irFunction
+        )
+    }
+
     private fun generateNonNullAssertions() {
         if (state.isParamAssertionsDisabled)
             return
@@ -229,14 +246,17 @@ class ExpressionCodegen(
                 irFunction.origin == IrDeclarationOrigin.BRIDGE_SPECIAL ||
                 irFunction.origin == JvmLoweredDeclarationOrigin.DEFAULT_IMPLS_BRIDGE ||
                 irFunction.origin == JvmLoweredDeclarationOrigin.JVM_STATIC_WRAPPER ||
-                irFunction.origin == JvmLoweredDeclarationOrigin.MULTIFILE_BRIDGE
+                irFunction.origin == JvmLoweredDeclarationOrigin.MULTIFILE_BRIDGE ||
+                irFunction.parentAsClass.origin == JvmLoweredDeclarationOrigin.CONTINUATION_CLASS ||
+                irFunction.parentAsClass.origin == JvmLoweredDeclarationOrigin.SUSPEND_LAMBDA
 
         if (notCallableFromJava)
             return
 
         // Do not generate non-null checks for suspend function views. When resumed the arguments
         // will be null and the actual values are taken from the continuation.
-        val isSuspendFunctionView = irFunction.origin == JvmLoweredDeclarationOrigin.SUSPEND_FUNCTION_VIEW
+        val isSuspendFunctionView = irFunction.origin == JvmLoweredDeclarationOrigin.SUSPEND_FUNCTION_VIEW ||
+                irFunction.origin == JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE_VIEW
 
         if (isSuspendFunctionView)
             return
@@ -381,7 +401,7 @@ class ExpressionCodegen(
             }
             expression.symbol.descriptor is ConstructorDescriptor ->
                 throw AssertionError("IrCall with ConstructorDescriptor: ${expression.javaClass.simpleName}")
-            callee.isSuspend && !irFunction.shouldNotContainSuspendMarkers(classCodegen.context) ->
+            callee.isSuspend && !irFunction.shouldNotContainSuspendMarkers() ->
                 addInlineMarker(mv, isStartNotEnd = true)
         }
 
@@ -407,13 +427,13 @@ class ExpressionCodegen(
         expression.markLineNumber(true)
 
         // Do not generate redundant markers in continuation class.
-        if (callee.isSuspend && !irFunction.shouldNotContainSuspendMarkers(classCodegen.context)) {
+        if (callee.isSuspend && !irFunction.shouldNotContainSuspendMarkers()) {
             addSuspendMarker(mv, isStartNotEnd = true)
         }
 
         callGenerator.genCall(callable, this, expression)
 
-        if (callee.isSuspend && !irFunction.shouldNotContainSuspendMarkers(classCodegen.context)) {
+        if (callee.isSuspend && !irFunction.shouldNotContainSuspendMarkers()) {
             addSuspendMarker(mv, isStartNotEnd = false)
             addInlineMarker(mv, isStartNotEnd = false)
         }

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

@@ -12,6 +12,8 @@ import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.codegen.AsmUtil
 import org.jetbrains.kotlin.codegen.mangleNameIfNeeded
 import org.jetbrains.kotlin.codegen.state.GenerationState
+import org.jetbrains.kotlin.codegen.ClassBuilderMode
+import org.jetbrains.kotlin.codegen.coroutines.SUSPEND_IMPL_NAME_SUFFIX
 import org.jetbrains.kotlin.codegen.visitAnnotableParameterCount
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.descriptors.Modality
@@ -28,6 +30,7 @@ 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.utils.addToStdlib.safeAs
+import org.jetbrains.kotlin.utils.sure
 import org.jetbrains.org.objectweb.asm.MethodVisitor
 import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter
@@ -77,34 +80,56 @@ open class FunctionCodegen(
             generateAnnotationDefaultValueIfNeeded(methodVisitor)
         } else {
             val frameMap = createFrameMapWithReceivers()
-            val irClass = context.suspendFunctionContinuations[irFunction]
-            val element = (irFunction.symbol.descriptor.psiElement
-                ?: context.suspendLambdaToOriginalFunctionMap[irFunction.parent]?.symbol?.descriptor?.psiElement) as? KtElement
-            val continuationClassBuilder = context.continuationClassBuilders[irClass]
             methodVisitor = when {
-                irFunction.isSuspend &&
-                        // We do not generate continuation and state-machine for synthetic accessors, bridges, and delegated members,
-                        // in a sense, they are tail-call
-                        !irFunction.isKnownToBeTailCall() &&
-                        // TODO: We should generate two versions of inline suspend function: one with state-machine and one without
-                        !irFunction.isInline ->
+                irFunction.hasContinuation() -> {
                     generateStateMachineForNamedFunction(
-                        irFunction, classCodegen, methodVisitor, flags, signature, continuationClassBuilder, element!!
+                        irFunction, classCodegen, methodVisitor,
+                        access = flags,
+                        signature = signature,
+                        obtainContinuationClassBuilder = {
+                            context.continuationClassBuilders[continuationClass().attributeOwnerId]!!
+                        },
+                        element = psiElement()
                     )
-                irFunction.isInvokeSuspendOfLambda(context) -> generateStateMachineForLambda(
-                    classCodegen, methodVisitor, flags, signature, element!!
+                }
+                irFunction.isInvokeSuspendOfLambda() -> generateStateMachineForLambda(
+                    classCodegen, methodVisitor, flags, signature, psiElement()
                 )
                 else -> methodVisitor
             }
             ExpressionCodegen(functionView, signature, frameMap, InstructionAdapter(methodVisitor), classCodegen, inlinedInto).generate()
             methodVisitor.visitMaxs(-1, -1)
-            continuationClassBuilder?.done()
+            if (irFunction.hasContinuation()) {
+                context.continuationClassBuilders[continuationClass().attributeOwnerId].sure {
+                    "Could not find continuation class builder for ${continuationClass().render()}"
+                }.done()
+            }
         }
         methodVisitor.visitEnd()
 
         return signature
     }
 
+    private fun psiElement(): KtElement =
+        if (irFunction.isSuspend) irFunction.symbol.descriptor.psiElement as KtElement
+        else context.suspendLambdaToOriginalFunctionMap[irFunction.parentAsClass.attributeOwnerId]!!.symbol.descriptor.psiElement as KtElement
+
+    private fun IrFunction.hasContinuation(): Boolean = isSuspend &&
+            // We do not generate continuation and state-machine for synthetic accessors, bridges, and delegated members,
+            // in a sense, they are tail-call
+            !isKnownToBeTailCall() &&
+            // TODO: We should generate two versions of inline suspend function: one with state-machine and one without
+            !isInline &&
+            // This is suspend lambda parameter of inline function
+            origin != IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA &&
+            // This is just a template for inliner
+            origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE &&
+            // Continuations are generated for suspendImpls
+            parentAsClass.functions.none { it.name.asString() == name.asString() + SUSPEND_IMPL_NAME_SUFFIX }
+
+    private fun continuationClass(): IrClass =
+        irFunction.body!!.statements[0] as IrClass
+
     private fun calculateMethodFlags(isStatic: Boolean): Int {
         if (irFunction.origin == IrDeclarationOrigin.FUNCTION_FOR_DEFAULT_PARAMETER) {
             return Opcodes.ACC_PUBLIC or Opcodes.ACC_SYNTHETIC.let {

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.backend.jvm.codegen
 
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
+import org.jetbrains.kotlin.backend.common.CodegenUtil
 import org.jetbrains.kotlin.backend.common.ir.ir2string
 import org.jetbrains.kotlin.codegen.BaseExpressionCodegen
 import org.jetbrains.kotlin.codegen.ClassBuilder
@@ -16,16 +17,24 @@ import org.jetbrains.kotlin.codegen.inline.*
 import org.jetbrains.kotlin.codegen.serialization.JvmSerializationBindings
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptorWithSource
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.expressions.IrMemberAccessExpression
+import org.jetbrains.kotlin.ir.util.isSuspend
+import org.jetbrains.kotlin.ir.util.nameForIrSerialization
+import org.jetbrains.kotlin.ir.util.parentAsClass
+import org.jetbrains.kotlin.ir.util.render
+import org.jetbrains.kotlin.load.java.JvmAbi
+import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodGenericSignature
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature
+import org.jetbrains.kotlin.resolve.source.getPsi
 import org.jetbrains.org.objectweb.asm.*
 import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter
 import org.jetbrains.org.objectweb.asm.commons.Method
@@ -61,7 +70,8 @@ class IrSourceCompilerForInline(
                 root.signature.asmMethod.descriptor,
                 //compilationContextFunctionDescriptor.isInlineOrInsideInline()
                 false,
-                compilationContextFunctionDescriptor.isSuspend
+                compilationContextFunctionDescriptor.isSuspend,
+                findElement()?.let { CodegenUtil.getLineNumberForElement(it, false) } ?: 0
             )
         }
 
@@ -127,13 +137,20 @@ class IrSourceCompilerForInline(
     override fun getContextLabels(): Set<String> {
         val name = codegen.irFunction.name.asString()
         if (name == INVOKE_SUSPEND_METHOD_NAME) {
-            codegen.context.suspendLambdaToOriginalFunctionMap[codegen.irFunction.parent]?.let {
+            codegen.context.suspendLambdaToOriginalFunctionMap[codegen.irFunction.parentAsClass.attributeOwnerId]?.let {
                 return setOf(it.name.asString())
             }
         }
         return setOf(name)
     }
 
+    // TODO: Find a way to avoid using PSI here
+    override fun reportSuspensionPointInsideMonitor(stackTraceElement: String) {
+        org.jetbrains.kotlin.codegen.coroutines.reportSuspensionPointInsideMonitor(findElement()!!, state, stackTraceElement)
+    }
+
+    private fun findElement() = (callElement.symbol.descriptor.original as? DeclarationDescriptorWithSource)?.source?.getPsi() as? KtElement
+
     internal val isPrimaryCopy: Boolean
         get() = codegen.classCodegen !is FakeClassCodegen
 

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

@@ -18,6 +18,7 @@ import org.jetbrains.kotlin.descriptors.ClassifierDescriptor
 import org.jetbrains.kotlin.descriptors.TypeParameterDescriptor
 import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrPackageFragment
+import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.declarations.IrTypeParameter
 import org.jetbrains.kotlin.ir.symbols.IrTypeParameterSymbol
 import org.jetbrains.kotlin.ir.types.*
@@ -25,6 +26,7 @@ import org.jetbrains.kotlin.ir.types.impl.originalKotlinType
 import org.jetbrains.kotlin.ir.util.defaultType
 import org.jetbrains.kotlin.ir.util.fqNameWhenAvailable
 import org.jetbrains.kotlin.ir.util.isSuspendFunction
+import org.jetbrains.kotlin.ir.util.parentAsClass
 import org.jetbrains.kotlin.load.kotlin.TypeMappingMode
 import org.jetbrains.kotlin.load.kotlin.computeExpandedTypeForInlineClass
 import org.jetbrains.kotlin.load.kotlin.mapBuiltInType

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

@@ -5,12 +5,14 @@
 
 package org.jetbrains.kotlin.backend.jvm.codegen
 
+import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.backend.jvm.ir.eraseTypeParameters
 import org.jetbrains.kotlin.backend.jvm.ir.erasedUpperBound
 import org.jetbrains.kotlin.backend.jvm.lower.inlineclasses.InlineClassAbi
 import org.jetbrains.kotlin.codegen.AsmUtil
 import org.jetbrains.kotlin.codegen.StackValue
 import org.jetbrains.kotlin.ir.types.*
+import org.jetbrains.kotlin.ir.util.parentAsClass
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.org.objectweb.asm.Label
@@ -90,7 +92,8 @@ fun PromisedValue.coerce(target: Type, irTarget: IrType): PromisedValue {
 
     // Boxing and unboxing kotlin.Result leads to CCE in generated code
     val doNotCoerceKotlinResultInContinuation =
-        (codegen.irFunction.isInvokeSuspendOfContinuation(codegen.context) || codegen.irFunction.isInvokeOfSuspendLambda(codegen.context))
+        (codegen.irFunction.parentAsClass.origin == JvmLoweredDeclarationOrigin.CONTINUATION_CLASS ||
+                codegen.irFunction.parentAsClass.origin == JvmLoweredDeclarationOrigin.SUSPEND_LAMBDA)
                 && (irType.isKotlinResult() || irTarget.isKotlinResult())
 
     // Coerce inline classes

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

@@ -5,19 +5,16 @@
 
 package org.jetbrains.kotlin.backend.jvm.lower
 
-import org.jetbrains.kotlin.backend.common.FileLoweringPass
-import org.jetbrains.kotlin.backend.common.descriptors.synthesizedName
+import org.jetbrains.kotlin.backend.common.*
 import org.jetbrains.kotlin.backend.common.ir.*
+import org.jetbrains.kotlin.backend.common.lower.LocalDeclarationsLowering
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
-import org.jetbrains.kotlin.backend.common.lower.parents
 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.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.backend.jvm.codegen.isInlineIrBlock
-import org.jetbrains.kotlin.backend.jvm.codegen.isInvokeOfSuspendCallableReference
 import org.jetbrains.kotlin.codegen.coroutines.*
+import org.jetbrains.kotlin.codegen.inline.coroutines.FOR_INLINE_SUFFIX
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
@@ -43,7 +40,6 @@ import org.jetbrains.kotlin.load.java.JavaVisibilities
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.util.OperatorNameConventions
-import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
 
 internal val addContinuationPhase = makeIrFilePhase(
     ::AddContinuationLowering,
@@ -51,50 +47,53 @@ internal val addContinuationPhase = makeIrFilePhase(
     "Add continuation classes to suspend functions and transform suspend lambdas into continuations"
 )
 
-private object CONTINUATION_CLASS : IrDeclarationOriginImpl("CONTINUATION_CLASS")
-
 private class AddContinuationLowering(private val context: JvmBackendContext) : FileLoweringPass {
+    private val functionsToAdd = mutableMapOf<IrClass, MutableSet<IrFunction>>()
+
     override fun lower(irFile: IrFile) {
-        val suspendLambdas = markSuspendLambdas(irFile)
-        val suspendFunctions = markSuspendFunctions(irFile, suspendLambdas.map { it.function }.toSet())
-        for (lambda in suspendLambdas) {
-            generateContinuationClassForLambda(lambda)
+        val (suspendLambdas, inlineLambdas) = markSuspendLambdas(irFile)
+        transformSuspendFunctions(irFile, (suspendLambdas.map { it.function } + inlineLambdas).toSet())
+        for ((clazz, functions) in functionsToAdd) {
+            for (function in functions) {
+                clazz.declarations.add(function)
+            }
         }
         transformReferencesToSuspendLambdas(irFile, suspendLambdas)
-        transformSuspendCalls(irFile)
-        for (suspendFunction in suspendFunctions) {
-            generateContinuationClassForNamedFunction(suspendFunction)
-        }
     }
 
     private fun transformReferencesToSuspendLambdas(irFile: IrFile, suspendLambdas: List<SuspendLambdaInfo>) {
-        irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
+        for (lambda in suspendLambdas) {
+            (lambda.function.parent as IrDeclarationContainer).declarations.remove(lambda.function)
+        }
+        irFile.transformChildrenVoid(object : IrElementTransformerVoidWithContext() {
             override fun visitFunctionReference(expression: IrFunctionReference): IrExpression {
-                expression.transformChildrenVoid(this)
-
                 if (!expression.isSuspend)
                     return expression
-                val constructor = suspendLambdas.singleOrNull { it.function == expression.symbol.owner }?.constructor
-                    ?: return expression
-                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)
+                val info = suspendLambdas.singleOrNull { it.function == expression.symbol.owner } ?: return expression
+                return context.createIrBuilder(expression.symbol, expression.startOffset, expression.endOffset).run {
+                    val expressionArguments = expression.getArguments().map { it.second }
+                    irBlock {
+                        +generateContinuationClassForLambda(info, currentDeclarationParent)
+                        val constructor = info.constructor
+                        assert(constructor.valueParameters.size == expressionArguments.size + 1) {
+                            "Inconsistency between callable reference to suspend lambda and the corresponding continuation"
+                        }
+                        +irCall(constructor.symbol).apply {
+                            expressionArguments.forEachIndexed { index, argument ->
+                                putValueArgument(index, argument)
+                            }
+                            // Pass null as completion parameter
+                            putValueArgument(expressionArguments.size, irNull())
                         }
                     }
-                }
+                }.also { it.transformChildrenVoid(this) }
             }
         })
     }
 
-    private fun generateContinuationClassForLambda(info: SuspendLambdaInfo) {
+    private fun generateContinuationClassForLambda(info: SuspendLambdaInfo, parent: IrDeclarationParent): IrClass {
         val suspendLambda = context.ir.symbols.suspendLambdaClass.owner
-        suspendLambda.createContinuationClassFor(info.function).apply {
+        return suspendLambda.createContinuationClassFor(parent, JvmLoweredDeclarationOrigin.SUSPEND_LAMBDA).apply {
             copyAttributes(info.reference)
             val functionNClass = context.ir.symbols.getJvmFunctionClass(info.arity + 1)
             superTypes.add(
@@ -120,11 +119,14 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                 .mapNotNull { (i, field) -> if (info.reference.getValueArgument(i) == null) field else null }
             val parametersWithArguments = parametersFields - parametersWithoutArguments
             val constructor = addPrimaryConstructorForLambda(info.arity, info.reference, parametersFields)
-            val secondaryConstructor = addSecondaryConstructorForLambda(constructor)
             val invokeToOverride = functionNClass.functions.single {
                 it.owner.valueParameters.size == info.arity + 1 && it.owner.name.asString() == "invoke"
             }
             val invokeSuspend = addInvokeSuspendForLambda(info.function, parametersFields, receiverField)
+            if (info.capturesCrossinline) {
+                addInvokeSuspendForInlineForLambda(invokeSuspend, info.function, parametersFields, receiverField)
+            }
+            info.function.parentAsClass.declarations.remove(info.function)
             if (info.arity <= 1) {
                 val singleParameterField = receiverField ?: parametersWithoutArguments.singleOrNull()
                 val create = addCreate(constructor, suspendLambda, info, parametersWithArguments, singleParameterField)
@@ -139,9 +141,9 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                 )
             }
 
-            context.suspendLambdaToOriginalFunctionMap[this] = info.function
+            context.suspendLambdaToOriginalFunctionMap[attributeOwnerId as IrFunctionReference] = info.function
 
-            info.constructor = secondaryConstructor
+            info.constructor = constructor
         }
     }
 
@@ -154,46 +156,73 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             it.owner.name.asString() == INVOKE_SUSPEND_METHOD_NAME && it.owner.valueParameters.size == 1 &&
                     it.owner.valueParameters[0].type.isKotlinResult()
         }.owner
-        return addFunctionOverride(superMethod)
-            .also { function ->
-                function.copyTypeParametersFrom(irFunction)
-                function.body = irFunction.body?.deepCopyWithSymbols(function)
-                function.body?.transformChildrenVoid(object : IrElementTransformerVoid() {
-                    override fun visitGetValue(expression: IrGetValue): IrExpression {
-                        if (expression.symbol.owner == irFunction.extensionReceiverParameter) {
-                            assert(receiverField != null)
-                            return IrGetFieldImpl(
-                                expression.startOffset,
-                                expression.endOffset,
-                                receiverField!!.symbol,
-                                receiverField.type
-                            ).also {
-                                it.receiver =
-                                    IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, function.dispatchReceiverParameter!!.symbol)
-                            }
-                        } else if (expression.symbol.owner == irFunction.dispatchReceiverParameter) {
-                            return IrGetValueImpl(expression.startOffset, expression.endOffset, function.dispatchReceiverParameter!!.symbol)
-                        }
-                        val field = fields.find { it.name == expression.symbol.owner.name } ?: return expression
-                        return IrGetFieldImpl(expression.startOffset, expression.endOffset, field.symbol, field.type).also {
-                            it.receiver = IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, function.dispatchReceiverParameter!!.symbol)
-                        }
-                    }
+        return addFunctionOverride(superMethod).also { it.copySuspendLambdaBodyFrom(irFunction, receiverField, fields) }
+    }
 
-                    // If the suspend lambda body contains declarations of other classes (for other lambdas),
-                    // do not rewrite those. In particular, that could lead to rewriting of returns in nested
-                    // lambdas to unintended non-local returns.
-                    override fun visitClass(declaration: IrClass): IrStatement {
-                        return declaration
-                    }
+    private fun IrClass.addInvokeSuspendForInlineForLambda(
+        invokeSuspend: IrFunction,
+        irFunction: IrFunction,
+        fields: List<IrField>,
+        receiverField: IrField?
+    ): IrFunction {
+        return addFunction(
+            INVOKE_SUSPEND_METHOD_NAME + FOR_INLINE_SUFFIX,
+            context.irBuiltIns.anyNType,
+            Modality.FINAL,
+            origin = JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE
+        ).apply {
+            invokeSuspend.valueParameters.mapTo(valueParameters) { it.copyTo(this) }
+        }.also { it.copySuspendLambdaBodyFrom(irFunction, receiverField, fields) }
+    }
 
-                    override fun visitReturn(expression: IrReturn): IrExpression {
-                        val ret = super.visitReturn(expression) as IrReturn
-                        return IrReturnImpl(ret.startOffset, ret.endOffset, ret.type, function.symbol, ret.value)
+    private fun IrSimpleFunction.copySuspendLambdaBodyFrom(
+        irFunction: IrFunction,
+        receiverField: IrField?,
+        fields: List<IrField>
+    ) {
+        copyTypeParametersFrom(irFunction)
+        body = irFunction.body?.deepCopyWithSymbols(this)
+        body?.transformChildrenVoid(object : IrElementTransformerVoid() {
+            override fun visitGetValue(expression: IrGetValue): IrExpression {
+                if (expression.symbol.owner == irFunction.extensionReceiverParameter) {
+                    assert(receiverField != null)
+                    return IrGetFieldImpl(
+                        expression.startOffset,
+                        expression.endOffset,
+                        receiverField!!.symbol,
+                        receiverField.type
+                    ).also {
+                        it.receiver =
+                            IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, dispatchReceiverParameter!!.symbol)
                     }
-                })
-                (irFunction.parent as IrDeclarationContainer).declarations.remove(irFunction)
+                } else if (expression.symbol.owner == irFunction.dispatchReceiverParameter) {
+                    return IrGetValueImpl(expression.startOffset, expression.endOffset, dispatchReceiverParameter!!.symbol)
+                }
+                val field = fields.find { it.name == expression.symbol.owner.name } ?: return expression
+                return IrGetFieldImpl(expression.startOffset, expression.endOffset, field.symbol, field.type).also {
+                    it.receiver = IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, dispatchReceiverParameter!!.symbol)
+                }
             }
+
+            // If the suspend lambda body contains declarations of other classes (for other lambdas),
+            // do not rewrite those. In particular, that could lead to rewriting of returns in nested
+            // lambdas to unintended non-local returns.
+            override fun visitClass(declaration: IrClass): IrStatement {
+                return declaration
+            }
+
+            override fun visitReturn(expression: IrReturn): IrExpression {
+                val ret = super.visitReturn(expression) as IrReturn
+                return IrReturnImpl(ret.startOffset, ret.endOffset, ret.type, symbol, ret.value)
+            }
+        })
+    }
+
+    private fun IrFunction.copyBodyFrom(oldFunction: IrFunction) {
+        val mapping: Map<IrValueParameter, IrValueParameter> =
+            (listOfNotNull(oldFunction.dispatchReceiverParameter, oldFunction.extensionReceiverParameter) + oldFunction.valueParameters)
+                .zip(listOfNotNull(dispatchReceiverParameter, extensionReceiverParameter) + valueParameters).toMap()
+        copyBodyWithParametersMapping(this, oldFunction, mapping)
     }
 
     private fun IrDeclarationContainer.addFunctionOverride(
@@ -311,17 +340,15 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         }
     }
 
-    private fun IrClass.createContinuationClassFor(irFunction: IrFunction): IrClass = buildClass {
-        name = "${irFunction.name}\$Continuation".synthesizedName
-        origin = CONTINUATION_CLASS
+    private fun IrClass.createContinuationClassFor(parent: IrDeclarationParent, newOrigin: IrDeclarationOrigin): IrClass = buildClass {
+        name = Name.special("<Continuation>")
+        origin = newOrigin
         visibility = JavaVisibilities.PACKAGE_VISIBILITY
     }.also { irClass ->
         irClass.createImplicitParameterDeclarationWithWrappedDescriptor()
         irClass.superTypes.add(defaultType)
 
-        val parent = irFunction.parents.firstIsInstance<IrDeclarationContainer>()
         irClass.parent = parent
-        parent.declarations.add(irClass)
     }
 
     // Primary constructor accepts parameters equal to function reference arguments + continuation and sets the fields.
@@ -354,51 +381,37 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             }
         }
 
-    // Secondary constructor accepts parameters equals to arguments of function reference and is used for callable references
-    // TODO: get rid of it and use primary constructor only
-    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) })
-
-            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())
-                }
-            }
-        }
-
     private fun IrFunction.addCompletionValueParameter(): IrValueParameter =
         addValueParameter(SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME, continuationType())
 
     private fun IrFunction.continuationType(): IrType =
         context.ir.symbols.continuationClass.typeWith(returnType).makeNullable()
 
-    private fun generateContinuationClassForNamedFunction(irFunction: IrFunction) {
-        context.ir.symbols.continuationImplClass.owner.createContinuationClassFor(irFunction).apply {
-            val resultField = addField(
-                context.state.languageVersionSettings.dataFieldName(),
-                context.irBuiltIns.anyType,
-                JavaVisibilities.PACKAGE_VISIBILITY
-            )
-            val capturedThisField = irFunction.dispatchReceiverParameter?.let { addField("this\$0", it.type) }
-            val labelField = addField(COROUTINE_LABEL_FIELD_NAME, context.irBuiltIns.intType, JavaVisibilities.PACKAGE_VISIBILITY)
-            addConstructorForNamedFunction(capturedThisField)
-            var function = irFunction
-            if (function is IrSimpleFunction && function.isOverridable && function.body != null) {
-                // Create static method for the suspend state machine method so that reentering the method
-                // does not lead to virtual dispatch to the wrong method.
-                context.suspendFunctionContinuations[function] = this
-                function = createStaticSuspendImpl(function)
+    private fun generateContinuationClassForNamedFunction(
+        irFunction: IrFunction,
+        dispatchReceiverParameter: IrValueParameter?,
+        attributeContainer: IrAttributeContainer
+    ): IrClass {
+        return context.ir.symbols.continuationImplClass.owner
+            .createContinuationClassFor(irFunction, JvmLoweredDeclarationOrigin.CONTINUATION_CLASS)
+            .apply {
+                val resultField = addField(
+                    context.state.languageVersionSettings.dataFieldName(),
+                    context.irBuiltIns.anyType,
+                    JavaVisibilities.PACKAGE_VISIBILITY
+                )
+                val capturedThisField = dispatchReceiverParameter?.let { addField("this\$0", it.type) }
+                val labelField = addField(COROUTINE_LABEL_FIELD_NAME, context.irBuiltIns.intType, JavaVisibilities.PACKAGE_VISIBILITY)
+                addConstructorForNamedFunction(capturedThisField)
+                addInvokeSuspendForNamedFunction(
+                    irFunction,
+                    resultField,
+                    labelField,
+                    capturedThisField,
+                    irFunction.origin == JvmLoweredDeclarationOrigin.SUSPEND_IMPL_STATIC_FUNCTION
+                )
+                copyAttributes(attributeContainer)
             }
-            addInvokeSuspendForNamedFunction(function, resultField, labelField, capturedThisField, function != irFunction)
-            context.suspendFunctionContinuations[function] = this
-        }
     }
 
     private fun IrClass.addConstructorForNamedFunction(capturedThisField: IrField?): IrConstructor = addConstructor {
@@ -433,7 +446,8 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         capturedThisField: IrField?,
         isStaticSuspendImpl: Boolean
     ) {
-        val invokeSuspend = context.ir.symbols.continuationImplClass.owner.functions.single { it.name == Name.identifier(INVOKE_SUSPEND_METHOD_NAME) }
+        val invokeSuspend = context.ir.symbols.continuationImplClass.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]))
@@ -495,7 +509,6 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             copyMetadata = false
         )
         copyBodyToStatic(irFunction, static)
-        (irFunction.parent as IrClass).declarations.add(static)
         // Rewrite the body of the original suspend method to forward to the new static method.
         irFunction.body = context.createIrBuilder(irFunction.symbol).irBlockBody {
             +irReturn(irCall(static).also {
@@ -521,30 +534,78 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
     }
 
     // TODO: Generate two copies of inline suspend functions
-    private fun markSuspendFunctions(irFile: IrFile, suspendLambdas: Set<IrFunction>): Set<IrFunction> {
-        val result = hashSetOf<IrFunction>()
+    private fun transformSuspendFunctions(irFile: IrFile, suspendAndInlineLambdas: Set<IrFunction>) {
+        irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
+            private val functionsStack = arrayListOf<IrFunction>()
+            private val suspendFunctionsCapturingCrossinline = mutableSetOf<IrFunction>()
 
-        irFile.acceptChildrenVoid(object : IrElementVisitorVoid {
-            override fun visitElement(element: IrElement) {
-                element.acceptChildrenVoid(this)
+            override fun visitFunction(declaration: IrFunction): IrStatement {
+                functionsStack.push(declaration)
+                val function = super.visitFunction(declaration) as IrFunction
+                functionsStack.pop()
+                if (skip(function)) return function
+
+                function as IrSimpleFunction
+                if (function in suspendFunctionsCapturingCrossinline) {
+                    val newFunction = buildFun {
+                        name = Name.identifier(function.name.asString() + FOR_INLINE_SUFFIX)
+                        returnType = function.returnType
+                        modality = function.modality
+                        isSuspend = function.isSuspend
+                        origin = JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
+                    }.apply {
+                        dispatchReceiverParameter = function.dispatchReceiverParameter?.copyTo(this)
+                        extensionReceiverParameter = function.extensionReceiverParameter?.copyTo(this)
+                        function.valueParameters.mapTo(valueParameters) { it.copyTo(this) }
+                        copyBodyFrom(function)
+                        copyAttributes(function)
+                    }
+                    functionsToAdd[function.parentAsClass] = mutableSetOf(newFunction)
+                }
+
+                val newFunction = if (function.isOverridable) {
+                    // Create static method for the suspend state machine method so that reentering the method
+                    // does not lead to virtual dispatch to the wrong method.
+                    functionsToAdd.getOrPut(function.parentAsClass) { mutableSetOf() }.add(function)
+                    createStaticSuspendImpl(function)
+                } else function
+
+                newFunction.body = context.createIrBuilder(newFunction.symbol).irBlockBody {
+                    +generateContinuationClassForNamedFunction(
+                        newFunction,
+                        function.dispatchReceiverParameter,
+                        declaration as IrAttributeContainer
+                    )
+                    for (statement in newFunction.body!!.statements) {
+                        +statement
+                    }
+                }
+                return newFunction
             }
 
-            override fun visitFunction(declaration: IrFunction) {
-                super.visitFunction(declaration)
-                if (declaration.isSuspend && declaration !in suspendLambdas && !declaration.isInline &&
-                    !declaration.isInvokeOfSuspendCallableReference()
+            private fun skip(function: IrFunction) =
+                !function.isSuspend || function in suspendAndInlineLambdas || function.isInline || function.body == null ||
+                        function.origin == JvmLoweredDeclarationOrigin.DEFAULT_IMPLS_BRIDGE ||
+                        function.origin == IrDeclarationOrigin.FUNCTION_FOR_DEFAULT_PARAMETER ||
+                        function.parentAsClass.origin == JvmLoweredDeclarationOrigin.FUNCTION_REFERENCE_IMPL
+
+            override fun visitFieldAccess(expression: IrFieldAccessExpression): IrExpression {
+                val result = super.visitFieldAccess(expression)
+                val function = functionsStack.peek() ?: return result
+                if (function.isSuspend &&
+                    expression.symbol.owner.origin == LocalDeclarationsLowering.DECLARATION_ORIGIN_FIELD_FOR_CROSSINLINE_CAPTURED_VALUE
                 ) {
-                    result.add(declaration)
+                    suspendFunctionsCapturingCrossinline += function
                 }
+                return result
             }
         })
-
-        return result
     }
 
-    private fun markSuspendLambdas(irElement: IrElement): List<SuspendLambdaInfo> {
+    private fun markSuspendLambdas(irElement: IrElement): Pair<List<SuspendLambdaInfo>, List<IrFunction>> {
         val suspendLambdas = arrayListOf<SuspendLambdaInfo>()
         val inlineLambdas = mutableSetOf<IrFunctionReference>()
+        val capturesCrossinline = mutableSetOf<IrFunctionReference>()
         irElement.acceptChildrenVoid(object : IrElementVisitorVoid {
             override fun visitElement(element: IrElement) {
                 element.acceptChildrenVoid(this)
@@ -552,6 +613,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
 
             override fun visitCall(expression: IrCall) {
                 val owner = expression.symbol.owner
+                // Collect inline lambdas
                 if (owner.isInline) {
                     for (i in 0 until expression.valueArgumentsCount) {
                         if (owner.valueParameters[i].isNoinline) continue
@@ -565,6 +627,20 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                         }
                     }
                 }
+                // collect lambdas, which capture crossinline
+                for (i in 0 until expression.valueArgumentsCount) {
+                    val valueArgument = expression.getValueArgument(i) as? IrBlock ?: continue
+                    if (valueArgument.origin != IrStatementOrigin.LAMBDA) continue
+                    val reference = valueArgument.statements.last() as? IrFunctionReference ?: continue
+                    for (j in 0 until reference.valueArgumentsCount) {
+                        val getValue = (reference.getValueArgument(j) as? IrGetValue) ?: continue
+                        val parameter = getValue.symbol.owner as? IrValueParameter ?: continue
+                        if (parameter.isCrossinline) {
+                            capturesCrossinline += reference
+                            break
+                        }
+                    }
+                }
                 expression.acceptChildrenVoid(this)
             }
 
@@ -575,32 +651,21 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                     suspendLambdas += SuspendLambdaInfo(
                         expression.symbol.owner,
                         (expression.type as IrSimpleType).arguments.size - 1,
-                        expression
+                        expression,
+                        expression in capturesCrossinline
                     )
                 }
             }
         })
-        return suspendLambdas
+        return suspendLambdas to inlineLambdas.map { it.symbol.owner }
     }
 
-    private fun transformSuspendCalls(irFile: IrFile) {
-        irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
-            private val functionStack = mutableListOf<IrFunction>()
-            override fun visitElement(element: IrElement): IrElement {
-                element.transformChildrenVoid(this)
-                return element
-            }
-
-            override fun visitFunction(declaration: IrFunction): IrStatement {
-                functionStack.push(declaration)
-                val res = super.visitFunction(declaration)
-                functionStack.pop()
-                return res
-            }
-        })
-    }
-
-    private class SuspendLambdaInfo(val function: IrFunction, val arity: Int, val reference: IrFunctionReference) {
+    private class SuspendLambdaInfo(
+        val function: IrFunction,
+        val arity: Int,
+        val reference: IrFunctionReference,
+        val capturesCrossinline: Boolean
+    ) {
         lateinit var constructor: IrConstructor
     }
 }

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

@@ -414,7 +414,4 @@ internal class CallableReferenceLowering(private val context: JvmBackendContext)
             }
         }
     }
-
-    private val currentDeclarationParent
-        get() = allScopes.last { it.irElement is IrDeclarationParent }.irElement as IrDeclarationParent
 }

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/FunctionGenerator.kt

@@ -56,10 +56,6 @@ class FunctionGenerator(declarationGenerator: DeclarationGenerator) : Declaratio
             ?.let {
                 declareSimpleFunctionInner(it, ktElement, IrDeclarationOrigin.FAKE_OVERRIDE).buildWithScope { irFunction ->
                     generateFunctionParameterDeclarationsAndReturnType(irFunction, ktElement, null)
-                    val overridenSymbol = irFunction.overriddenSymbols.first()
-                    if (overridenSymbol.isBound) {
-                        irFunction.copyAttributes(overridenSymbol.owner)
-                    }
                 }
             }
 

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/lazy/IrLazyFunction.kt

@@ -8,10 +8,7 @@ package org.jetbrains.kotlin.ir.declarations.lazy
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibility
-import org.jetbrains.kotlin.ir.declarations.IrAttributeContainer
-import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
-import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
-import org.jetbrains.kotlin.ir.declarations.IrTypeParameter
+import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.symbols.IrPropertySymbol
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.util.DeclarationStubGenerator