mirror of
https://github.com/jlengrand/kotlin.git
synced 2026-04-25 00:21:28 +00:00
JVM_IR: move SuspendLambdaLowering to a separate file
This commit is contained in:
pyos
@@ -16,16 +16,12 @@ import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
|
||||
import org.jetbrains.kotlin.backend.jvm.codegen.continuationParameter
|
||||
import org.jetbrains.kotlin.backend.jvm.codegen.hasContinuation
|
||||
import org.jetbrains.kotlin.backend.jvm.codegen.isInvokeSuspendOfContinuation
|
||||
import org.jetbrains.kotlin.backend.jvm.codegen.isReadOfCrossinline
|
||||
import org.jetbrains.kotlin.backend.jvm.ir.IrInlineReferenceLocator
|
||||
import org.jetbrains.kotlin.backend.jvm.ir.defaultValue
|
||||
import org.jetbrains.kotlin.backend.jvm.ir.isStaticInlineClassReplacement
|
||||
import org.jetbrains.kotlin.backend.jvm.localDeclarationsPhase
|
||||
import org.jetbrains.kotlin.codegen.coroutines.*
|
||||
import org.jetbrains.kotlin.codegen.inline.coroutines.FOR_INLINE_SUFFIX
|
||||
import org.jetbrains.kotlin.descriptors.Modality
|
||||
import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
|
||||
import org.jetbrains.kotlin.ir.IrElement
|
||||
import org.jetbrains.kotlin.ir.IrStatement
|
||||
import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
|
||||
import org.jetbrains.kotlin.ir.builders.*
|
||||
@@ -39,21 +35,13 @@ import org.jetbrains.kotlin.ir.types.*
|
||||
import org.jetbrains.kotlin.ir.util.*
|
||||
import org.jetbrains.kotlin.ir.visitors.IrElementTransformer
|
||||
import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
|
||||
import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
|
||||
import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
|
||||
import org.jetbrains.kotlin.load.java.JavaDescriptorVisibilities
|
||||
import org.jetbrains.kotlin.name.Name
|
||||
import org.jetbrains.kotlin.name.SpecialNames
|
||||
import org.jetbrains.kotlin.util.OperatorNameConventions
|
||||
import org.jetbrains.kotlin.utils.DFS
|
||||
import org.jetbrains.org.objectweb.asm.Type
|
||||
|
||||
internal val suspendLambdaPhase = makeIrFilePhase(
|
||||
::SuspendLambdaLowering,
|
||||
"SuspendLambda",
|
||||
"Transform suspend lambdas into continuation classes"
|
||||
)
|
||||
|
||||
internal val addContinuationPhase = makeIrFilePhase(
|
||||
::AddContinuationLowering,
|
||||
"AddContinuation",
|
||||
@@ -61,261 +49,7 @@ internal val addContinuationPhase = makeIrFilePhase(
|
||||
prerequisite = setOf(suspendLambdaPhase, localDeclarationsPhase, tailCallOptimizationPhase)
|
||||
)
|
||||
|
||||
private abstract class SuspendLoweringBase(protected val context: JvmBackendContext) : FileLoweringPass {
|
||||
protected fun IrClass.addFunctionOverride(function: IrSimpleFunction): IrSimpleFunction =
|
||||
addFunction(function.name.asString(), function.returnType).apply {
|
||||
overriddenSymbols += function.symbol
|
||||
valueParameters = function.valueParameters.map { it.copyTo(this) }
|
||||
}
|
||||
|
||||
protected fun IrClass.addFunctionOverride(
|
||||
function: IrSimpleFunction,
|
||||
makeBody: IrBlockBodyBuilder.(IrFunction) -> Unit
|
||||
): IrSimpleFunction =
|
||||
addFunctionOverride(function).apply {
|
||||
body = context.createIrBuilder(symbol).irBlockBody { makeBody(this@apply) }
|
||||
}
|
||||
|
||||
protected fun IrSimpleFunction.generateErrorForInlineBody() {
|
||||
val message = "This is a stub representing a copy of a suspend method without the state machine " +
|
||||
"(used by the inliner). Since the difference is at the bytecode level, the body is " +
|
||||
"still on the original function. Use suspendForInlineToOriginal() to retrieve it."
|
||||
body = IrExpressionBodyImpl(startOffset, endOffset, IrErrorExpressionImpl(startOffset, endOffset, returnType, message))
|
||||
}
|
||||
|
||||
protected fun IrFunction.addCompletionValueParameter(): IrValueParameter =
|
||||
addValueParameter(SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME, continuationType())
|
||||
|
||||
protected fun IrFunction.continuationType(): IrType =
|
||||
context.ir.symbols.continuationClass.typeWith(returnType).makeNullable()
|
||||
}
|
||||
|
||||
private fun IrFunction.capturesCrossinline(): Boolean {
|
||||
var result = false
|
||||
accept(object : IrElementVisitorVoid {
|
||||
override fun visitElement(element: IrElement) {
|
||||
if (!result) element.acceptChildren(this, null)
|
||||
}
|
||||
|
||||
override fun visitFunction(declaration: IrFunction) {
|
||||
functions.add(declaration)
|
||||
super.visitFunction(declaration)
|
||||
functions.remove(declaration)
|
||||
}
|
||||
|
||||
override fun visitGetValue(expression: IrGetValue) {
|
||||
result = result || (expression.isReadOfCrossinline() && expression.symbol.owner.parent !in functions)
|
||||
}
|
||||
|
||||
private val functions = mutableSetOf<IrFunction>()
|
||||
}, null)
|
||||
return result
|
||||
}
|
||||
|
||||
private class SuspendLambdaLowering(context: JvmBackendContext) : SuspendLoweringBase(context) {
|
||||
override fun lower(irFile: IrFile) {
|
||||
val inlineReferences = IrInlineReferenceLocator.scan(context, irFile)
|
||||
irFile.transformChildrenVoid(object : IrElementTransformerVoidWithContext() {
|
||||
private fun IrFunctionReference.shouldBeTreatedAsSuspendLambda() =
|
||||
isSuspend && (origin == IrStatementOrigin.LAMBDA || origin == IrStatementOrigin.SUSPEND_CONVERSION)
|
||||
|
||||
override fun visitBlock(expression: IrBlock): IrExpression {
|
||||
val reference = expression.statements.lastOrNull() as? IrFunctionReference ?: return super.visitBlock(expression)
|
||||
if (reference.shouldBeTreatedAsSuspendLambda() && reference !in inlineReferences) {
|
||||
assert(expression.statements.size == 2 && expression.statements[0] is IrFunction)
|
||||
expression.transformChildrenVoid(this)
|
||||
val parent = currentDeclarationParent ?: error("No current declaration parent at ${reference.dump()}")
|
||||
return generateAnonymousObjectForLambda(reference, parent)
|
||||
}
|
||||
return super.visitBlock(expression)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
private fun generateAnonymousObjectForLambda(reference: IrFunctionReference, parent: IrDeclarationParent) =
|
||||
context.createIrBuilder(reference.symbol).irBlock(reference.startOffset, reference.endOffset) {
|
||||
assert(reference.getArguments().isEmpty()) { "lambda with bound arguments: ${reference.render()}" }
|
||||
val continuation = generateContinuationClassForLambda(reference, parent)
|
||||
+continuation
|
||||
+irCall(continuation.constructors.single().symbol).apply {
|
||||
// Pass null as completion parameter
|
||||
putValueArgument(0, irNull())
|
||||
}
|
||||
}
|
||||
|
||||
private fun generateContinuationClassForLambda(reference: IrFunctionReference, parent: IrDeclarationParent): IrClass =
|
||||
context.irFactory.buildClass {
|
||||
name = SpecialNames.NO_NAME_PROVIDED
|
||||
origin = JvmLoweredDeclarationOrigin.SUSPEND_LAMBDA
|
||||
visibility = DescriptorVisibilities.LOCAL
|
||||
}.apply {
|
||||
this.parent = parent
|
||||
createImplicitParameterDeclarationWithWrappedDescriptor()
|
||||
copyAttributes(reference)
|
||||
|
||||
val function = reference.symbol.owner
|
||||
val isRestricted = reference.symbol.owner.extensionReceiverParameter?.type?.classOrNull?.owner?.annotations?.any {
|
||||
it.type.classOrNull?.signature == IdSignature.PublicSignature("kotlin.coroutines", "RestrictsSuspension", null, 0)
|
||||
} == true
|
||||
val suspendLambda =
|
||||
if (isRestricted) context.ir.symbols.restrictedSuspendLambdaClass.owner
|
||||
else context.ir.symbols.suspendLambdaClass.owner
|
||||
val arity = (reference.type as IrSimpleType).arguments.size - 1
|
||||
val functionNClass = context.ir.symbols.getJvmFunctionClass(arity + 1)
|
||||
superTypes += suspendLambda.defaultType
|
||||
superTypes += functionNClass.typeWith(
|
||||
function.explicitParameters.subList(0, arity).map { it.type }
|
||||
+ function.continuationType()
|
||||
+ context.irBuiltIns.anyNType
|
||||
)
|
||||
|
||||
addField(COROUTINE_LABEL_FIELD_NAME, context.irBuiltIns.intType, JavaDescriptorVisibilities.PACKAGE_VISIBILITY)
|
||||
|
||||
val parametersFields = function.explicitParameters.map {
|
||||
addField {
|
||||
// Rename `$this` to avoid being caught by inlineCodegenUtils.isCapturedFieldName()
|
||||
name = if (it.index < 0) Name.identifier("p\$") else it.name
|
||||
type = it.type
|
||||
origin = LocalDeclarationsLowering.DECLARATION_ORIGIN_FIELD_FOR_CAPTURED_VALUE
|
||||
isFinal = false
|
||||
visibility = if (it.index < 0) DescriptorVisibilities.PRIVATE else JavaDescriptorVisibilities.PACKAGE_VISIBILITY
|
||||
}
|
||||
}
|
||||
val constructor = addPrimaryConstructorForLambda(suspendLambda, arity)
|
||||
val invokeToOverride = functionNClass.functions.single {
|
||||
it.owner.valueParameters.size == arity + 1 && it.owner.name.asString() == "invoke"
|
||||
}
|
||||
val createToOverride = suspendLambda.symbol.functions.singleOrNull {
|
||||
it.owner.valueParameters.size == arity + 1 && it.owner.name.asString() == "create"
|
||||
}
|
||||
val invokeSuspend = addInvokeSuspendForLambda(function, parametersFields)
|
||||
if (function.capturesCrossinline()) {
|
||||
addInvokeSuspendForInlineForLambda(invokeSuspend)
|
||||
}
|
||||
if (createToOverride != null) {
|
||||
addInvokeCallingCreate(addCreate(constructor, createToOverride, parametersFields), invokeSuspend, invokeToOverride)
|
||||
} else {
|
||||
addInvokeCallingConstructor(constructor, invokeSuspend, invokeToOverride, parametersFields)
|
||||
}
|
||||
|
||||
context.suspendLambdaToOriginalFunctionMap[attributeOwnerId as IrFunctionReference] = function
|
||||
}
|
||||
|
||||
private fun IrClass.addInvokeSuspendForLambda(irFunction: IrFunction, fields: List<IrField>): IrSimpleFunction {
|
||||
val superMethod = context.ir.symbols.suspendLambdaClass.functions.single {
|
||||
it.owner.name.asString() == INVOKE_SUSPEND_METHOD_NAME && it.owner.valueParameters.size == 1 &&
|
||||
it.owner.valueParameters[0].type.isKotlinResult()
|
||||
}.owner
|
||||
return addFunctionOverride(superMethod).apply {
|
||||
body = irFunction.moveBodyTo(this, mapOf())?.transform(object : IrElementTransformerVoid() {
|
||||
override fun visitGetValue(expression: IrGetValue): IrExpression {
|
||||
val parameter = (expression.symbol.owner as? IrValueParameter)?.takeIf { it.parent == irFunction }
|
||||
?: return expression
|
||||
val field = fields[parameter.index + if (irFunction.extensionReceiverParameter != null) 1 else 0]
|
||||
val receiver = IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, dispatchReceiverParameter!!.symbol)
|
||||
return IrGetFieldImpl(expression.startOffset, expression.endOffset, field.symbol, field.type, receiver)
|
||||
}
|
||||
}, null)
|
||||
}
|
||||
}
|
||||
|
||||
private fun IrClass.addInvokeSuspendForInlineForLambda(invokeSuspend: IrSimpleFunction): IrSimpleFunction {
|
||||
return addFunction(
|
||||
INVOKE_SUSPEND_METHOD_NAME + FOR_INLINE_SUFFIX,
|
||||
context.irBuiltIns.anyNType,
|
||||
Modality.FINAL,
|
||||
origin = JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE
|
||||
).apply {
|
||||
copyAttributes(invokeSuspend)
|
||||
generateErrorForInlineBody()
|
||||
valueParameters = invokeSuspend.valueParameters.map { it.copyTo(this) }
|
||||
}
|
||||
}
|
||||
|
||||
// Invoke function in lambdas is responsible for
|
||||
// 1) calling `create`
|
||||
// 2) starting newly created coroutine by calling `invokeSuspend`.
|
||||
// Thus, it creates a clone of suspend lambda and starts it.
|
||||
// TODO: fix the generic signature -- type parameters of FunctionN should be substituted
|
||||
private fun IrClass.addInvokeCallingCreate(
|
||||
create: IrFunction,
|
||||
invokeSuspend: IrSimpleFunction,
|
||||
invokeToOverride: IrSimpleFunctionSymbol
|
||||
) = addFunctionOverride(invokeToOverride.owner) { function ->
|
||||
val newlyCreatedObject = irCall(create).also { createCall ->
|
||||
createCall.dispatchReceiver = irGet(function.dispatchReceiverParameter!!)
|
||||
for ((index, param) in function.valueParameters.withIndex()) {
|
||||
createCall.putValueArgument(index, irGet(param))
|
||||
}
|
||||
}
|
||||
+irReturn(callInvokeSuspend(invokeSuspend, irImplicitCast(newlyCreatedObject, defaultType)))
|
||||
}
|
||||
|
||||
// Same as above, but with `create` inlined. `create` is only defined in `SuspendLambda` in unary and binary
|
||||
// versions; for other lambdas, there's no point in generating a non-overriding `create` separately.
|
||||
private fun IrClass.addInvokeCallingConstructor(
|
||||
constructor: IrFunction,
|
||||
invokeSuspend: IrSimpleFunction,
|
||||
invokeToOverride: IrSimpleFunctionSymbol,
|
||||
fieldsForUnbound: List<IrField>
|
||||
) = addFunctionOverride(invokeToOverride.owner) { function ->
|
||||
+irReturn(callInvokeSuspend(invokeSuspend, cloneLambda(function, constructor, fieldsForUnbound)))
|
||||
}
|
||||
|
||||
private fun IrClass.addCreate(
|
||||
constructor: IrFunction,
|
||||
createToOverride: IrSimpleFunctionSymbol,
|
||||
fieldsForUnbound: List<IrField>
|
||||
) = addFunctionOverride(createToOverride.owner) { function ->
|
||||
+irReturn(cloneLambda(function, constructor, fieldsForUnbound))
|
||||
}
|
||||
|
||||
private fun IrBlockBodyBuilder.cloneLambda(
|
||||
scope: IrFunction,
|
||||
constructor: IrFunction,
|
||||
fieldsForUnbound: List<IrField>
|
||||
): IrExpression {
|
||||
val constructorCall = irCall(constructor).also {
|
||||
for (typeParameter in constructor.parentAsClass.typeParameters) {
|
||||
it.putTypeArgument(typeParameter.index, typeParameter.defaultType)
|
||||
}
|
||||
it.putValueArgument(0, irGet(scope.valueParameters.last()))
|
||||
}
|
||||
if (fieldsForUnbound.isEmpty()) {
|
||||
return constructorCall
|
||||
}
|
||||
val result = irTemporary(constructorCall, "result")
|
||||
for ((index, field) in fieldsForUnbound.withIndex()) {
|
||||
+irSetField(irGet(result), field, irGet(scope.valueParameters[index]))
|
||||
}
|
||||
return irGet(result)
|
||||
}
|
||||
|
||||
private fun IrBlockBodyBuilder.callInvokeSuspend(invokeSuspend: IrSimpleFunction, lambda: IrExpression): IrExpression =
|
||||
irCallOp(invokeSuspend.symbol, invokeSuspend.returnType, lambda, irUnit())
|
||||
|
||||
// Primary constructor accepts parameters equal to function reference arguments + continuation and sets the fields.
|
||||
private fun IrClass.addPrimaryConstructorForLambda(superClass: IrClass, arity: Int): IrConstructor =
|
||||
addConstructor {
|
||||
origin = JvmLoweredDeclarationOrigin.SUSPEND_LAMBDA
|
||||
isPrimary = true
|
||||
returnType = defaultType
|
||||
visibility = DescriptorVisibilities.LOCAL
|
||||
}.also { constructor ->
|
||||
val completionParameterSymbol = constructor.addCompletionValueParameter()
|
||||
val superClassConstructor = superClass.constructors.single {
|
||||
it.valueParameters.size == 2 && it.valueParameters[0].type.isInt() && it.valueParameters[1].type.isNullableContinuation()
|
||||
}
|
||||
constructor.body = context.createIrBuilder(constructor.symbol).irBlockBody {
|
||||
+irDelegatingConstructorCall(superClassConstructor).also {
|
||||
it.putValueArgument(0, irInt(arity + 1))
|
||||
it.putValueArgument(1, irGet(completionParameterSymbol))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private class AddContinuationLowering(context: JvmBackendContext) : SuspendLoweringBase(context) {
|
||||
private class AddContinuationLowering(context: JvmBackendContext) : SuspendLoweringUtils(context), FileLoweringPass {
|
||||
override fun lower(irFile: IrFile) {
|
||||
// This should be done after converting lambdas into classes to avoid breaking the invariant that
|
||||
// each lambda is referenced at most once while creating `$$forInline` methods.
|
||||
@@ -619,7 +353,7 @@ private fun IrSimpleFunction.isOrOverridesDefaultParameterStub(): Boolean =
|
||||
{ it.origin == IrDeclarationOrigin.FUNCTION_FOR_DEFAULT_PARAMETER }
|
||||
)
|
||||
|
||||
val defaultImplsOrigins = setOf(
|
||||
private val defaultImplsOrigins = setOf(
|
||||
IrDeclarationOrigin.FUNCTION_FOR_DEFAULT_PARAMETER,
|
||||
JvmLoweredDeclarationOrigin.DEFAULT_IMPLS_WITH_MOVED_RECEIVERS,
|
||||
JvmLoweredDeclarationOrigin.DEFAULT_IMPLS_WITH_MOVED_RECEIVERS_SYNTHETIC,
|
||||
|
||||
@@ -0,0 +1,305 @@
|
||||
/*
|
||||
* Copyright 2010-2020 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.
|
||||
*/
|
||||
|
||||
package org.jetbrains.kotlin.backend.jvm.lower
|
||||
|
||||
import org.jetbrains.kotlin.backend.common.FileLoweringPass
|
||||
import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
|
||||
import org.jetbrains.kotlin.backend.common.ir.copyTo
|
||||
import org.jetbrains.kotlin.backend.common.ir.createImplicitParameterDeclarationWithWrappedDescriptor
|
||||
import org.jetbrains.kotlin.backend.common.ir.isSuspend
|
||||
import org.jetbrains.kotlin.backend.common.ir.moveBodyTo
|
||||
import org.jetbrains.kotlin.backend.common.lower.LocalDeclarationsLowering
|
||||
import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
|
||||
import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
|
||||
import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
|
||||
import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
|
||||
import org.jetbrains.kotlin.backend.jvm.codegen.isReadOfCrossinline
|
||||
import org.jetbrains.kotlin.backend.jvm.ir.IrInlineReferenceLocator
|
||||
import org.jetbrains.kotlin.codegen.coroutines.COROUTINE_LABEL_FIELD_NAME
|
||||
import org.jetbrains.kotlin.codegen.coroutines.INVOKE_SUSPEND_METHOD_NAME
|
||||
import org.jetbrains.kotlin.codegen.coroutines.SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME
|
||||
import org.jetbrains.kotlin.codegen.inline.coroutines.FOR_INLINE_SUFFIX
|
||||
import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
|
||||
import org.jetbrains.kotlin.descriptors.Modality
|
||||
import org.jetbrains.kotlin.ir.IrElement
|
||||
import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
|
||||
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.IrErrorExpressionImpl
|
||||
import org.jetbrains.kotlin.ir.expressions.impl.IrExpressionBodyImpl
|
||||
import org.jetbrains.kotlin.ir.expressions.impl.IrGetFieldImpl
|
||||
import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
|
||||
import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
|
||||
import org.jetbrains.kotlin.ir.types.*
|
||||
import org.jetbrains.kotlin.ir.util.*
|
||||
import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
|
||||
import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
|
||||
import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
|
||||
import org.jetbrains.kotlin.load.java.JavaDescriptorVisibilities
|
||||
import org.jetbrains.kotlin.name.Name
|
||||
import org.jetbrains.kotlin.name.SpecialNames
|
||||
|
||||
internal val suspendLambdaPhase = makeIrFilePhase(
|
||||
::SuspendLambdaLowering,
|
||||
"SuspendLambda",
|
||||
"Transform suspend lambdas into continuation classes"
|
||||
)
|
||||
|
||||
private fun IrFunction.capturesCrossinline(): Boolean {
|
||||
var result = false
|
||||
accept(object : IrElementVisitorVoid {
|
||||
override fun visitElement(element: IrElement) {
|
||||
if (!result) element.acceptChildren(this, null)
|
||||
}
|
||||
|
||||
override fun visitFunction(declaration: IrFunction) {
|
||||
functions.add(declaration)
|
||||
super.visitFunction(declaration)
|
||||
functions.remove(declaration)
|
||||
}
|
||||
|
||||
override fun visitGetValue(expression: IrGetValue) {
|
||||
result = result || (expression.isReadOfCrossinline() && expression.symbol.owner.parent !in functions)
|
||||
}
|
||||
|
||||
private val functions = mutableSetOf<IrFunction>()
|
||||
}, null)
|
||||
return result
|
||||
}
|
||||
|
||||
internal abstract class SuspendLoweringUtils(protected val context: JvmBackendContext) {
|
||||
protected fun IrClass.addFunctionOverride(function: IrSimpleFunction): IrSimpleFunction =
|
||||
addFunction(function.name.asString(), function.returnType).apply {
|
||||
overriddenSymbols = listOf(function.symbol)
|
||||
valueParameters = function.valueParameters.map { it.copyTo(this) }
|
||||
}
|
||||
|
||||
protected fun IrClass.addFunctionOverride(
|
||||
function: IrSimpleFunction,
|
||||
makeBody: IrBlockBodyBuilder.(IrFunction) -> Unit
|
||||
): IrSimpleFunction =
|
||||
addFunctionOverride(function).apply {
|
||||
body = context.createIrBuilder(symbol).irBlockBody { makeBody(this@apply) }
|
||||
}
|
||||
|
||||
protected fun IrSimpleFunction.generateErrorForInlineBody() {
|
||||
val message = "This is a stub representing a copy of a suspend method without the state machine " +
|
||||
"(used by the inliner). Since the difference is at the bytecode level, the body is " +
|
||||
"still on the original function. Use suspendForInlineToOriginal() to retrieve it."
|
||||
body = IrExpressionBodyImpl(startOffset, endOffset, IrErrorExpressionImpl(startOffset, endOffset, returnType, message))
|
||||
}
|
||||
|
||||
protected fun IrFunction.addCompletionValueParameter(): IrValueParameter =
|
||||
addValueParameter(SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME, continuationType())
|
||||
|
||||
protected fun IrFunction.continuationType(): IrType =
|
||||
context.ir.symbols.continuationClass.typeWith(returnType).makeNullable()
|
||||
}
|
||||
|
||||
private class SuspendLambdaLowering(context: JvmBackendContext) : SuspendLoweringUtils(context), FileLoweringPass {
|
||||
override fun lower(irFile: IrFile) {
|
||||
val inlineReferences = IrInlineReferenceLocator.scan(context, irFile)
|
||||
irFile.transformChildrenVoid(object : IrElementTransformerVoidWithContext() {
|
||||
private fun IrFunctionReference.shouldBeTreatedAsSuspendLambda() =
|
||||
isSuspend && (origin == IrStatementOrigin.LAMBDA || origin == IrStatementOrigin.SUSPEND_CONVERSION)
|
||||
|
||||
override fun visitBlock(expression: IrBlock): IrExpression {
|
||||
val reference = expression.statements.lastOrNull() as? IrFunctionReference ?: return super.visitBlock(expression)
|
||||
if (reference.shouldBeTreatedAsSuspendLambda() && reference !in inlineReferences) {
|
||||
assert(expression.statements.size == 2 && expression.statements[0] is IrFunction)
|
||||
expression.transformChildrenVoid(this)
|
||||
val parent = currentDeclarationParent ?: error("No current declaration parent at ${reference.dump()}")
|
||||
return generateAnonymousObjectForLambda(reference, parent)
|
||||
}
|
||||
return super.visitBlock(expression)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
private fun generateAnonymousObjectForLambda(reference: IrFunctionReference, parent: IrDeclarationParent) =
|
||||
context.createIrBuilder(reference.symbol).irBlock(reference.startOffset, reference.endOffset) {
|
||||
assert(reference.getArguments().isEmpty()) { "lambda with bound arguments: ${reference.render()}" }
|
||||
val continuation = generateContinuationClassForLambda(reference, parent)
|
||||
+continuation
|
||||
+irCall(continuation.constructors.single().symbol).apply {
|
||||
// Pass null as completion parameter
|
||||
putValueArgument(0, irNull())
|
||||
}
|
||||
}
|
||||
|
||||
private fun generateContinuationClassForLambda(reference: IrFunctionReference, parent: IrDeclarationParent): IrClass =
|
||||
context.irFactory.buildClass {
|
||||
name = SpecialNames.NO_NAME_PROVIDED
|
||||
origin = JvmLoweredDeclarationOrigin.SUSPEND_LAMBDA
|
||||
visibility = DescriptorVisibilities.LOCAL
|
||||
}.apply {
|
||||
this.parent = parent
|
||||
createImplicitParameterDeclarationWithWrappedDescriptor()
|
||||
copyAttributes(reference)
|
||||
|
||||
val function = reference.symbol.owner
|
||||
val isRestricted = reference.symbol.owner.extensionReceiverParameter?.type?.classOrNull?.owner?.annotations?.any {
|
||||
it.type.classOrNull?.signature == IdSignature.PublicSignature("kotlin.coroutines", "RestrictsSuspension", null, 0)
|
||||
} == true
|
||||
val suspendLambda =
|
||||
if (isRestricted) context.ir.symbols.restrictedSuspendLambdaClass.owner
|
||||
else context.ir.symbols.suspendLambdaClass.owner
|
||||
val arity = (reference.type as IrSimpleType).arguments.size - 1
|
||||
val functionNClass = context.ir.symbols.getJvmFunctionClass(arity + 1)
|
||||
val functionNType = functionNClass.typeWith(
|
||||
function.explicitParameters.subList(0, arity).map { it.type }
|
||||
+ function.continuationType()
|
||||
+ context.irBuiltIns.anyNType
|
||||
)
|
||||
superTypes = listOf(suspendLambda.defaultType, functionNType)
|
||||
|
||||
addField(COROUTINE_LABEL_FIELD_NAME, context.irBuiltIns.intType, JavaDescriptorVisibilities.PACKAGE_VISIBILITY)
|
||||
|
||||
val parametersFields = function.explicitParameters.map {
|
||||
addField {
|
||||
// Rename `$this` to avoid being caught by inlineCodegenUtils.isCapturedFieldName()
|
||||
name = if (it.index < 0) Name.identifier("p\$") else it.name
|
||||
type = it.type
|
||||
origin = LocalDeclarationsLowering.DECLARATION_ORIGIN_FIELD_FOR_CAPTURED_VALUE
|
||||
isFinal = false
|
||||
visibility = if (it.index < 0) DescriptorVisibilities.PRIVATE else JavaDescriptorVisibilities.PACKAGE_VISIBILITY
|
||||
}
|
||||
}
|
||||
val constructor = addPrimaryConstructorForLambda(suspendLambda, arity)
|
||||
val invokeToOverride = functionNClass.functions.single {
|
||||
it.owner.valueParameters.size == arity + 1 && it.owner.name.asString() == "invoke"
|
||||
}
|
||||
val createToOverride = suspendLambda.symbol.functions.singleOrNull {
|
||||
it.owner.valueParameters.size == arity + 1 && it.owner.name.asString() == "create"
|
||||
}
|
||||
val invokeSuspend = addInvokeSuspendForLambda(function, parametersFields)
|
||||
if (function.capturesCrossinline()) {
|
||||
addInvokeSuspendForInlineForLambda(invokeSuspend)
|
||||
}
|
||||
if (createToOverride != null) {
|
||||
addInvokeCallingCreate(addCreate(constructor, createToOverride, parametersFields), invokeSuspend, invokeToOverride)
|
||||
} else {
|
||||
addInvokeCallingConstructor(constructor, invokeSuspend, invokeToOverride, parametersFields)
|
||||
}
|
||||
|
||||
context.suspendLambdaToOriginalFunctionMap[attributeOwnerId as IrFunctionReference] = function
|
||||
}
|
||||
|
||||
private fun IrClass.addInvokeSuspendForLambda(irFunction: IrFunction, fields: List<IrField>): IrSimpleFunction {
|
||||
val superMethod = context.ir.symbols.suspendLambdaClass.functions.single {
|
||||
it.owner.name.asString() == INVOKE_SUSPEND_METHOD_NAME && it.owner.valueParameters.size == 1 &&
|
||||
it.owner.valueParameters[0].type.isKotlinResult()
|
||||
}.owner
|
||||
return addFunctionOverride(superMethod).apply {
|
||||
body = irFunction.moveBodyTo(this, mapOf())?.transform(object : IrElementTransformerVoid() {
|
||||
override fun visitGetValue(expression: IrGetValue): IrExpression {
|
||||
val parameter = (expression.symbol.owner as? IrValueParameter)?.takeIf { it.parent == irFunction }
|
||||
?: return expression
|
||||
val field = fields[parameter.index + if (irFunction.extensionReceiverParameter != null) 1 else 0]
|
||||
val receiver = IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, dispatchReceiverParameter!!.symbol)
|
||||
return IrGetFieldImpl(expression.startOffset, expression.endOffset, field.symbol, field.type, receiver)
|
||||
}
|
||||
}, null)
|
||||
}
|
||||
}
|
||||
|
||||
private fun IrClass.addInvokeSuspendForInlineForLambda(invokeSuspend: IrSimpleFunction): IrSimpleFunction {
|
||||
return addFunction(
|
||||
INVOKE_SUSPEND_METHOD_NAME + FOR_INLINE_SUFFIX,
|
||||
context.irBuiltIns.anyNType,
|
||||
Modality.FINAL,
|
||||
origin = JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE
|
||||
).apply {
|
||||
copyAttributes(invokeSuspend)
|
||||
generateErrorForInlineBody()
|
||||
valueParameters = invokeSuspend.valueParameters.map { it.copyTo(this) }
|
||||
}
|
||||
}
|
||||
|
||||
// Invoke function in lambdas is responsible for
|
||||
// 1) calling `create`
|
||||
// 2) starting newly created coroutine by calling `invokeSuspend`.
|
||||
// Thus, it creates a clone of suspend lambda and starts it.
|
||||
// TODO: fix the generic signature -- type parameters of FunctionN should be substituted
|
||||
private fun IrClass.addInvokeCallingCreate(
|
||||
create: IrFunction,
|
||||
invokeSuspend: IrSimpleFunction,
|
||||
invokeToOverride: IrSimpleFunctionSymbol
|
||||
) = addFunctionOverride(invokeToOverride.owner) { function ->
|
||||
val newlyCreatedObject = irCall(create).also { createCall ->
|
||||
createCall.dispatchReceiver = irGet(function.dispatchReceiverParameter!!)
|
||||
for ((index, param) in function.valueParameters.withIndex()) {
|
||||
createCall.putValueArgument(index, irGet(param))
|
||||
}
|
||||
}
|
||||
+irReturn(callInvokeSuspend(invokeSuspend, irImplicitCast(newlyCreatedObject, defaultType)))
|
||||
}
|
||||
|
||||
// Same as above, but with `create` inlined. `create` is only defined in `SuspendLambda` in unary and binary
|
||||
// versions; for other lambdas, there's no point in generating a non-overriding `create` separately.
|
||||
private fun IrClass.addInvokeCallingConstructor(
|
||||
constructor: IrFunction,
|
||||
invokeSuspend: IrSimpleFunction,
|
||||
invokeToOverride: IrSimpleFunctionSymbol,
|
||||
fieldsForUnbound: List<IrField>
|
||||
) = addFunctionOverride(invokeToOverride.owner) { function ->
|
||||
+irReturn(callInvokeSuspend(invokeSuspend, cloneLambda(function, constructor, fieldsForUnbound)))
|
||||
}
|
||||
|
||||
private fun IrClass.addCreate(
|
||||
constructor: IrFunction,
|
||||
createToOverride: IrSimpleFunctionSymbol,
|
||||
fieldsForUnbound: List<IrField>
|
||||
) = addFunctionOverride(createToOverride.owner) { function ->
|
||||
+irReturn(cloneLambda(function, constructor, fieldsForUnbound))
|
||||
}
|
||||
|
||||
private fun IrBlockBodyBuilder.cloneLambda(
|
||||
scope: IrFunction,
|
||||
constructor: IrFunction,
|
||||
fieldsForUnbound: List<IrField>
|
||||
): IrExpression {
|
||||
val constructorCall = irCall(constructor).also {
|
||||
for (typeParameter in constructor.parentAsClass.typeParameters) {
|
||||
it.putTypeArgument(typeParameter.index, typeParameter.defaultType)
|
||||
}
|
||||
it.putValueArgument(0, irGet(scope.valueParameters.last()))
|
||||
}
|
||||
if (fieldsForUnbound.isEmpty()) {
|
||||
return constructorCall
|
||||
}
|
||||
val result = irTemporary(constructorCall, "result")
|
||||
for ((index, field) in fieldsForUnbound.withIndex()) {
|
||||
+irSetField(irGet(result), field, irGet(scope.valueParameters[index]))
|
||||
}
|
||||
return irGet(result)
|
||||
}
|
||||
|
||||
private fun IrBlockBodyBuilder.callInvokeSuspend(invokeSuspend: IrSimpleFunction, lambda: IrExpression): IrExpression =
|
||||
irCallOp(invokeSuspend.symbol, invokeSuspend.returnType, lambda, irUnit())
|
||||
|
||||
// Primary constructor accepts parameters equal to function reference arguments + continuation and sets the fields.
|
||||
private fun IrClass.addPrimaryConstructorForLambda(superClass: IrClass, arity: Int): IrConstructor =
|
||||
addConstructor {
|
||||
origin = JvmLoweredDeclarationOrigin.SUSPEND_LAMBDA
|
||||
isPrimary = true
|
||||
returnType = defaultType
|
||||
visibility = DescriptorVisibilities.LOCAL
|
||||
}.also { constructor ->
|
||||
val completionParameterSymbol = constructor.addCompletionValueParameter()
|
||||
val superClassConstructor = superClass.constructors.single {
|
||||
it.valueParameters.size == 2 && it.valueParameters[0].type.isInt() && it.valueParameters[1].type.isNullableContinuation()
|
||||
}
|
||||
constructor.body = context.createIrBuilder(constructor.symbol).irBlockBody {
|
||||
+irDelegatingConstructorCall(superClassConstructor).also {
|
||||
it.putValueArgument(0, irInt(arity + 1))
|
||||
it.putValueArgument(1, irGet(completionParameterSymbol))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user