JVM_IR: Support inline suspend lambdas

2019-09-16 20:20:10 +03:00
parent 61ae98df40
commit 0e3f0c98e5
7 changed files with 67 additions and 21 deletions
@@ -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,
-    languageVersionSettings: LanguageVersionSettings,
+    private val typeMapper: KotlinTypeMapper,
+    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
+        )
+    }
 }
@@ -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)
                        }
@@ -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,
@@ -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()) +
-                            it.argumentTypes.drop((if (isExtensionLambda) 1 else 0) + capturedVars.size)
-                    ).toTypedArray()
+            ((if (isExtensionLambda) it.argumentTypes.take(1) else emptyList()) +
+                    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
@@ -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,
@@ -1,5 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
 // FILE: test.kt
 // COMMON_COROUTINES_TEST
 // WITH_RUNTIME
@@ -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)