Fix false positive ENCLOSING_SUSPEND_FUNCTION_FOR_SUSPEND_FUNCTION_CALL

Previously added additional processing at findEnclosingSuspendFunction
seems unnecessary anymore

^KT-43258 Fixed

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

@@ -5601,6 +5601,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                 runTest("compiler/testData/diagnostics/tests/coroutines/suspendInvokeInsideWhen.kt");
             }
 
+            @Test
+            @TestMetadata("suspendInvokeWithReceiver.kt")
+            public void testSuspendInvokeWithReceiver() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/coroutines/suspendInvokeWithReceiver.kt");
+            }
+
             @Nested
             @TestMetadata("compiler/testData/diagnostics/tests/coroutines/callableReference")
             @TestDataPath("$PROJECT_ROOT")

compiler/frontend.java/src/org/jetbrains/kotlin/resolve/jvm/checkers/SuspensionPointInsideMutexLockChecker.kt

@@ -26,7 +26,7 @@ class SuspensionPointInsideMutexLockChecker : CallChecker {
         if (descriptor !is FunctionDescriptor || !descriptor.isSuspend) return
 
         val enclosingSuspendFunctionSource =
-            findEnclosingSuspendFunction(context, resolvedCall.call.callElement)?.source?.getPsi() ?: return
+            findEnclosingSuspendFunction(context)?.source?.getPsi() ?: return
 
         // Search for `synchronized` call
         var parent = reportOn
@@ -78,4 +78,4 @@ class SuspensionPointInsideMutexLockChecker : CallChecker {
     private fun reportProblem(context: CallCheckerContext, reportOn: PsiElement, resolvedCall: ResolvedCall<*>) {
         context.trace.report(ErrorsJvm.SUSPENSION_POINT_INSIDE_CRITICAL_SECTION.on(reportOn, resolvedCall.resultingDescriptor))
     }
-}
+}

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/checkers/coroutineCallChecker.kt

@@ -14,13 +14,11 @@ import org.jetbrains.kotlin.diagnostics.DiagnosticSink
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtCodeFragment
-import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtThisExpression
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.calls.callUtil.isCallableReference
-import org.jetbrains.kotlin.resolve.calls.context.CallResolutionContext
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.resolve.inline.InlineUtil
@@ -52,28 +50,10 @@ fun PropertyDescriptor.isBuiltInCoroutineContext(languageVersionSettings: Langua
 
 private val ALLOWED_SCOPE_KINDS = setOf(LexicalScopeKind.FUNCTION_INNER_SCOPE, LexicalScopeKind.FUNCTION_HEADER_FOR_DESTRUCTURING)
 
-fun findEnclosingSuspendFunction(context: CallCheckerContext, checkingCall: KtElement): FunctionDescriptor? {
-    /*
-     * If checking call isn't equal to call in resolution context, we should look at lexical scope from trace.
-     * It means there is a parent function analysis of which isn't completed yet
-     * and their lexical scope in the resolution context isn't recorded yet (but there is lexical scope with not completed descriptor in trace).
-     * Example (suggest that we're analyzing the last expression of lambda now):
-     *      fun main() {
-     *          runBlocking {
-     *              retry { 1 } // `fun main` lexical scope in the resolution context, `runBlocking { ... }` one in the recorded in trace lexical scope
-     *          }
-     *      }
-     */
-    val scope = if (context.resolutionContext !is CallResolutionContext<*> || context.resolutionContext.call.callElement == checkingCall) {
-        context.scope
-    } else {
-        context.trace.get(BindingContext.LEXICAL_SCOPE, checkingCall) ?: context.scope
-    }
-
-    return scope.parentsWithSelf.firstOrNull {
+fun findEnclosingSuspendFunction(context: CallCheckerContext): FunctionDescriptor? =
+    context.scope.parentsWithSelf.firstOrNull {
         it is LexicalScope && it.kind in ALLOWED_SCOPE_KINDS && it.ownerDescriptor.safeAs<FunctionDescriptor>()?.isSuspend == true
     }?.cast<LexicalScope>()?.ownerDescriptor?.cast()
-}
 
 object CoroutineSuspendCallChecker : CallChecker {
     override fun check(resolvedCall: ResolvedCall<*>, reportOn: PsiElement, context: CallCheckerContext) {
@@ -89,7 +69,7 @@ object CoroutineSuspendCallChecker : CallChecker {
         }
 
         val callElement = resolvedCall.call.callElement as KtExpression
-        val enclosingSuspendFunction = findEnclosingSuspendFunction(context, callElement)
+        val enclosingSuspendFunction = findEnclosingSuspendFunction(context)
 
         when {
             enclosingSuspendFunction != null -> {

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/ResolvedAtomCompleter.kt

@@ -117,7 +117,8 @@ class ResolvedAtomCompleter(
 
         clearPartiallyResolvedCall(resolvedCallAtom)
 
-        if (resolvedCallAtom.atom.psiKotlinCall is PSIKotlinCallForVariable) return null
+        val atom = resolvedCallAtom.atom
+        if (atom.psiKotlinCall is PSIKotlinCallForVariable) return null
 
         val allDiagnostics = diagnostics + diagnosticsFromPartiallyResolvedCall
 
@@ -135,8 +136,14 @@ class ResolvedAtomCompleter(
             return resolvedCall
         }
 
+        val psiCallForResolutionContext = when (atom) {
+            // PARTIAL_CALL_RESOLUTION_CONTEXT has been written for the baseCall
+            is PSIKotlinCallForInvoke -> atom.baseCall.psiCall
+            else -> atom.psiKotlinCall.psiCall
+        }
+
         val resolutionContextForPartialCall =
-            topLevelCallContext.trace[BindingContext.PARTIAL_CALL_RESOLUTION_CONTEXT, resolvedCallAtom.atom.psiKotlinCall.psiCall]
+            topLevelCallContext.trace[BindingContext.PARTIAL_CALL_RESOLUTION_CONTEXT, psiCallForResolutionContext]
 
         val callCheckerContext = if (resolutionContextForPartialCall != null)
             CallCheckerContext(

compiler/testData/diagnostics/tests/coroutines/suspendInvokeWithReceiver.kt

@@ -0,0 +1,13 @@
+// FIR_IDENTICAL
+// SKIP_TXT
+fun <T> r(x: suspend () -> T): T = null!!
+
+fun nonReproducer1(): String = r {
+    MyObject.someProperty()
+}
+
+object MyObject {
+    val someProperty: String = "TODO()"
+}
+
+suspend operator fun <R> String.invoke(): R = null!!

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

@@ -5607,6 +5607,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                 runTest("compiler/testData/diagnostics/tests/coroutines/suspendInvokeInsideWhen.kt");
             }
 
+            @Test
+            @TestMetadata("suspendInvokeWithReceiver.kt")
+            public void testSuspendInvokeWithReceiver() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/coroutines/suspendInvokeWithReceiver.kt");
+            }
+
             @Nested
             @TestMetadata("compiler/testData/diagnostics/tests/coroutines/callableReference")
             @TestDataPath("$PROJECT_ROOT")