Fix codegen issue of safe-qualified suspension points

#KT-15527 Fixed

compiler/backend/src/org/jetbrains/kotlin/codegen/ExpressionCodegen.java

@@ -2912,7 +2912,9 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
             boolean isSuspensionPoint,
             boolean isConstructor
     ) {
-        if (isSuspensionPoint) {
+        boolean isSafeCall = receiver instanceof StackValue.SafeCall;
+
+        if (isSuspensionPoint && !isSafeCall) {
             // Inline markers are used to spill the stack before coroutine suspension
             addInlineMarker(v, true);
         }
@@ -2920,6 +2922,35 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
         if (!isConstructor) { // otherwise already
             receiver = StackValue.receiver(resolvedCall, receiver, this, callableMethod);
             receiver.put(receiver.type, v);
+
+            // In regular cases we add an inline marker just before receiver is loaded (to spill the stack before a suspension)
+            // But in case of safe call things we get the following bytecode:
+
+            // ---- inlineMarkerBefore()
+            // LOAD $receiver
+            // IFNULL L1
+            // ---- load the rest of the arguments
+            // INVOKEVIRTUAL suspendCall()
+            // ---- inlineMarkerBefore()
+            // GOTO L2
+            // L1
+            // ACONST_NULL
+            // L2
+            // ...
+            //
+            // The problem is that the stack before the call is not restored in case of null receiver.
+            // The solution is to spill stack just after receiver is loaded (after IFNULL) in case of safe call.
+            // But the problem is that we should leave the receiver itself on the stack, so we store it in a temporary variable.
+            if (isSuspensionPoint && isSafeCall) {
+                int tmpVar = myFrameMap.enterTemp(receiver.type);
+
+                v.store(tmpVar, receiver.type);
+                addInlineMarker(v, true);
+                v.load(tmpVar, receiver.type);
+
+                myFrameMap.leaveTemp(receiver.type);
+            }
+
             callableMethod.afterReceiverGeneration(v);
         }
     }

compiler/testData/codegen/box/coroutines/suspensionInsideSafeCall.kt

@@ -0,0 +1,40 @@
+// WITH_RUNTIME
+// WITH_COROUTINES
+// IGNORE_BACKEND: JS
+import kotlin.coroutines.*
+import kotlin.coroutines.intrinsics.*
+
+class TestClass {
+    suspend fun suspendHere(): String = suspendCoroutineOrReturn { x ->
+        x.resume("K")
+        SUSPENDED_MARKER
+    }
+}
+
+fun builder(c: suspend () -> Unit) {
+    c.startCoroutine(EmptyContinuation)
+}
+
+fun foo(x: String, y: String?) = x + (y ?: "")
+
+fun box(): String {
+    var result = ""
+
+    var instance: TestClass? = null
+
+    builder {
+        result = foo("OK", instance?.suspendHere())
+    }
+
+    if (result != "OK") return "fail 1: $result"
+
+    result = ""
+    instance = TestClass()
+    builder {
+        result = foo("O", instance?.suspendHere())
+    }
+
+    if (result != "OK") return "fail 2: $result"
+
+    return "OK"
+}

compiler/testData/codegen/box/coroutines/suspensionInsideSafeCallWithElvis.kt

@@ -0,0 +1,37 @@
+// WITH_RUNTIME
+// WITH_COROUTINES
+// IGNORE_BACKEND: JS
+import kotlin.coroutines.*
+import kotlin.coroutines.intrinsics.*
+
+class TestClass {
+    suspend fun toInt(): Int = suspendCoroutineOrReturn { x ->
+        x.resume(14)
+        SUSPENDED_MARKER
+    }
+}
+
+fun builder(c: suspend () -> Unit) {
+    c.startCoroutine(EmptyContinuation)
+}
+
+fun box(): String {
+    var result = 0
+
+    var instance: TestClass? = null
+
+    builder {
+        result = 42 + (instance?.toInt() ?: 0)
+    }
+
+    if (result != 42) return "fail 1: $result"
+
+    instance = TestClass()
+    builder {
+        result = 42 + (instance?.toInt() ?: 0)
+    }
+
+    if (result != 56) return "fail 2: $result"
+
+    return "OK"
+}

compiler/testData/codegen/light-analysis/coroutines/suspensionInsideSafeCall.txt

@@ -0,0 +1,37 @@
+@kotlin.Metadata
+public final class CoroutineUtilKt {
+    public final static @org.jetbrains.annotations.NotNull method handleExceptionContinuation(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): kotlin.coroutines.Continuation
+    public final static @org.jetbrains.annotations.NotNull method handleResultContinuation(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): kotlin.coroutines.Continuation
+}
+
+@kotlin.Metadata
+public class EmptyContinuation {
+    public final static field Companion: EmptyContinuation.Companion
+    private final @org.jetbrains.annotations.NotNull field context: kotlin.coroutines.CoroutineContext
+    inner class EmptyContinuation/Companion
+    public @synthetic.kotlin.jvm.GeneratedByJvmOverloads method <init>(): void
+    public method <init>(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.CoroutineContext): void
+    public synthetic method <init>(p0: kotlin.coroutines.CoroutineContext, p1: int, p2: kotlin.jvm.internal.DefaultConstructorMarker): void
+    public @org.jetbrains.annotations.NotNull method getContext(): kotlin.coroutines.CoroutineContext
+    public method resume(@org.jetbrains.annotations.Nullable p0: java.lang.Object): void
+    public method resumeWithException(@org.jetbrains.annotations.NotNull p0: java.lang.Throwable): void
+}
+
+@kotlin.Metadata
+public final static class EmptyContinuation/Companion {
+    inner class EmptyContinuation/Companion
+    private method <init>(): void
+}
+
+@kotlin.Metadata
+public final class SuspensionInsideSafeCallKt {
+    public final static @org.jetbrains.annotations.NotNull method box(): java.lang.String
+    public final static method builder(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): void
+    public final static @org.jetbrains.annotations.NotNull method foo(@org.jetbrains.annotations.NotNull p0: java.lang.String, @org.jetbrains.annotations.Nullable p1: java.lang.String): java.lang.String
+}
+
+@kotlin.Metadata
+public final class TestClass {
+    public method <init>(): void
+    public final @org.jetbrains.annotations.Nullable method suspendHere(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.Continuation): java.lang.Object
+}

compiler/testData/codegen/light-analysis/coroutines/suspensionInsideSafeCallWithElvis.txt

@@ -0,0 +1,36 @@
+@kotlin.Metadata
+public final class CoroutineUtilKt {
+    public final static @org.jetbrains.annotations.NotNull method handleExceptionContinuation(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): kotlin.coroutines.Continuation
+    public final static @org.jetbrains.annotations.NotNull method handleResultContinuation(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): kotlin.coroutines.Continuation
+}
+
+@kotlin.Metadata
+public class EmptyContinuation {
+    public final static field Companion: EmptyContinuation.Companion
+    private final @org.jetbrains.annotations.NotNull field context: kotlin.coroutines.CoroutineContext
+    inner class EmptyContinuation/Companion
+    public @synthetic.kotlin.jvm.GeneratedByJvmOverloads method <init>(): void
+    public method <init>(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.CoroutineContext): void
+    public synthetic method <init>(p0: kotlin.coroutines.CoroutineContext, p1: int, p2: kotlin.jvm.internal.DefaultConstructorMarker): void
+    public @org.jetbrains.annotations.NotNull method getContext(): kotlin.coroutines.CoroutineContext
+    public method resume(@org.jetbrains.annotations.Nullable p0: java.lang.Object): void
+    public method resumeWithException(@org.jetbrains.annotations.NotNull p0: java.lang.Throwable): void
+}
+
+@kotlin.Metadata
+public final static class EmptyContinuation/Companion {
+    inner class EmptyContinuation/Companion
+    private method <init>(): void
+}
+
+@kotlin.Metadata
+public final class SuspensionInsideSafeCallWithElvisKt {
+    public final static @org.jetbrains.annotations.NotNull method box(): java.lang.String
+    public final static method builder(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): void
+}
+
+@kotlin.Metadata
+public final class TestClass {
+    public method <init>(): void
+    public final @org.jetbrains.annotations.Nullable method toInt(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.Continuation): java.lang.Object
+}

compiler/tests-ir-jvm/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -4823,6 +4823,18 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             doTest(fileName);
         }
 
+        @TestMetadata("suspensionInsideSafeCall.kt")
+        public void testSuspensionInsideSafeCall() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspensionInsideSafeCall.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("suspensionInsideSafeCallWithElvis.kt")
+        public void testSuspensionInsideSafeCallWithElvis() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspensionInsideSafeCallWithElvis.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("tryCatchFinallyWithHandleResult.kt")
         public void testTryCatchFinallyWithHandleResult() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/tryCatchFinallyWithHandleResult.kt");

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -4823,6 +4823,18 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("suspensionInsideSafeCall.kt")
+        public void testSuspensionInsideSafeCall() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspensionInsideSafeCall.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("suspensionInsideSafeCallWithElvis.kt")
+        public void testSuspensionInsideSafeCallWithElvis() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspensionInsideSafeCallWithElvis.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("tryCatchFinallyWithHandleResult.kt")
         public void testTryCatchFinallyWithHandleResult() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/tryCatchFinallyWithHandleResult.kt");

compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeCodegenTestGenerated.java

@@ -4823,6 +4823,18 @@ public class LightAnalysisModeCodegenTestGenerated extends AbstractLightAnalysis
             doTest(fileName);
         }
 
+        @TestMetadata("suspensionInsideSafeCall.kt")
+        public void testSuspensionInsideSafeCall() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspensionInsideSafeCall.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("suspensionInsideSafeCallWithElvis.kt")
+        public void testSuspensionInsideSafeCallWithElvis() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspensionInsideSafeCallWithElvis.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("tryCatchFinallyWithHandleResult.kt")
         public void testTryCatchFinallyWithHandleResult() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/tryCatchFinallyWithHandleResult.kt");

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -5562,6 +5562,30 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             doTest(fileName);
         }
 
+        @TestMetadata("suspensionInsideSafeCall.kt")
+        public void testSuspensionInsideSafeCall() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspensionInsideSafeCall.kt");
+            try {
+                doTest(fileName);
+            }
+            catch (Throwable ignore) {
+                return;
+            }
+            throw new AssertionError("Looks like this test can be unmuted. Remove IGNORE_BACKEND directive for that.");
+        }
+
+        @TestMetadata("suspensionInsideSafeCallWithElvis.kt")
+        public void testSuspensionInsideSafeCallWithElvis() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspensionInsideSafeCallWithElvis.kt");
+            try {
+                doTest(fileName);
+            }
+            catch (Throwable ignore) {
+                return;
+            }
+            throw new AssertionError("Looks like this test can be unmuted. Remove IGNORE_BACKEND directive for that.");
+        }
+
         @TestMetadata("tryCatchFinallyWithHandleResult.kt")
         public void testTryCatchFinallyWithHandleResult() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/tryCatchFinallyWithHandleResult.kt");