diff --git a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosisCompilerTestFE10TestdataTestGenerated.java b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosisCompilerTestFE10TestdataTestGenerated.java index 03af46eb488..254fa10ece8 100644 --- a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosisCompilerTestFE10TestdataTestGenerated.java +++ b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosisCompilerTestFE10TestdataTestGenerated.java @@ -35410,6 +35410,12 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/extensions.kt"); } + @Test + @TestMetadata("invokeExtension.kt") + public void testInvokeExtension() throws Exception { + runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/invokeExtension.kt"); + } + @Test @TestMetadata("memberExtension.kt") public void testMemberExtension() throws Exception { diff --git a/compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsTestGenerated.java b/compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsTestGenerated.java index 11497c892dd..64f6d47efdf 100644 --- a/compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsTestGenerated.java +++ b/compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsTestGenerated.java @@ -35410,6 +35410,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/extensions.kt"); } + @Test + @TestMetadata("invokeExtension.kt") + public void testInvokeExtension() throws Exception { + runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/invokeExtension.kt"); + } + @Test @TestMetadata("memberExtension.kt") public void testMemberExtension() throws Exception { diff --git a/compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsWithLightTreeTestGenerated.java b/compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsWithLightTreeTestGenerated.java index 1b26fcbef03..50cf1150cfe 100644 --- a/compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsWithLightTreeTestGenerated.java +++ b/compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsWithLightTreeTestGenerated.java @@ -35410,6 +35410,12 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/extensions.kt"); } + @Test + @TestMetadata("invokeExtension.kt") + public void testInvokeExtension() throws Exception { + runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/invokeExtension.kt"); + } + @Test @TestMetadata("memberExtension.kt") public void testMemberExtension() throws Exception { diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirSuspendCallChecker.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirSuspendCallChecker.kt index b2418a74a8f..c9e64f47866 100644 --- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirSuspendCallChecker.kt +++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirSuspendCallChecker.kt @@ -23,16 +23,12 @@ import org.jetbrains.kotlin.fir.expressions.impl.FirNoReceiverExpression import org.jetbrains.kotlin.fir.references.FirResolvedCallableReference import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference import org.jetbrains.kotlin.fir.resolve.fullyExpandedType -import org.jetbrains.kotlin.fir.types.isSuspendFunctionType import org.jetbrains.kotlin.fir.resolve.toFirRegularClassSymbol import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol -import org.jetbrains.kotlin.fir.types.ConeClassLikeType -import org.jetbrains.kotlin.fir.types.ConeKotlinType -import org.jetbrains.kotlin.fir.types.ConeTypeParameterType -import org.jetbrains.kotlin.fir.types.coneType +import org.jetbrains.kotlin.fir.types.* import org.jetbrains.kotlin.name.CallableId import org.jetbrains.kotlin.name.StandardClassIds import org.jetbrains.kotlin.utils.addToStdlib.lastIsInstanceOrNull @@ -144,11 +140,14 @@ object FirSuspendCallChecker : FirQualifiedAccessExpressionChecker() { context: CheckerContext ): Boolean { val session = context.session + val enclosingSuspendFunctionDispatchReceiverOwnerSymbol = (enclosingSuspendFunction.dispatchReceiverType as? ConeClassLikeType)?.lookupTag?.toFirRegularClassSymbol(session) val enclosingSuspendFunctionExtensionReceiverOwnerSymbol = enclosingSuspendFunction.takeIf { it.receiverTypeRef != null }?.symbol - val dispatchReceiverExpression = expression.dispatchReceiver.takeIf { it !is FirNoReceiverExpression } - val extensionReceiverExpression = expression.extensionReceiver.takeIf { it !is FirNoReceiverExpression } + + val (dispatchReceiverExpression, extensionReceiverExpression, extensionReceiverParameterType) = + expression.computeReceiversInfo(session, calledDeclarationSymbol) + for (receiverExpression in listOfNotNull(dispatchReceiverExpression, extensionReceiverExpression)) { if (!receiverExpression.typeRef.coneType.isRestrictSuspensionReceiver(session)) continue if (sameInstanceOfReceiver(receiverExpression, enclosingSuspendFunctionDispatchReceiverOwnerSymbol)) continue @@ -156,14 +155,17 @@ object FirSuspendCallChecker : FirQualifiedAccessExpressionChecker() { return false } + if (enclosingSuspendFunctionExtensionReceiverOwnerSymbol?.resolvedReceiverTypeRef?.coneType?.isRestrictSuspensionReceiver(session) != true) { return true } + if (sameInstanceOfReceiver(dispatchReceiverExpression, enclosingSuspendFunctionExtensionReceiverOwnerSymbol)) { return true } + if (sameInstanceOfReceiver(extensionReceiverExpression, enclosingSuspendFunctionExtensionReceiverOwnerSymbol)) { - if (calledDeclarationSymbol.resolvedReceiverTypeRef?.coneType?.isRestrictSuspensionReceiver(session) == true) { + if (extensionReceiverParameterType?.isRestrictSuspensionReceiver(session) == true) { return true } } @@ -196,4 +198,32 @@ object FirSuspendCallChecker : FirQualifiedAccessExpressionChecker() { } return false } + + // Triple + private fun FirQualifiedAccessExpression.computeReceiversInfo( + session: FirSession, + calledDeclarationSymbol: FirCallableSymbol<*> + ): Triple { + if (this is FirImplicitInvokeCall && + dispatchReceiver.typeRef.coneTypeSafe()?.isSuspendFunctionType(session) == true + ) { + val variableForInvoke = dispatchReceiver + val variableForInvokeType = variableForInvoke.typeRef.coneType + if (!variableForInvokeType.isExtensionFunctionType) return Triple(null, null, null) + + // `a.foo()` is resolved to invokeExtension, so it's been desugared to `foo.invoke(a)` + // And we use the first argument (`a`) as an extension receiver + return Triple( + null, + argumentList.arguments.getOrNull(0), + variableForInvokeType.typeArguments.getOrNull(0) as? ConeKotlinType + ) + } + + return Triple( + dispatchReceiver.takeIf { it !is FirNoReceiverExpression }, + extensionReceiver.takeIf { it !is FirNoReceiverExpression }, + calledDeclarationSymbol.resolvedReceiverTypeRef?.coneType, + ) + } } diff --git a/compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/invokeExtension.kt b/compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/invokeExtension.kt new file mode 100644 index 00000000000..81c4800d084 --- /dev/null +++ b/compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/invokeExtension.kt @@ -0,0 +1,18 @@ +// FIR_IDENTICAL +// SKIP_TXT +import kotlin.coroutines.RestrictsSuspension + +@RestrictsSuspension +interface Scope + +suspend fun Scope.foo( + a1: suspend Scope.() -> Unit, + a2: suspend () -> Unit, + a3: suspend String.() -> Unit, +) { + a1() + this.a1() + a2() + + "".a3() +} diff --git a/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java b/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java index de9b104675f..971ebf4c269 100644 --- a/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java +++ b/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java @@ -35506,6 +35506,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest { runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/extensions.kt"); } + @Test + @TestMetadata("invokeExtension.kt") + public void testInvokeExtension() throws Exception { + runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/restrictSuspension/invokeExtension.kt"); + } + @Test @TestMetadata("memberExtension.kt") public void testMemberExtension() throws Exception {