JVM_IR: Support inline suspend lambdas

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

@@ -13,8 +13,10 @@ import org.jetbrains.kotlin.codegen.binding.CalculatedClosure
 import org.jetbrains.kotlin.codegen.binding.CodegenBinding.*
 import org.jetbrains.kotlin.codegen.binding.MutableClosure
 import org.jetbrains.kotlin.codegen.context.EnclosedValueDescriptor
+import org.jetbrains.kotlin.codegen.coroutines.getOrCreateJvmSuspendFunctionView
 import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
 import org.jetbrains.kotlin.config.LanguageVersionSettings
+import org.jetbrains.kotlin.config.isReleaseCoroutines
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.psi.KtCallableReferenceExpression
@@ -190,12 +192,14 @@ abstract class ExpressionLambda(isCrossInline: Boolean) : LambdaInfo(isCrossInli
     override fun generateLambdaBody(sourceCompiler: SourceCompilerForInline, reifiedTypeInliner: ReifiedTypeInliner<*>) {
         node = sourceCompiler.generateLambdaBody(this)
     }
+
+    abstract fun getInlineSuspendLambdaViewDescriptor(): FunctionDescriptor
 }
 
 class PsiExpressionLambda(
     expression: KtExpression,
-    typeMapper: KotlinTypeMapper,
+    private val typeMapper: KotlinTypeMapper,
-    languageVersionSettings: LanguageVersionSettings,
+    private val languageVersionSettings: LanguageVersionSettings,
     isCrossInline: Boolean,
     override val isBoundCallableReference: Boolean
 ) : ExpressionLambda(isCrossInline) {
@@ -292,4 +296,12 @@ class PsiExpressionLambda(
 
     val isPropertyReference: Boolean
         get() = propertyReferenceInfo != null
+
+    override fun getInlineSuspendLambdaViewDescriptor(): FunctionDescriptor {
+        return getOrCreateJvmSuspendFunctionView(
+            invokeMethodDescriptor,
+            languageVersionSettings.isReleaseCoroutines(),
+            typeMapper.bindingContext
+        )
+    }
 }

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

@@ -236,13 +236,16 @@ class MethodInliner(
                     var coroutineDesc = desc
                     val actualInvokeDescriptor: FunctionDescriptor
                     if (info.invokeMethodDescriptor.isSuspend) {
-                        actualInvokeDescriptor = getOrCreateJvmSuspendFunctionView(info.invokeMethodDescriptor, inliningContext.state)
+                        actualInvokeDescriptor = (info as ExpressionLambda).getInlineSuspendLambdaViewDescriptor()
                         val parametersSize = actualInvokeDescriptor.valueParameters.size +
                                 (if (actualInvokeDescriptor.extensionReceiverParameter != null) 1 else 0)
                         // And here we expect invoke(...Ljava/lang/Object;) be replaced with invoke(...Lkotlin/coroutines/Continuation;)
                         // if this does not happen, insert fake continuation, since we could not have one yet.
                         val argumentTypes = Type.getArgumentTypes(desc)
-                        if (argumentTypes.size != parametersSize) {
+                        if (argumentTypes.size != parametersSize &&
+                            // TODO: Workaround IR-related problem. In IR we already have lowered lambda, while in Old BE we don't.
+                            !(inliningContext.root.state.isIrBackend && desc.endsWith("Lkotlin/coroutines/Continuation;)Ljava/lang/Object;"))
+                        ) {
                             addFakeContinuationMarker(this)
                             coroutineDesc = Type.getMethodDescriptor(Type.getReturnType(desc), *argumentTypes, AsmTypes.OBJECT_TYPE)
                         }

compiler/backend/src/org/jetbrains/kotlin/codegen/state/GenerationState.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.
  */
 
@@ -196,7 +196,7 @@ class GenerationState private constructor(
     )
     val bindingContext: BindingContext = bindingTrace.bindingContext
     val mainFunctionDetector = MainFunctionDetector(bindingContext, languageVersionSettings)
-    private val isIrBackend = configuration.get(JVMConfigurationKeys.IR) ?: false
+    val isIrBackend = configuration.get(JVMConfigurationKeys.IR) ?: false
     val typeMapper: KotlinTypeMapper = KotlinTypeMapper(
         this.bindingContext,
         classBuilderMode,

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

@@ -126,7 +126,7 @@ class IrInlineCodegen(
     ): LambdaInfo {
         val referencedFunction = irReference.symbol.owner
         return IrExpressionLambdaImpl(
-            irReference, referencedFunction, codegen.typeMapper, codegen.methodSignatureMapper, parameter.isCrossinline,
+            irReference, referencedFunction, codegen.typeMapper, codegen.methodSignatureMapper, codegen.context, parameter.isCrossinline,
             boundReceiver != null, parameter.type.isExtensionFunctionType
         ).also { lambda ->
             val closureInfo = invocationParamBuilder.addNextValueParameter(type, true, null, parameter.index)
@@ -141,6 +141,7 @@ class IrExpressionLambdaImpl(
     val function: IrFunction,
     private val typeMapper: IrTypeMapper,
     methodSignatureMapper: MethodSignatureMapper,
+    private val context: JvmBackendContext,
     isCrossInline: Boolean,
     override val isBoundCallableReference: Boolean,
     override val isExtensionLambda: Boolean
@@ -170,7 +171,7 @@ class IrExpressionLambdaImpl(
             }
         }
 
-    private val loweredMethod = methodSignatureMapper.mapAsmMethod(function)
+    private val loweredMethod = methodSignatureMapper.mapAsmMethod(function.getOrCreateSuspendFunctionViewIfNeeded(context))
 
     val capturedParamsInDesc: List<Type> =
         loweredMethod.argumentTypes.drop(if (isExtensionLambda) 1 else 0).take(capturedVars.size)
@@ -179,21 +180,23 @@ class IrExpressionLambdaImpl(
         Method(
             it.name,
             it.returnType,
-            (
+            ((if (isExtensionLambda) it.argumentTypes.take(1) else emptyList()) +
-                    (if (isExtensionLambda) it.argumentTypes.take(1) else emptyList()) +
+                    it.argumentTypes.drop((if (isExtensionLambda) 1 else 0) + capturedVars.size)).toTypedArray()
-                            it.argumentTypes.drop((if (isExtensionLambda) 1 else 0) + capturedVars.size)
-                    ).toTypedArray()
         )
     }
 
     override val invokeMethodDescriptor: FunctionDescriptor = function.descriptor
 
     override val hasDispatchReceiver: Boolean = false
+
+    override fun getInlineSuspendLambdaViewDescriptor(): FunctionDescriptor {
+        return function.getOrCreateSuspendFunctionViewIfNeeded(context).descriptor
+    }
 }
 
 fun isInlineIrExpression(argumentExpression: IrExpression) =
     when (argumentExpression) {
-        is IrBlock -> (argumentExpression.origin == IrStatementOrigin.LAMBDA || argumentExpression.origin == IrStatementOrigin.ANONYMOUS_FUNCTION)
+        is IrBlock -> argumentExpression.isInlineIrBlock()
         is IrCallableReference -> true.also {
             assert((0 until argumentExpression.valueArgumentsCount).count { argumentExpression.getValueArgument(it) != null } == 0) {
                 "Expecting 0 value arguments for bounded callable reference: ${argumentExpression.dump()}"
@@ -202,5 +205,7 @@ fun isInlineIrExpression(argumentExpression: IrExpression) =
         else -> false
     }
 
+fun IrBlock.isInlineIrBlock(): Boolean = origin == IrStatementOrigin.LAMBDA || origin == IrStatementOrigin.ANONYMOUS_FUNCTION
+
 fun IrFunction.isInlineFunctionCall(context: JvmBackendContext) =
     (!context.state.isInlineDisabled || typeParameters.any { it.isReified }) && isInline

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

@@ -7,13 +7,17 @@ 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.ir.*
+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.common.pop
 import org.jetbrains.kotlin.backend.common.push
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.backend.jvm.codegen.isInlineIrBlock
 import org.jetbrains.kotlin.codegen.coroutines.*
 import org.jetbrains.kotlin.config.coroutinesPackageFqName
 import org.jetbrains.kotlin.descriptors.Modality
@@ -24,13 +28,18 @@ import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.builders.declarations.*
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
-import org.jetbrains.kotlin.ir.expressions.impl.*
+import org.jetbrains.kotlin.ir.expressions.impl.IrGetFieldImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrReturnImpl
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
 import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
 import org.jetbrains.kotlin.ir.util.*
-import org.jetbrains.kotlin.ir.visitors.*
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
+import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
+import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.load.java.JavaVisibilities
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.types.Variance
@@ -76,7 +85,8 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
 
                 if (!expression.isSuspend)
                     return expression
-                val constructor = suspendLambdas.single { it.function == expression.symbol.owner }.constructor
+                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"
@@ -407,15 +417,34 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
 
     private fun markSuspendLambdas(irElement: IrElement): List<SuspendLambdaInfo> {
         val suspendLambdas = arrayListOf<SuspendLambdaInfo>()
+        val inlineLambdas = mutableSetOf<IrFunctionReference>()
         irElement.acceptChildrenVoid(object : IrElementVisitorVoid {
             override fun visitElement(element: IrElement) {
                 element.acceptChildrenVoid(this)
             }
 
+            override fun visitCall(expression: IrCall) {
+                val owner = expression.symbol.owner
+                if (owner.isInline) {
+                    for (i in 0 until expression.valueArgumentsCount) {
+                        if (owner.valueParameters[i].isNoinline) continue
+
+                        val valueArgument = expression.getValueArgument(i) ?: continue
+                        if (valueArgument is IrBlock && valueArgument.isInlineIrBlock()) {
+                            assert(valueArgument !is IrCallableReference) {
+                                "callable references should be lowered to function references"
+                            }
+                            inlineLambdas += valueArgument.statements.filterIsInstance<IrFunctionReference>().single()
+                        }
+                    }
+                }
+                expression.acceptChildrenVoid(this)
+            }
+
             override fun visitFunctionReference(expression: IrFunctionReference) {
                 expression.acceptChildrenVoid(this)
 
-                if (expression.isSuspend) {
+                if (expression.isSuspend && expression !in inlineLambdas) {
                     suspendLambdas += SuspendLambdaInfo(
                         expression.symbol.owner,
                         (expression.type as IrSimpleType).arguments.size - 1,

compiler/testData/codegen/boxInline/suspend/inlineSuspendOfSuspend.kt

@@ -1,5 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
 // FILE: test.kt
 // COMMON_COROUTINES_TEST
 // WITH_RUNTIME

compiler/testData/codegen/bytecodeText/coroutines/destructuringInLambda/inlineSeparateFiles.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 data class A(val x: String, val y: String)
 
 suspend inline fun foo(a: A, block: suspend (A) -> String): String = block(a)