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)