Support calling experimental coroutines API in JVM

The support is very limited, though

 #KT-25683 Fixed

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

@@ -298,6 +298,8 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
 
             StackValue stackValue = selector.accept(visitor, receiver);
 
+            stackValue = suspendFunctionTypeWrapperIfNeeded(selector, stackValue);
+
             RuntimeAssertionInfo runtimeAssertionInfo = null;
             if (selector instanceof KtExpression) {
                 KtExpression expression = (KtExpression) selector;
@@ -321,6 +323,32 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
         }
     }
 
+    private StackValue suspendFunctionTypeWrapperIfNeeded(KtElement selector, StackValue stackValue) {
+        Type functionTypeForWrapper =
+                selector instanceof KtExpression
+                ? bindingContext.get(CodegenBinding.FUNCTION_TYPE_FOR_SUSPEND_WRAPPER, (KtExpression) selector)
+                : null;
+
+        if (functionTypeForWrapper == null) return stackValue;
+
+        StackValue stackValueToWrap = stackValue;
+        stackValue = new StackValue(stackValue.type, stackValue.kotlinType) {
+            @Override
+            public void putSelector(
+                    @NotNull Type type, @Nullable KotlinType kotlinType, @NotNull InstructionAdapter v
+            ) {
+                stackValueToWrap.put(functionTypeForWrapper, null, v);
+                invokeCoroutineMigrationMethod(
+                        v,
+                        "toExperimentalSuspendFunction",
+                        Type.getMethodDescriptor(functionTypeForWrapper, functionTypeForWrapper)
+                );
+                coerce(functionTypeForWrapper, type, v);
+            }
+        };
+        return stackValue;
+    }
+
     public StackValue gen(KtElement expr) {
         StackValue tempVar = tempVariables.get(expr);
         return tempVar != null ? tempVar : genQualified(StackValue.none(), expr);
@@ -2299,9 +2327,43 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
                 coroutineInstanceValueForSuspensionPoint != null
                 ? coroutineInstanceValueForSuspensionPoint
                 : getContinuationParameterFromEnclosingSuspendFunction(resolvedCall);
+
+        if (coroutineInstanceValue != null && needsExperimentalCoroutinesWrapper(resolvedCall.getCandidateDescriptor())) {
+            StackValue releaseContinuation = coroutineInstanceValue;
+            coroutineInstanceValue = new StackValue(CoroutineCodegenUtilKt.EXPERIMENTAL_CONTINUATION_ASM_TYPE) {
+                @Override
+                public void putSelector(
+                        @NotNull Type type, @Nullable KotlinType kotlinType, @NotNull InstructionAdapter v
+                ) {
+                    releaseContinuation.put(CoroutineCodegenUtilKt.RELEASE_CONTINUATION_ASM_TYPE, v);
+                    invokeCoroutineMigrationMethod(
+                            v,
+                            "toExperimentalContinuation",
+                            Type.getMethodDescriptor(
+                                    CoroutineCodegenUtilKt.EXPERIMENTAL_CONTINUATION_ASM_TYPE,
+                                    CoroutineCodegenUtilKt.RELEASE_CONTINUATION_ASM_TYPE
+                            )
+                    );
+                }
+            };
+        }
+
         tempVariables.put(continuationExpression, coroutineInstanceValue);
     }
 
+    private static void invokeCoroutineMigrationMethod(
+            @NotNull InstructionAdapter v,
+            String methodName,
+            String descriptor
+    ) {
+        v.invokestatic(
+                "kotlin/coroutines/experimental/migration/CoroutinesMigrationKt",
+                methodName,
+                descriptor,
+                false
+        );
+    }
+
     private StackValue getContinuationParameterFromEnclosingSuspendFunction(@NotNull ResolvedCall<?> resolvedCall) {
         FunctionDescriptor enclosingSuspendFunction =
                 bindingContext.get(BindingContext.ENCLOSING_SUSPEND_FUNCTION_FOR_SUSPEND_FUNCTION_CALL, resolvedCall.getCall());

compiler/backend/src/org/jetbrains/kotlin/codegen/binding/CodegenAnnotatingVisitor.java

@@ -15,6 +15,7 @@ import kotlin.Pair;
 import kotlin.collections.CollectionsKt;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.kotlin.builtins.FunctionTypesKt;
 import org.jetbrains.kotlin.builtins.ReflectionTypes;
 import org.jetbrains.kotlin.cfg.WhenChecker;
 import org.jetbrains.kotlin.codegen.*;
@@ -52,6 +53,7 @@ import org.jetbrains.kotlin.resolve.constants.EnumValue;
 import org.jetbrains.kotlin.resolve.constants.NullValue;
 import org.jetbrains.kotlin.resolve.descriptorUtil.DescriptorUtilsKt;
 import org.jetbrains.kotlin.resolve.jvm.RuntimeAssertionsOnDeclarationBodyChecker;
+import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver;
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue;
 import org.jetbrains.kotlin.resolve.scopes.receivers.TransientReceiver;
 import org.jetbrains.kotlin.types.KotlinType;
@@ -61,6 +63,7 @@ import org.jetbrains.org.objectweb.asm.Type;
 import java.util.*;
 
 import static org.jetbrains.kotlin.codegen.binding.CodegenBinding.*;
+import static org.jetbrains.kotlin.codegen.inline.InlineCodegenUtilsKt.NUMBERED_FUNCTION_PREFIX;
 import static org.jetbrains.kotlin.lexer.KtTokens.*;
 import static org.jetbrains.kotlin.name.SpecialNames.safeIdentifier;
 import static org.jetbrains.kotlin.resolve.BindingContext.*;
@@ -645,6 +648,54 @@ class CodegenAnnotatingVisitor extends KtVisitorVoid {
         super.visitCallExpression(expression);
         checkSamCall(expression);
         checkCrossinlineSuspendCall(expression);
+        recordSuspendFunctionTypeWrapperForArguments(expression);
+    }
+
+    private void recordSuspendFunctionTypeWrapperForArguments(@NotNull KtCallExpression expression) {
+        ResolvedCall<?> call = CallUtilKt.getResolvedCall(expression, bindingContext);
+        if (call == null) return;
+
+        CallableDescriptor descriptor = call.getResultingDescriptor();
+        if (!CodegenUtilKt.needsExperimentalCoroutinesWrapper(descriptor)) return;
+
+        List<ResolvedValueArgument> argumentsByIndex = call.getValueArgumentsByIndex();
+        if (argumentsByIndex == null) return;
+
+        for (ValueParameterDescriptor parameter : descriptor.getValueParameters()) {
+            ResolvedValueArgument resolvedValueArgument = argumentsByIndex.get(parameter.getIndex());
+            if (!(resolvedValueArgument instanceof ExpressionValueArgument)) continue;
+            ValueArgument valueArgument = ((ExpressionValueArgument) resolvedValueArgument).getValueArgument();
+            if (valueArgument == null) continue;
+            KtExpression argumentExpression = valueArgument.getArgumentExpression();
+            if (argumentExpression == null) continue;
+
+            recordSuspendFunctionTypeWrapperForArgument(parameter, argumentExpression);
+        }
+
+        ReceiverValue receiver = call.getExtensionReceiver();
+        if (descriptor.getExtensionReceiverParameter() != null && receiver instanceof ExpressionReceiver) {
+            recordSuspendFunctionTypeWrapperForArgument(
+                    descriptor.getExtensionReceiverParameter(),
+                    ((ExpressionReceiver) receiver).getExpression()
+            );
+        }
+    }
+
+    private void recordSuspendFunctionTypeWrapperForArgument(ParameterDescriptor parameter, KtExpression argumentExpression) {
+        if (FunctionTypesKt.isSuspendFunctionTypeOrSubtype(parameter.getType())) {
+
+            // SuspendFunctionN type is mapped to is mapped to FunctionTypeN+1, but we also need to remove an argument for return type
+            // So, it could be parameter.getType().getArguments().size() + 1 - 1
+            int functionTypeArity = parameter.getType().getArguments().size();
+
+            Type functionType = Type.getObjectType(NUMBERED_FUNCTION_PREFIX + functionTypeArity);
+
+            bindingTrace.record(
+                    FUNCTION_TYPE_FOR_SUSPEND_WRAPPER,
+                    argumentExpression,
+                    functionType
+            );
+        }
     }
 
     private void checkCrossinlineSuspendCall(@NotNull KtCallExpression expression) {

compiler/backend/src/org/jetbrains/kotlin/codegen/binding/CodegenBinding.java

@@ -43,6 +43,7 @@ public class CodegenBinding {
     private static final WritableSlice<ClassDescriptor, Collection<ClassDescriptor>> INNER_CLASSES = Slices.createSimpleSlice();
 
     public static final WritableSlice<KtExpression, SamType> SAM_VALUE = Slices.createSimpleSlice();
+    public static final WritableSlice<KtExpression, Type> FUNCTION_TYPE_FOR_SUSPEND_WRAPPER = Slices.createSimpleSlice();
 
     public static final WritableSlice<KtCallElement, KtExpression> SAM_CONSTRUCTOR_TO_ARGUMENT = Slices.createSimpleSlice();
 

compiler/backend/src/org/jetbrains/kotlin/codegen/codegenUtil.kt

@@ -41,6 +41,8 @@ import org.jetbrains.kotlin.resolve.isInlineClassType
 import org.jetbrains.kotlin.resolve.jvm.JvmClassName
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin
 import org.jetbrains.kotlin.resolve.scopes.receivers.TransientReceiver
+import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedMemberDescriptor
+import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedMemberDescriptor.CoroutinesCompatibilityMode
 import org.jetbrains.kotlin.types.ErrorUtils
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeUtils
@@ -431,4 +433,7 @@ fun KotlinType.isInlineClassTypeWithPrimitiveEquality(): Boolean {
 
     // TODO support other inline classes that can be compared as underlying primitives
     return false
-}
+}
+
+fun CallableDescriptor.needsExperimentalCoroutinesWrapper() =
+    (this as? DeserializedMemberDescriptor)?.coroutinesExperimentalCompatibilityMode == CoroutinesCompatibilityMode.NEEDS_WRAPPER

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

@@ -295,7 +295,7 @@ fun <D : FunctionDescriptor> D.createCustomCopy(
 }
 
 private fun FunctionDescriptor.getContinuationParameterTypeOfSuspendFunction(isReleaseCoroutines: Boolean) =
-    module.getContinuationOfTypeOrAny(returnType!!, isReleaseCoroutines)
+    module.getContinuationOfTypeOrAny(returnType!!, if (this.needsExperimentalCoroutinesWrapper()) false else isReleaseCoroutines)
 
 fun ModuleDescriptor.getSuccessOrFailure(kotlinType: KotlinType) =
     module.resolveTopLevelClass(
@@ -498,3 +498,9 @@ fun FunctionDescriptor.isSuspendLambdaOrLocalFunction() = this.isSuspend && when
 }
 
 fun FunctionDescriptor.isLocalSuspendFunctionNotSuspendLambda() = isSuspendLambdaOrLocalFunction() && this !is AnonymousFunctionDescriptor
+
+@JvmField
+val EXPERIMENTAL_CONTINUATION_ASM_TYPE = DescriptorUtils.CONTINUATION_INTERFACE_FQ_NAME_EXPERIMENTAL.topLevelClassAsmType()
+
+@JvmField
+val RELEASE_CONTINUATION_ASM_TYPE = DescriptorUtils.CONTINUATION_INTERFACE_FQ_NAME_RELEASE.topLevelClassAsmType()

compiler/testData/compileKotlinAgainstKotlin/coroutines/builder.kt

@@ -1,9 +1,18 @@
 // FILE: A.kt
 // LANGUAGE_VERSION: 1.2
 // TODO: Unmute when automatic conversion experimental <-> release will be implemented
-// IGNORE_BACKEND: JVM, JS, NATIVE, JVM_IR, JS_IR
+// IGNORE_BACKEND: JS, NATIVE, JVM_IR, JS_IR
 import kotlin.coroutines.experimental.*
 
+var callback: (() -> Unit)? = null
+fun join() {
+    while (callback != null) {
+        val x = callback!!
+        callback = null
+        x()
+    }
+}
+
 fun builder1(c: suspend () -> String): String {
     var res = "FAIL"
     c.startCoroutine(object : Continuation<String> {
@@ -15,6 +24,7 @@ fun builder1(c: suspend () -> String): String {
             throw exception
         }
     })
+    join()
     return res
 }
 
@@ -29,23 +39,31 @@ fun builder2(c: suspend String.() -> String): String {
             throw exception
         }
     })
+    join()
     return res
 }
 
 // FILE: B.kt
 // LANGUAGE_VERSION: 1.3
-import kotlin.coroutines.experimental.*
+import kotlin.coroutines.*
+import kotlin.coroutines.intrinsics.*
 
-fun ok(continuation: Continuation<String>): Any? {
-    return "OK"
+suspend fun ok(): String {
+    return o() + k()
+}
+
+suspend fun o(): String = suspendCoroutine { continuation ->
+    callback = { continuation.resume("O") }
+}
+
+suspend fun k(): String = suspendCoroutine { continuation ->
+    callback = { continuation.resume("K") }
 }
 
 fun box(): String {
-//    if (builder1(::ok) != "OK") return "FAIL 1"
-//    if (builder1 { cont: Continuation<String> -> "OK" } != "OK") return "FAIL 2"
-//    if (builder1(fun (cont: Continuation<String>): Any? = "OK") != "OK") return "FAIL 3"
-//
-//    if (builder2 { cont: Continuation<String> -> this + "K" } != "OK") return "FAIL 5"
-//    if (builder2(fun String.(cont: Continuation<String>): Any? = this + "K") != "OK") return "FAIL 6"
+    if (builder1(::ok) != "OK") return "FAIL 1"
+    if (builder1 { ok() } != "OK") return "FAIL 2"
+
+    if (builder2 { this + k() } != "OK") return "FAIL 5"
     return "OK"
 }

compiler/testData/compileKotlinAgainstKotlin/coroutines/receiver.kt

@@ -1,9 +1,18 @@
 // FILE: A.kt
 // LANGUAGE_VERSION: 1.2
 // TODO: Unmute when automatic conversion experimental <-> release will be implemented
-// IGNORE_BACKEND: JVM, JS, NATIVE, JVM_IR, JS_IR
+// IGNORE_BACKEND: JS, NATIVE, JVM_IR, JS_IR
 import kotlin.coroutines.experimental.*
 
+var callback: (() -> Unit)? = null
+fun join() {
+    while (callback != null) {
+        val x = callback!!
+        callback = null
+        x()
+    }
+}
+
 fun (suspend () -> String).builder(): String {
     var res = "FAIL"
     startCoroutine(object : Continuation<String> {
@@ -15,20 +24,31 @@ fun (suspend () -> String).builder(): String {
             throw exception
         }
     })
+
+    join()
+
     return res
 }
 
 // FILE: B.kt
 // LANGUAGE_VERSION: 1.3
-import kotlin.coroutines.experimental.*
+import kotlin.coroutines.*
+import kotlin.coroutines.intrinsics.*
 
-fun ok(continuation: Continuation<String>): Any? {
-    return "OK"
+suspend fun ok(): String {
+    return o() + k()
+}
+
+suspend fun o(): String = suspendCoroutine { continuation ->
+    callback = { continuation.resume("O") }
+}
+
+suspend fun k(): String = suspendCoroutine { continuation ->
+    callback = { continuation.resume("K") }
 }
 
 fun box(): String {
-//    if ((::ok).builder() != "OK") return "FAIL 1"
-//    if (({ cont: Continuation<String> -> "OK" }).builder() != "OK") return "FAIL 2"
-//    if ((fun (cont: Continuation<String>): Any? = "OK").builder() != "OK") return "FAIL 3"
+    if ((::ok).builder() != "OK") return "FAIL 1"
+    if (suspend { ok() }.builder() != "OK") return "FAIL 2"
     return "OK"
 }

compiler/testData/compileKotlinAgainstKotlin/coroutines/simple.kt

@@ -1,49 +1,42 @@
 // FILE: A.kt
 // LANGUAGE_VERSION: 1.2
 // TODO: Unmute when automatic conversion experimental <-> release will be implemented
-// IGNORE_BACKEND: JVM, JS, NATIVE, JVM_IR, JS_IR
+// IGNORE_BACKEND: JS, NATIVE, JVM_IR, JS_IR
+import kotlin.coroutines.experimental.*
+import kotlin.coroutines.experimental.intrinsics.*
 
-suspend fun dummy() = "OK"
-
-suspend fun String.dummy() = this + "K"
-
-suspend fun String.dummy(s: String) = this + s
-
-class C {
-    suspend fun dummy() = "OK"
+var callback: (() -> Unit)? = null
+suspend fun dummy(): String = suspendCoroutine {
+    callback = { it.resume("OK") }
 }
 
-class WithNested {
-    class Nested {
-        suspend fun dummy() = "OK"
-    }
-}
-
-class WithInner {
-    inner class Inner {
-        suspend fun dummy() = "OK"
-    }
-}
-
-
 // FILE: B.kt
 // LANGUAGE_VERSION: 1.3
-import kotlin.coroutines.experimental.*
+import kotlin.coroutines.*
+import kotlin.coroutines.intrinsics.*
+
+fun builder(x: suspend () -> Unit) {
+    x.startCoroutine(object : Continuation<Any?> {
+        override val context: CoroutineContext = EmptyCoroutineContext
+
+        override fun resumeWith(result: SuccessOrFailure<Any?>) {
+            result.getOrThrow()
+        }
+    })
+
+    while (callback != null) {
+        val x = callback!!
+        callback = null
+        x()
+    }
+}
 
 fun box(): String {
-//    val continuation = object : Continuation<String> {
-//        override val context = EmptyCoroutineContext
-//        override fun resume(value: String) {
-//        }
-//        override fun resumeWithException(exception: Throwable) {
-//            throw exception
-//        }
-//    }
-//    if (dummy(continuation) != "OK") return "FAIL 1"
-//    if ("O".dummy(continuation) != "OK") return "FAIL 2"
-//    if ("O".dummy("K", continuation) != "OK") return "FAIL 3"
-//    if (C().dummy(continuation) != "OK") return "FAIL 4"
-//    if (WithNested.Nested().dummy(continuation) != "OK") return "FAIL 5"
-//    if (WithInner().Inner().dummy(continuation) != "OK") return "FAIL 6"
-    return "OK"
+    var res = "fail"
+
+    builder {
+        res = dummy()
+    }
+
+    return res
 }

libraries/stdlib/coroutines/jvm/src/kotlin/coroutines/experimental/migration/CoroutinesMigration.kt

@@ -12,6 +12,8 @@ import kotlin.coroutines.experimental.EmptyCoroutineContext as ExperimentalEmpty
 import kotlin.coroutines.experimental.ContinuationInterceptor as ExperimentalContinuationInterceptor
 import kotlin.coroutines.experimental.intrinsics.COROUTINE_SUSPENDED as EXPERIMENTAL_COROUTINE_SUSPENDED
 import kotlin.coroutines.*
+import kotlin.coroutines.intrinsics.*
+
 
 /**
  * Converts [Continuation] to [ExperimentalContinuation].
@@ -109,3 +111,38 @@ private class ContinuationInterceptorMigration(val interceptor: ExperimentalCont
         interceptor.interceptContinuation(continuation.toExperimentalContinuation()).toContinuation()
 }
 
+internal fun <R> ((Continuation<R>) -> Any?).toExperimentalSuspendFunction(): (ExperimentalContinuation<R>) -> Any? =
+    ExperimentalSuspendFunction0Migration(this)
+
+internal fun <T1, R> ((T1, Continuation<R>) -> Any?).toExperimentalSuspendFunction(): (T1, ExperimentalContinuation<R>) -> Any? =
+    ExperimentalSuspendFunction1Migration(this)
+
+internal fun <T1, T2, R> ((T1, T2, Continuation<R>) -> Any?).toExperimentalSuspendFunction(): (T1, T2, ExperimentalContinuation<R>) -> Any? =
+    ExperimentalSuspendFunction2Migration(this)
+
+private class ExperimentalSuspendFunction0Migration<R>(
+    val function: (Continuation<R>) -> Any?
+) : (ExperimentalContinuation<R>) -> Any? {
+    override fun invoke(continuation: ExperimentalContinuation<R>): Any? {
+        val result = function(continuation.toContinuation())
+        return if (result === COROUTINE_SUSPENDED) EXPERIMENTAL_COROUTINE_SUSPENDED else result
+    }
+}
+
+private class ExperimentalSuspendFunction1Migration<T1, R>(
+    val function: (T1, Continuation<R>) -> Any?
+) : (T1, ExperimentalContinuation<R>) -> Any? {
+    override fun invoke(t1: T1, continuation: ExperimentalContinuation<R>): Any? {
+        val result = function(t1, continuation.toContinuation())
+        return if (result === COROUTINE_SUSPENDED) EXPERIMENTAL_COROUTINE_SUSPENDED else result
+    }
+}
+
+private class ExperimentalSuspendFunction2Migration<T1, T2, R>(
+    val function: (T1, T2, Continuation<R>) -> Any?
+) : (T1, T2, ExperimentalContinuation<R>) -> Any? {
+    override fun invoke(t1: T1, t2: T2, continuation: ExperimentalContinuation<R>): Any? {
+        val result = function(t1, t2, continuation.toContinuation())
+        return if (result === COROUTINE_SUSPENDED) EXPERIMENTAL_COROUTINE_SUSPENDED else result
+    }
+}