Refine definition of whether call is suspension point

It's not about all calls to suspend functions, they may called from non-coroutine code

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

@@ -2560,8 +2560,8 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
             @NotNull CallGenerator callGenerator,
             @NotNull ArgumentGenerator argumentGenerator
     ) {
-        boolean isCallToSuspendFunction = isCallToSuspendFunction(resolvedCall);
+        boolean isSuspensionPoint = CoroutineCodegenUtilKt.isSuspensionPoint(resolvedCall);
-        if (isCallToSuspendFunction) {
+        if (isSuspensionPoint) {
             // Inline markers are used to spill the stack before coroutine suspension
             addInlineMarker(v, true);
         }
@@ -2597,7 +2597,7 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
             }
         }
 
-        if (isCallToSuspendFunction) {
+        if (isSuspensionPoint) {
             v.invokestatic(
                     CoroutineCodegenUtilKt.SUSPENSION_POINT_MARKER_OWNER,
                     CoroutineCodegenUtilKt.SUSPENSION_POINT_MARKER_NAME, "()V", false);
@@ -2605,7 +2605,7 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
 
         callGenerator.genCall(callableMethod, resolvedCall, defaultMaskWasGenerated, this);
 
-        if (isCallToSuspendFunction) {
+        if (isSuspensionPoint) {
             addInlineMarker(v, false);
         }
 
@@ -2616,11 +2616,6 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
         }
     }
 
-    private static boolean isCallToSuspendFunction(@NotNull ResolvedCall<?> resolvedCall) {
-        CallableDescriptor descriptor = resolvedCall.getResultingDescriptor();
-        return descriptor instanceof FunctionDescriptor && ((FunctionDescriptor) descriptor).isSuspend();
-    }
-
     @NotNull
     private CallGenerator getOrCreateCallGenerator(
             @NotNull CallableDescriptor descriptor,

compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/coroutineCodegenUtil.kt

@@ -28,6 +28,7 @@ import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCallImpl
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.tasks.TracingStrategy
+import org.jetbrains.kotlin.resolve.coroutine.SUSPENSION_POINT_KEY
 import org.jetbrains.kotlin.types.TypeConstructorSubstitution
 import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
 
@@ -52,13 +53,14 @@ fun ResolvedCall<*>.replaceSuspensionFunctionViewWithRealDescriptor(
         project: Project
 ): ResolvedCallWithRealDescriptor? {
     val function = candidateDescriptor as? FunctionDescriptor ?: return null
-    if (!function.isSuspend) return null
+    if (!isSuspensionPoint()) return null
 
     val initialSignatureDescriptor = function.initialSignatureDescriptor ?: return null
     val newCandidateDescriptor =
             initialSignatureDescriptor.createCustomCopy {
                 // Here we know that last parameter should be Continuation<T> where T is return type
                 setReturnType(it.valueParameters.last().type.arguments.single().type)
+                putUserData(SUSPENSION_POINT_KEY, true)
             }
 
     val newCall = ResolvedCallImpl(
@@ -85,6 +87,10 @@ fun ResolvedCall<*>.replaceSuspensionFunctionViewWithRealDescriptor(
     return ResolvedCallWithRealDescriptor(newCall, thisExpression)
 }
 
+fun ResolvedCall<*>.isSuspensionPoint() =
+        (candidateDescriptor as? FunctionDescriptor)?.let { it.isSuspend && it.getUserData(SUSPENSION_POINT_KEY) ?: false }
+        ?: false
+
 private fun FunctionDescriptor.createCustomCopy(
         copySettings: FunctionDescriptor.CopyBuilder<out FunctionDescriptor>.(FunctionDescriptor) -> FunctionDescriptor.CopyBuilder<out FunctionDescriptor>
 ): FunctionDescriptor {

compiler/resolution/src/org/jetbrains/kotlin/resolve/coroutine/coroutineUtil.kt

@@ -17,10 +17,13 @@
 package org.jetbrains.kotlin.resolve.coroutine
 
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
 import org.jetbrains.kotlin.resolve.scopes.receivers.ExtensionReceiver
 import org.jetbrains.kotlin.types.KotlinType
 
+val SUSPENSION_POINT_KEY: FunctionDescriptor.UserDataKey<Boolean> = object : FunctionDescriptor.UserDataKey<Boolean> {}
+
 // Returns suspension function as it's visible within coroutines:
 // E.g. `fun <V> await(f: CompletableFuture<V>): V` instead of `fun <V> await(f: CompletableFuture<V>, machine: Continuation<V>): Unit`
 fun SimpleFunctionDescriptor.createCoroutineSuspensionFunctionView(): SimpleFunctionDescriptor? {
@@ -37,6 +40,7 @@ fun SimpleFunctionDescriptor.createCoroutineSuspensionFunctionView(): SimpleFunc
         setOriginal(newOriginal)
         setValueParameters(valueParameters.subList(0, valueParameters.size - 1))
         setSignatureChange()
+        putUserData(SUSPENSION_POINT_KEY, true)
     }.build()!!
 }
 

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

@@ -0,0 +1,36 @@
+class Controller {
+    suspend fun suspendHere(x: Continuation<String>) {
+        x.resume("OK")
+    }
+}
+
+fun builder(coroutine c: Controller.() -> Continuation<Unit>) {
+    c(Controller()).resume(Unit)
+}
+
+fun box(): String {
+    var result = "fail"
+
+    val lambda: Controller.() -> Continuation<Unit> = {
+        object : Continuation<Any?> {
+            override fun resume(data: Any?) {
+                if (data == Unit) {
+                    suspendHere(this)
+                    return
+                }
+
+                if (data != "OK") {
+                    throw java.lang.RuntimeException("fail: $data")
+                }
+
+                result = "OK"
+            }
+
+            override fun resumeWithException(exception: Throwable) = throw exception
+        }
+    }
+
+    builder(lambda)
+
+    return result
+}

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

@@ -0,0 +1,23 @@
+class Controller {
+    suspend fun suspendHere(x: Continuation<String>) {
+        suspendThere(x)
+    }
+
+    suspend fun suspendThere(x: Continuation<String>) {
+        x.resume("OK")
+    }
+}
+
+fun builder(coroutine c: Controller.() -> Continuation<Unit>) {
+    c(Controller()).resume(Unit)
+}
+
+fun box(): String {
+    var result = ""
+
+    builder {
+        result = suspendHere()
+    }
+
+    return result
+}

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

@@ -4129,6 +4129,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("manualContinuationImpl.kt")
+        public void testManualContinuationImpl() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/manualContinuationImpl.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("nonLocalReturnFromInlineLambda.kt")
         public void testNonLocalReturnFromInlineLambda() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/nonLocalReturnFromInlineLambda.kt");
@@ -4165,6 +4171,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("suspendDelegation.kt")
+        public void testSuspendDelegation() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendDelegation.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("suspendFromInlineLambda.kt")
         public void testSuspendFromInlineLambda() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendFromInlineLambda.kt");