Support restricted suspend lambdas in JVM_IR

#KT-40135 Fixed

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");

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)

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<IrField>,
         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() }

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"
+}

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");

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");

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");