NI: add fallback strategy to get lexical scope to checking coroutine call legality

^KT-40247 Fixed
^KT-40337 Fixed

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

@@ -4691,6 +4691,16 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirOldFronte
             KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/coroutines"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
         }
 
+        @TestMetadata("suspendInvokeInsideTry.kt")
+        public void testSuspendInvokeInsideTry() throws Exception {
+            runTest("compiler/testData/diagnostics/tests/coroutines/suspendInvokeInsideTry.kt");
+        }
+
+        @TestMetadata("suspendInvokeInsideWhen.kt")
+        public void testSuspendInvokeInsideWhen() throws Exception {
+            runTest("compiler/testData/diagnostics/tests/coroutines/suspendInvokeInsideWhen.kt");
+        }
+
         @TestMetadata("compiler/testData/diagnostics/tests/coroutines/callableReference")
         @TestDataPath("$PROJECT_ROOT")
         @RunWith(JUnit3RunnerWithInners.class)

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

@@ -64,10 +64,10 @@ fun findEnclosingSuspendFunction(context: CallCheckerContext, checkingCall: KtEl
     val scope = if (context.resolutionContext !is CallResolutionContext<*> || context.resolutionContext.call.callElement == checkingCall) {
         context.scope
     } else {
-        context.trace.get(BindingContext.LEXICAL_SCOPE, checkingCall)
+        context.trace.get(BindingContext.LEXICAL_SCOPE, checkingCall) ?: context.scope
     }
 
-    return scope?.parentsWithSelf?.firstOrNull {
+    return scope.parentsWithSelf.firstOrNull {
         it is LexicalScope && it.kind in ALLOWED_SCOPE_KINDS && it.ownerDescriptor.safeAs<FunctionDescriptor>()?.isSuspend == true
     }?.cast<LexicalScope>()?.ownerDescriptor?.cast()
 }

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

@@ -0,0 +1,13 @@
+// FIR_IDENTICAL
+// !DIAGNOSTICS: -UNUSED_PARAMETER -UNCHECKED_CAST
+
+class Foo {
+    suspend operator fun <T> invoke(body: () -> T) = null as T
+}
+
+suspend fun main() {
+    val retry = Foo()
+    try {
+        retry { 1 } // Before the fix: "Suspend function 'invoke' should be called only from a coroutine or another suspend function"
+    } catch (e: Exception) { }
+}

compiler/testData/diagnostics/tests/coroutines/suspendInvokeInsideTry.txt

@@ -0,0 +1,11 @@
+package
+
+public suspend fun main(): kotlin.Unit
+
+public final class Foo {
+    public constructor Foo()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public final suspend operator fun </*0*/ T> invoke(/*0*/ body: () -> T): T
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

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

@@ -0,0 +1,16 @@
+// FIR_IDENTICAL
+// !DIAGNOSTICS: -UNUSED_PARAMETER -UNCHECKED_CAST
+
+internal class Demo() {
+    suspend operator fun <T> invoke(name: String, block: suspend () -> T): T {
+        TODO()
+    }
+}
+
+suspend fun demo(callback: suspend () -> Unit) = when {
+    true -> {
+        val demo = Demo()
+        demo("test") { callback() } // Before the fix: "Suspend function 'invoke' should be called only from a coroutine or another suspend function"
+    }
+    else -> TODO()
+}

compiler/testData/diagnostics/tests/coroutines/suspendInvokeInsideWhen.txt

@@ -0,0 +1,11 @@
+package
+
+public suspend fun demo(/*0*/ callback: suspend () -> kotlin.Unit): kotlin.Unit
+
+internal final class Demo {
+    public constructor Demo()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public final suspend operator fun </*0*/ T> invoke(/*0*/ name: kotlin.String, /*1*/ block: suspend () -> T): T
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

compiler/tests/org/jetbrains/kotlin/checkers/DiagnosticsTestGenerated.java

@@ -4698,6 +4698,16 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTestWithFirVali
                 KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/coroutines"), Pattern.compile("^(.*)\\.kts?$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
             }
 
+            @TestMetadata("suspendInvokeInsideTry.kt")
+            public void testSuspendInvokeInsideTry() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/coroutines/suspendInvokeInsideTry.kt");
+            }
+
+            @TestMetadata("suspendInvokeInsideWhen.kt")
+            public void testSuspendInvokeInsideWhen() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/coroutines/suspendInvokeInsideWhen.kt");
+            }
+
             @TestMetadata("compiler/testData/diagnostics/tests/coroutines/callableReference")
             @TestDataPath("$PROJECT_ROOT")
             @RunWith(JUnit3RunnerWithInners.class)

compiler/tests/org/jetbrains/kotlin/checkers/javac/DiagnosticsUsingJavacTestGenerated.java

@@ -4693,6 +4693,16 @@ public class DiagnosticsUsingJavacTestGenerated extends AbstractDiagnosticsUsing
                 KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/coroutines"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
             }
 
+            @TestMetadata("suspendInvokeInsideTry.kt")
+            public void testSuspendInvokeInsideTry() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/coroutines/suspendInvokeInsideTry.kt");
+            }
+
+            @TestMetadata("suspendInvokeInsideWhen.kt")
+            public void testSuspendInvokeInsideWhen() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/coroutines/suspendInvokeInsideWhen.kt");
+            }
+
             @TestMetadata("compiler/testData/diagnostics/tests/coroutines/callableReference")
             @TestDataPath("$PROJECT_ROOT")
             @RunWith(JUnit3RunnerWithInners.class)