From 862cd5665b1428fd5d3a2aa30962f52b9ea1fce7 Mon Sep 17 00:00:00 2001 From: pyos Date: Fri, 6 Mar 2020 16:14:44 +0100 Subject: [PATCH] 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. --- .../backend/jvm/codegen/IrInlineCodegen.kt | 9 +- .../jvm/codegen/IrSourceCompilerForInline.kt | 23 +++--- .../jvm/lower/AddContinuationLowering.kt | 82 +++++++++++-------- 3 files changed, 60 insertions(+), 54 deletions(-) diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrInlineCodegen.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrInlineCodegen.kt index f9a6374ca02..a98162819e9 100644 --- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrInlineCodegen.kt +++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrInlineCodegen.kt @@ -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 = capturedParameters.keys.toList() - private val loweredMethod = methodSignatureMapper.mapAsmMethod(function.suspendFunctionViewOrStub(context)) + private val loweredMethod = methodSignatureMapper.mapAsmMethod(function) val capturedParamsInDesc: List = 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 diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrSourceCompilerForInline.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrSourceCompilerForInline.kt index b8fb1978869..7574d6bfc31 100644 --- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrSourceCompilerForInline.kt +++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrSourceCompilerForInline.kt @@ -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() diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/AddContinuationLowering.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/AddContinuationLowering.kt index cbbcacb0064..dcc55a29907 100644 --- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/AddContinuationLowering.kt +++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/AddContinuationLowering.kt @@ -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 { 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): 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.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 ` | 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) } }