diff --git a/compiler/fir/fir2ir/tests/org/jetbrains/kotlin/codegen/ir/FirBlackBoxCodegenTestGenerated.java b/compiler/fir/fir2ir/tests/org/jetbrains/kotlin/codegen/ir/FirBlackBoxCodegenTestGenerated.java index beb96dd2ccc..2c345588754 100644 --- a/compiler/fir/fir2ir/tests/org/jetbrains/kotlin/codegen/ir/FirBlackBoxCodegenTestGenerated.java +++ b/compiler/fir/fir2ir/tests/org/jetbrains/kotlin/codegen/ir/FirBlackBoxCodegenTestGenerated.java @@ -6660,6 +6660,11 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/recursiveSuspend.kt", "kotlin.coroutines"); } + @TestMetadata("restrictedSuspendLambda.kt") + public void testRestrictedSuspendLambda() throws Exception { + runTest("compiler/testData/codegen/box/coroutines/restrictedSuspendLambda.kt"); + } + @TestMetadata("returnByLabel.kt") public void testReturnByLabel_1_3() throws Exception { runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/returnByLabel.kt", "kotlin.coroutines"); diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmSymbols.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmSymbols.kt index 01f45dfe575..707f869dac1 100644 --- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmSymbols.kt +++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmSymbols.kt @@ -270,26 +270,35 @@ class JvmSymbols( val suspendLambdaClass: IrClassSymbol = createClass(FqName("kotlin.coroutines.jvm.internal.SuspendLambda"), classModality = Modality.ABSTRACT) { klass -> - klass.superTypes += suspendFunctionInterface.defaultType - klass.addConstructor().apply { - addValueParameter("arity", irBuiltIns.intType) - addValueParameter( - SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME, - continuationClass.typeWith(irBuiltIns.anyNType).makeNullable() - ) - } - klass.addFunction(INVOKE_SUSPEND_METHOD_NAME, irBuiltIns.anyNType, Modality.ABSTRACT, Visibilities.PROTECTED).apply { - addValueParameter(SUSPEND_CALL_RESULT_NAME, resultClassStub.typeWith(irBuiltIns.anyNType)) - } - klass.addFunction(SUSPEND_FUNCTION_CREATE_METHOD_NAME, continuationClass.typeWith(irBuiltIns.unitType), Modality.OPEN).apply { - addValueParameter(SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME, continuationClass.typeWith(irBuiltIns.nothingType)) - } - klass.addFunction(SUSPEND_FUNCTION_CREATE_METHOD_NAME, continuationClass.typeWith(irBuiltIns.unitType), Modality.OPEN).apply { - addValueParameter("value", irBuiltIns.anyNType) - addValueParameter(SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME, continuationClass.typeWith(irBuiltIns.nothingType)) - } + addSuspendLambdaInterfaceFunctions(klass) } + val restrictedSuspendLambdaClass: IrClassSymbol = + createClass(FqName("kotlin.coroutines.jvm.internal.RestrictedSuspendLambda"), classModality = Modality.ABSTRACT) { klass -> + addSuspendLambdaInterfaceFunctions(klass) + } + + private fun addSuspendLambdaInterfaceFunctions(klass: IrClass) { + klass.superTypes += suspendFunctionInterface.defaultType + klass.addConstructor().apply { + addValueParameter("arity", irBuiltIns.intType) + addValueParameter( + SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME, + continuationClass.typeWith(irBuiltIns.anyNType).makeNullable() + ) + } + klass.addFunction(INVOKE_SUSPEND_METHOD_NAME, irBuiltIns.anyNType, Modality.ABSTRACT, Visibilities.PROTECTED).apply { + addValueParameter(SUSPEND_CALL_RESULT_NAME, resultClassStub.typeWith(irBuiltIns.anyNType)) + } + klass.addFunction(SUSPEND_FUNCTION_CREATE_METHOD_NAME, continuationClass.typeWith(irBuiltIns.unitType), Modality.OPEN).apply { + addValueParameter(SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME, continuationClass.typeWith(irBuiltIns.nothingType)) + } + klass.addFunction(SUSPEND_FUNCTION_CREATE_METHOD_NAME, continuationClass.typeWith(irBuiltIns.unitType), Modality.OPEN).apply { + addValueParameter("value", irBuiltIns.anyNType) + addValueParameter(SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME, continuationClass.typeWith(irBuiltIns.nothingType)) + } + } + private fun generateCallableReferenceMethods(klass: IrClass) { klass.addFunction("getSignature", irBuiltIns.stringType, Modality.OPEN) klass.addFunction("getName", irBuiltIns.stringType, Modality.OPEN) 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 8e2642ca860..b98539ba43f 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 @@ -102,7 +102,10 @@ private class AddContinuationLowering(private val context: JvmBackendContext) : override fun visitFunctionReference(expression: IrFunctionReference) { expression.acceptChildrenVoid(this) if (expression.isSuspend && expression.shouldBeTreatedAsSuspendLambda() && expression !in inlineReferences) { - suspendLambdas[expression] = SuspendLambdaInfo(expression) + val isRestricted = expression.symbol.owner.extensionReceiverParameter?.type?.classOrNull?.owner?.annotations?.any { + it.type.classOrNull?.signature == IdSignature.PublicSignature("kotlin.coroutines", "RestrictsSuspension", null, 0) + } == true + suspendLambdas[expression] = SuspendLambdaInfo(expression, isRestricted) } } @@ -151,7 +154,9 @@ private class AddContinuationLowering(private val context: JvmBackendContext) : parent: IrDeclarationParent, insideInlineFunction: Boolean ): IrClass { - val suspendLambda = context.ir.symbols.suspendLambdaClass.owner + val suspendLambda = + if (info.isRestricted) context.ir.symbols.restrictedSuspendLambdaClass.owner + else context.ir.symbols.suspendLambdaClass.owner return suspendLambda.createContinuationClassFor( parent, JvmLoweredDeclarationOrigin.SUSPEND_LAMBDA, @@ -200,7 +205,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) : } val fieldsForBound = parametersFields.filter { it.isFinal } val fieldsForUnbound = listOfNotNull(receiverField) + parametersFields.filter { !it.isFinal } - val constructor = addPrimaryConstructorForLambda(info.arity, info.reference, fieldsForBound, insideInlineFunction) + val constructor = addPrimaryConstructorForLambda(info, fieldsForBound, insideInlineFunction) val invokeToOverride = functionNClass.functions.single { it.owner.valueParameters.size == info.arity + 1 && it.owner.name.asString() == "invoke" } @@ -369,8 +374,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) : // Primary constructor accepts parameters equal to function reference arguments + continuation and sets the fields. private fun IrClass.addPrimaryConstructorForLambda( - arity: Int, - reference: IrFunctionReference, + info: SuspendLambdaInfo, fields: List, insideInlineFunction: Boolean ): IrConstructor = @@ -379,17 +383,20 @@ private class AddContinuationLowering(private val context: JvmBackendContext) : returnType = defaultType visibility = if (insideInlineFunction) Visibilities.PUBLIC else JavaVisibilities.PACKAGE_VISIBILITY }.also { constructor -> - for ((param, arg) in reference.getArguments()) { + for ((param, arg) in info.reference.getArguments()) { constructor.addValueParameter(name = param.name.asString(), type = arg.type) } val completionParameterSymbol = constructor.addCompletionValueParameter() - val superClassConstructor = context.ir.symbols.suspendLambdaClass.owner.constructors.single { + val superClass = + if (info.isRestricted) context.ir.symbols.restrictedSuspendLambdaClass + else context.ir.symbols.suspendLambdaClass + val superClassConstructor = superClass.owner.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(0, irInt(info.arity + 1)) it.putValueArgument(1, irGet(completionParameterSymbol)) } @@ -636,7 +643,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) : }, null) } - private class SuspendLambdaInfo(val reference: IrFunctionReference) { + private class SuspendLambdaInfo(val reference: IrFunctionReference, val isRestricted: Boolean) { val function = reference.symbol.owner val arity = (reference.type as IrSimpleType).arguments.size - 1 val capturesCrossinline = function.valueParameters.any { reference.getValueArgument(it.index).isReadOfCrossinline() } diff --git a/compiler/testData/codegen/box/coroutines/restrictedSuspendLambda.kt b/compiler/testData/codegen/box/coroutines/restrictedSuspendLambda.kt new file mode 100644 index 00000000000..41c6d210314 --- /dev/null +++ b/compiler/testData/codegen/box/coroutines/restrictedSuspendLambda.kt @@ -0,0 +1,22 @@ +// WITH_RUNTIME +// TARGET_BACKEND: JVM + +import kotlin.coroutines.* + +@RestrictsSuspension +interface Marker { + fun restricted() +} + +var lambda: Any? = null + +fun acceptsRestricted(c: suspend Marker.() -> Unit) { + lambda = c +} + +fun box(): String { + acceptsRestricted {} + @Suppress("INVISIBLE_REFERENCE") + return if (lambda is kotlin.coroutines.jvm.internal.RestrictedSuspendLambda) "OK" + else "FAIL" +} \ No newline at end of file diff --git a/compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java b/compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java index 587ded0d4c5..1fa96890f47 100644 --- a/compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java +++ b/compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java @@ -7005,6 +7005,11 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest { runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/recursiveSuspend.kt", "kotlin.coroutines"); } + @TestMetadata("restrictedSuspendLambda.kt") + public void testRestrictedSuspendLambda() throws Exception { + runTest("compiler/testData/codegen/box/coroutines/restrictedSuspendLambda.kt"); + } + @TestMetadata("returnByLabel.kt") public void testReturnByLabel_1_2() throws Exception { runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/returnByLabel.kt", "kotlin.coroutines.experimental"); diff --git a/compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java b/compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java index b03b515865d..7c4b6261449 100644 --- a/compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java +++ b/compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java @@ -7005,6 +7005,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/recursiveSuspend.kt", "kotlin.coroutines"); } + @TestMetadata("restrictedSuspendLambda.kt") + public void testRestrictedSuspendLambda() throws Exception { + runTest("compiler/testData/codegen/box/coroutines/restrictedSuspendLambda.kt"); + } + @TestMetadata("returnByLabel.kt") public void testReturnByLabel_1_2() throws Exception { runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/returnByLabel.kt", "kotlin.coroutines.experimental"); diff --git a/compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java b/compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java index 804d0c9e38c..64807b760fd 100644 --- a/compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java +++ b/compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java @@ -6660,6 +6660,11 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/recursiveSuspend.kt", "kotlin.coroutines"); } + @TestMetadata("restrictedSuspendLambda.kt") + public void testRestrictedSuspendLambda() throws Exception { + runTest("compiler/testData/codegen/box/coroutines/restrictedSuspendLambda.kt"); + } + @TestMetadata("returnByLabel.kt") public void testReturnByLabel_1_3() throws Exception { runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/returnByLabel.kt", "kotlin.coroutines");