JVM_IR: add continuations to inline suspend lambdas in the lowering

Delaying suspend view generation until the inliner is not a good idea because it prevents lowerings after AddContinuationLowering from further transforming the declarations.
2020-03-06 16:14:44 +01:00
parent b393f2f680
commit 862cd5665b
3 changed files with 60 additions and 54 deletions
@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.backend.jvm.codegen
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.codegen.*
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
-import org.jetbrains.kotlin.backend.jvm.lower.suspendFunctionViewOrStub
 import org.jetbrains.kotlin.backend.jvm.ir.isInlineParameter
 import org.jetbrains.kotlin.backend.jvm.ir.isLambda
 import org.jetbrains.kotlin.codegen.inline.*
@@ -222,7 +221,7 @@ class IrExpressionLambdaImpl(
    val function: IrFunction,
    private val typeMapper: IrTypeMapper,
    methodSignatureMapper: MethodSignatureMapper,
-    private val context: JvmBackendContext,
+    context: JvmBackendContext,
    isCrossInline: Boolean,
    override val isBoundCallableReference: Boolean,
    override val isExtensionLambda: Boolean
@@ -248,7 +247,7 @@ class IrExpressionLambdaImpl(

    override val capturedVars: List<CapturedParamDesc> = capturedParameters.keys.toList()

-    private val loweredMethod = methodSignatureMapper.mapAsmMethod(function.suspendFunctionViewOrStub(context))
+    private val loweredMethod = methodSignatureMapper.mapAsmMethod(function)

    val capturedParamsInDesc: List<Type> = if (isBoundCallableReference) {
        loweredMethod.argumentTypes.take(1)
@@ -270,9 +269,7 @@ class IrExpressionLambdaImpl(

    override val hasDispatchReceiver: Boolean = false

-    override fun getInlineSuspendLambdaViewDescriptor(): FunctionDescriptor {
-        return function.suspendFunctionViewOrStub(context).descriptor
-    }
+    override fun getInlineSuspendLambdaViewDescriptor(): FunctionDescriptor = function.descriptor

    override fun isCapturedSuspend(desc: CapturedParamDesc, inliningContext: InliningContext): Boolean =
        capturedParameters[desc]?.let { it.isInlineParameter() && it.type.isSuspendFunctionTypeOrSubtype() } == true
@@ -10,7 +10,6 @@ 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.backend.jvm.lower.suspendFunctionView
 import org.jetbrains.kotlin.codegen.BaseExpressionCodegen
 import org.jetbrains.kotlin.codegen.ClassBuilder
 import org.jetbrains.kotlin.codegen.OwnerKind
@@ -100,16 +99,7 @@ class IrSourceCompilerForInline(

    private fun makeInlineNode(function: IrFunction, classCodegen: ClassCodegen, isLambda: Boolean): SMAPAndMethodNode {
        var node: MethodNode? = null
-        // Do not inline the generated state-machine, which was generated to support java interop of inline suspend functions.
-        // Instead, find its $$forInline companion (they share the same attributeOwnerId), which is generated for the inliner to use.
-        val forInlineFunction =
-            if (function.isSuspend)
-                function.parentAsClass.functions.find {
-                    it.name.asString() == function.name.asString() + FOR_INLINE_SUFFIX &&
-                            it.attributeOwnerId == (function as? IrAttributeContainer)?.attributeOwnerId
-                } ?: function.suspendFunctionView(codegen.classCodegen.context)
-            else function
-        val functionCodegen = object : FunctionCodegen(forInlineFunction, classCodegen, codegen.takeIf { isLambda }) {
+        val functionCodegen = object : FunctionCodegen(function, classCodegen, codegen.takeIf { isLambda }) {
            override fun createMethod(flags: Int, signature: JvmMethodGenericSignature): MethodVisitor {
                val asmMethod = signature.asmMethod
                node = MethodNode(
@@ -137,7 +127,16 @@ class IrSourceCompilerForInline(
        asmMethod: Method
    ): SMAPAndMethodNode {
        assert(callableDescriptor == callee.symbol.descriptor.original) { "Expected $callableDescriptor got ${callee.descriptor.original}" }
-        return makeInlineNode(callee, FakeClassCodegen(callee, codegen.classCodegen), false)
+        // Do not inline the generated state-machine, which was generated to support java interop of inline suspend functions.
+        // Instead, find its $$forInline companion (they share the same attributeOwnerId), which is generated for the inliner to use.
+        val forInlineFunction = if (callee.isSuspend)
+            callee.parentAsClass.functions.find {
+                it.name.asString() == callee.name.asString() + FOR_INLINE_SUFFIX &&
+                        it.attributeOwnerId == (callee as IrAttributeContainer).attributeOwnerId
+            } ?: callee
+        else
+            callee
+        return makeInlineNode(forInlineFunction, FakeClassCodegen(forInlineFunction, codegen.classCodegen), false)
    }

    override fun hasFinallyBlocks() = data.hasFinallyBlocks()
@@ -28,6 +28,7 @@ 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.IrCallImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionReferenceImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrGetFieldImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
@@ -72,12 +73,20 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                return super.visitFunction(declaration).also { functionStack.pop() }
            }

+            override fun visitFunctionReference(expression: IrFunctionReference): IrExpression {
+                val transformed = super.visitFunctionReference(expression) as IrFunctionReference
+                // The only references not yet transformed into objects are inline lambdas; the continuation
+                // for those will be taken from the inline functions they are passed to, not the enclosing scope.
+                return transformed.retargetToSuspendView(context, null) {
+                    IrFunctionReferenceImpl(startOffset, endOffset, type, it, typeArgumentsCount, reflectionTarget, origin)
+                }
+            }
+
            override fun visitCall(expression: IrCall): IrExpression {
-                // This is a property, no need to add continuation parameter, since this cannot be suspend call
-                if (functionStack.isEmpty()) return super.visitCall(expression)
-                val caller = functionStack.peek()!!
-                return (super.visitCall(expression) as IrCall)
-                    .createSuspendFunctionCallViewIfNeeded(context, caller)
+                val transformed = super.visitCall(expression) as IrCall
+                return transformed.retargetToSuspendView(context, functionStack.peek() ?: return transformed) {
+                    IrCallImpl(startOffset, endOffset, type, it, origin, superQualifierSymbol)
+                }
            }
        })
    }
@@ -548,7 +557,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
        irFile.accept(object : IrElementTransformer<MutableFlag?> {
            override fun visitClass(declaration: IrClass, data: MutableFlag?): IrStatement {
                declaration.transformDeclarationsFlat {
-                    if (it is IrSimpleFunction && it.isSuspend && it.origin != IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA)
+                    if (it is IrSimpleFunction && it.isSuspend)
                        return@transformDeclarationsFlat transformToView(it)
                    it.accept(this, null)
                    null
@@ -560,7 +569,11 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                val flag = MutableFlag(false)
                function.accept(this, flag)

-                val view = function.suspendFunctionView(context) as IrSimpleFunction
+                val view = function.suspendFunctionViewOrStub(context) as IrSimpleFunction
+                val continuationParameter = view.continuationParameter()
+                val parameterMap = function.explicitParameters.zip(view.explicitParameters.filter { it != continuationParameter }).toMap()
+                view.body = function.moveBodyTo(view, parameterMap)
+
                val result = mutableListOf(view)
                if (function.body == null || !function.hasContinuation()) return result

@@ -621,19 +634,9 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :

 // 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.suspendFunctionView(context: JvmBackendContext): IrFunction {
+private fun IrFunction.suspendFunctionViewOrStub(context: JvmBackendContext): IrFunction {
    if (!isSuspend) return this
-    val stub = suspendFunctionViewOrStub(context)
-    if (stub.body == null) {
-        val continuationParameter = stub.continuationParameter()
-        stub.body = moveBodyTo(stub, explicitParameters.zip(stub.explicitParameters.filter { it != continuationParameter }).toMap())
-    }
-    return stub
-}
-
-internal fun IrFunction.suspendFunctionViewOrStub(context: JvmBackendContext): IrFunction {
-    if (!isSuspend) return this
-    return context.suspendFunctionOriginalToView.getOrPut(suspendFunctionOriginal()) { suspendFunctionStub(context) }
+    return context.suspendFunctionOriginalToView.getOrPut(suspendFunctionOriginal()) { createSuspendFunctionStub(context) }
 }

 internal fun IrFunction.suspendFunctionOriginal(): IrFunction {
@@ -641,12 +644,12 @@ internal fun IrFunction.suspendFunctionOriginal(): IrFunction {
    return attributeOwnerId as IrFunction
 }

-private fun IrFunction.suspendFunctionStub(context: JvmBackendContext): IrFunction {
+private fun IrFunction.createSuspendFunctionStub(context: JvmBackendContext): IrFunction {
    require(this.isSuspend && this is IrSimpleFunction)
    return buildFunWithDescriptorForInlining(descriptor) {
-        updateFrom(this@suspendFunctionStub)
-        name = this@suspendFunctionStub.name
-        origin = this@suspendFunctionStub.origin
+        updateFrom(this@createSuspendFunctionStub)
+        name = this@createSuspendFunctionStub.name
+        origin = this@createSuspendFunctionStub.origin
        returnType = context.irBuiltIns.anyNType
    }.also { function ->
        function.parent = parent
@@ -688,31 +691,38 @@ private fun IrFunction.continuationType(context: JvmBackendContext): IrType {
        context.ir.symbols.continuationClass.typeWith(returnType)
 }

-private fun IrCall.createSuspendFunctionCallViewIfNeeded(context: JvmBackendContext, caller: IrFunction): IrCall {
+private fun <T : IrFunctionAccessExpression> T.retargetToSuspendView(
+    context: JvmBackendContext,
+    caller: IrFunction?,
+    copyWithTargetSymbol: T.(IrSimpleFunctionSymbol) -> T
+): T {
    // Calls inside continuation are already generated with continuation parameter as well as calls to suspendImpls
-    if (!isSuspend || caller.isInvokeSuspendOfContinuation()
+    if (!symbol.owner.isSuspend || caller?.isInvokeSuspendOfContinuation() == true
        || symbol.owner.origin == JvmLoweredDeclarationOrigin.SUSPEND_IMPL_STATIC_FUNCTION
-        || (symbol.owner.valueParameters.lastOrNull()?.name?.asString() == SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME)
+        || symbol.owner.continuationParameter() != null
    ) return this
-    val view = (symbol.owner as IrSimpleFunction).suspendFunctionViewOrStub(context)
+    val view = symbol.owner.suspendFunctionViewOrStub(context)
    if (view == symbol.owner) return this

    // While the new callee technically returns `<original type> | COROUTINE_SUSPENDED`, the latter case is handled
    // by a method visitor so at an IR overview we don't need to consider it.
-    return IrCallImpl(startOffset, endOffset, type, view.symbol, origin, superQualifierSymbol).also {
+    return copyWithTargetSymbol(view.symbol as IrSimpleFunctionSymbol).also {
+        it.copyAttributes(this)
        it.copyTypeArgumentsFrom(this)
        it.dispatchReceiver = dispatchReceiver
        it.extensionReceiver = extensionReceiver
        val continuationIndex = view.continuationParameter()!!.index
        for (i in 0 until valueArgumentsCount) {
-            it.putValueArgument(i + if (i >= continuationIndex) 1 else 0, getValueArgument(i)!!)
+            it.putValueArgument(i + if (i >= continuationIndex) 1 else 0, getValueArgument(i))
+        }
+        if (caller != null) {
+            // At this point the only LOCAL_FUNCTION_FOR_LAMBDAs are inline and crossinline lambdas.
+            val continuation = if (caller.origin == IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA)
+                context.fakeContinuation
+            else
+                IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, caller.continuationParameter()?.symbol
+                    ?: throw AssertionError("${caller.render()} has no continuation; can't call ${symbol.owner.render()}"))
+            it.putValueArgument(continuationIndex, continuation)
        }
-        // At this point the only LOCAL_FUNCTION_FOR_LAMBDAs are inline and crossinline lambdas
-        val continuation = if (caller.origin == IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA)
-            context.fakeContinuation
-        else
-            IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, caller.continuationParameter()?.symbol
-                ?: throw AssertionError("${caller.render()} has no continuation; can't call ${symbol.owner.render()}"))
-        it.putValueArgument(continuationIndex, continuation)
    }
 }