[FIR] Fix false positive "suspension point is inside a critical section"

^KT-55072 Fixed

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosticCompilerTestFE10TestdataTestGenerated.java

@@ -38578,6 +38578,12 @@ public class DiagnosticCompilerTestFE10TestdataTestGenerated extends AbstractDia
                 runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/kt41430.kt");
             }
 
+            @Test
+            @TestMetadata("kt55072.kt")
+            public void testKt55072() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/kt55072.kt");
+            }
+
             @Test
             @TestMetadata("lambdaExpectedType.kt")
             public void testLambdaExpectedType() throws Exception {

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated.java

@@ -38578,6 +38578,12 @@ public class LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated
                 runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/kt41430.kt");
             }
 
+            @Test
+            @TestMetadata("kt55072.kt")
+            public void testKt55072() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/kt55072.kt");
+            }
+
             @Test
             @TestMetadata("lambdaExpectedType.kt")
             public void testLambdaExpectedType() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirLightTreeOldFrontendDiagnosticsTestGenerated.java

@@ -38578,6 +38578,12 @@ public class FirLightTreeOldFrontendDiagnosticsTestGenerated extends AbstractFir
                 runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/kt41430.kt");
             }
 
+            @Test
+            @TestMetadata("kt55072.kt")
+            public void testKt55072() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/kt55072.kt");
+            }
+
             @Test
             @TestMetadata("lambdaExpectedType.kt")
             public void testLambdaExpectedType() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirPsiOldFrontendDiagnosticsTestGenerated.java

@@ -38692,6 +38692,12 @@ public class FirPsiOldFrontendDiagnosticsTestGenerated extends AbstractFirPsiDia
                 runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/kt41430.kt");
             }
 
+            @Test
+            @TestMetadata("kt55072.kt")
+            public void testKt55072() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/kt55072.kt");
+            }
+
             @Test
             @TestMetadata("lambdaExpectedType.kt")
             public void testLambdaExpectedType() throws Exception {

compiler/fir/checkers/checkers.jvm/src/org/jetbrains/kotlin/fir/analysis/jvm/checkers/expression/FirJvmSuspensionPointInsideMutexLockChecker.kt

@@ -5,46 +5,72 @@
 
 package org.jetbrains.kotlin.fir.analysis.jvm.checkers.expression
 
-import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
-import org.jetbrains.kotlin.fir.analysis.checkers.context.findClosest
-import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirFunctionCallChecker
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
-import org.jetbrains.kotlin.fir.analysis.diagnostics.jvm.FirJvmErrors
 import org.jetbrains.kotlin.diagnostics.reportOn
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirFunctionCallChecker
+import org.jetbrains.kotlin.fir.analysis.diagnostics.jvm.FirJvmErrors
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
+import org.jetbrains.kotlin.fir.declarations.FirFunction
 import org.jetbrains.kotlin.fir.declarations.utils.isSuspend
-import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.resolvedArgumentMapping
 import org.jetbrains.kotlin.fir.references.toResolvedCallableSymbol
+import org.jetbrains.kotlin.fir.resolve.transformers.unwrapAnonymousFunctionExpression
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.isSuspendOrKSuspendFunctionType
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.utils.addToStdlib.runIf
 
 object FirJvmSuspensionPointInsideMutexLockChecker : FirFunctionCallChecker() {
     private val synchronizedCallableId = CallableId(FqName("kotlin"), Name.identifier("synchronized"))
     private val withLockCallableId = CallableId(FqName("kotlin.concurrent"), Name.identifier("withLock"))
+    private val synchronizedBlockParamName = Name.identifier("block")
 
     override fun check(expression: FirFunctionCall, context: CheckerContext, reporter: DiagnosticReporter) {
         val symbol = expression.calleeReference.toResolvedCallableSymbol() ?: return
         if (!symbol.isSuspend) return
-        val closestAnonymousFunction = context.findClosest<FirAnonymousFunction>() ?: return
+        var anonymousFunctionArg: FirAnonymousFunction? = null
+        var isMutexLockFound = false
+        var isSuspendFunctionFound = false
 
-        for (call in context.callsOrAssignments.asReversed()) {
+        for (element in context.containingElements.asReversed()) {
-            if (call is FirFunctionCall) {
+            if (element is FirFunctionCall) {
-                val callableSymbol = call.calleeReference.toResolvedCallableSymbol() ?: continue
+                val callableSymbol = element.calleeReference.toResolvedCallableSymbol() ?: continue
-                if (callableSymbol.callableId == synchronizedCallableId) {
+                val enclosingAnonymousFuncParam = element.resolvedArgumentMapping?.firstNotNullOfOrNull { entry ->
-                    val unwrappedFirstArgument = call.arguments.elementAtOrNull(1)?.unwrapArgument() ?: return
+                    entry.key.unwrapAnonymousFunctionExpression()?.let {
-                    val firstArgumentAnonymousFunction =
+                        runIf(it == anonymousFunctionArg) { entry.value }
-                        (unwrappedFirstArgument as? FirAnonymousFunctionExpression)?.anonymousFunction ?: return
-
-                    if (closestAnonymousFunction == firstArgumentAnonymousFunction) {
-                        reporter.reportOn(expression.source, FirJvmErrors.SUSPENSION_POINT_INSIDE_CRITICAL_SECTION, symbol, context)
                     }
-                    return
+                }
-                } else if (callableSymbol.callableId == withLockCallableId) {
+
-                    reporter.reportOn(expression.source, FirJvmErrors.SUSPENSION_POINT_INSIDE_CRITICAL_SECTION, symbol, context)
+                if ((enclosingAnonymousFuncParam?.returnTypeRef as? FirResolvedTypeRef)?.type?.isSuspendOrKSuspendFunctionType(context.session) == true) {
-                    return
+                    isSuspendFunctionFound = true
+                    break
+                }
+
+                if (callableSymbol.callableId == synchronizedCallableId &&
+                    enclosingAnonymousFuncParam?.name == synchronizedBlockParamName ||
+                    callableSymbol.callableId == withLockCallableId
+                ) {
+                    isMutexLockFound = true
+                }
+            } else if (element is FirFunction) {
+                if (element.isSuspend) {
+                    isSuspendFunctionFound = true
+                    break
+                }
+                if (element is FirAnonymousFunction) {
+                    anonymousFunctionArg = element // For anonymous function argument `isSuspend` can be detected from the respective parameter
                 }
             }
         }
+
+        // There is no need to report SUSPENSION_POINT_INSIDE_CRITICAL_SECTION if enclosing suspend function is not found
+        // Because ILLEGAL_SUSPEND_FUNCTION_CALL is reported in this case
+        if (isMutexLockFound && isSuspendFunctionFound) {
+            reporter.reportOn(expression.source, FirJvmErrors.SUSPENSION_POINT_INSIDE_CRITICAL_SECTION, symbol, context)
+        }
     }
 }

compiler/testData/diagnostics/testsWithStdLib/coroutines/kt55072.kt

@@ -0,0 +1,29 @@
+// FIR_IDENTICAL
+// ISSUE: KT-55072
+
+import java.util.concurrent.locks.ReentrantLock
+import kotlin.concurrent.withLock
+
+class Player
+{
+    fun play() {
+        ReentrantLock().withLock {
+            launch {
+                pumpEvents()
+            }
+            suspend fun launch2() {
+                pumpEvents()
+            }
+            suspend {
+                pumpEvents()
+            }
+            run {
+                <!ILLEGAL_SUSPEND_FUNCTION_CALL!>pumpEvents<!>()
+            }
+        }
+    }
+
+    private suspend fun pumpEvents() {}
+
+    private fun launch(block: suspend Any.() -> Unit) {}
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java

@@ -39554,6 +39554,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                 runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/kt41430.kt");
             }
 
+            @Test
+            @TestMetadata("kt55072.kt")
+            public void testKt55072() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/coroutines/kt55072.kt");
+            }
+
             @Test
             @TestMetadata("lambdaExpectedType.kt")
             public void testLambdaExpectedType() throws Exception {