JS: support new coroutine convention (see KT-15058)

compiler/backend-common/src/org/jetbrains/kotlin/backend/common/commonCoroutineCodegenUtil.kt

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2010-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.backend.common
+
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.incremental.components.NoLookupLocation
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
+
+val SUSPEND_WITH_CURRENT_CONTINUATION_NAME = Name.identifier("suspendWithCurrentContinuation")
+
+fun FunctionDescriptor.getBuiltInSuspendWithCurrentContinuation() =
+            builtIns.builtInsPackageFragments.singleOrNull { it.fqName == KotlinBuiltIns.COROUTINES_PACKAGE_FQ_NAME }
+                    ?.getMemberScope()
+                    ?.getContributedFunctions(SUSPEND_WITH_CURRENT_CONTINUATION_NAME, NoLookupLocation.FROM_BACKEND)
+                    ?.singleOrNull()
+
+

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

@@ -17,6 +17,9 @@
 package org.jetbrains.kotlin.codegen.coroutines
 
 import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.backend.common.SUSPEND_WITH_CURRENT_CONTINUATION_NAME
+import org.jetbrains.kotlin.backend.common.findInterceptResume
+import org.jetbrains.kotlin.backend.common.getBuiltInSuspendWithCurrentContinuation
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.codegen.binding.CodegenBinding
 import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
@@ -246,17 +249,10 @@ fun KotlinType.hasNoinlineInterceptResume() =
 fun findOperatorInController(controllerType: KotlinType, name: Name): SimpleFunctionDescriptor? =
         controllerType.memberScope.getContributedFunctions(name, NoLookupLocation.FROM_BACKEND).singleOrNull { it.isOperator }
 
-val SUSPEND_WITH_CURRENT_CONTINUATION_NAME = Name.identifier("suspendWithCurrentContinuation")
-
 fun FunctionDescriptor.isBuiltInSuspendWithCurrentContinuation(): Boolean {
     if (name != SUSPEND_WITH_CURRENT_CONTINUATION_NAME) return false
 
-    val originalDeclaration =
-            builtIns.builtInsPackageFragments.singleOrNull { it.fqName == KotlinBuiltIns.COROUTINES_PACKAGE_FQ_NAME }
-                    ?.getMemberScope()
-                    ?.getContributedFunctions(SUSPEND_WITH_CURRENT_CONTINUATION_NAME, NoLookupLocation.FROM_BACKEND)
-                    ?.singleOrNull()
-                    ?: return false
+    val originalDeclaration = getBuiltInSuspendWithCurrentContinuation() ?: return false
 
     return DescriptorEquivalenceForOverrides.areEquivalent(
             originalDeclaration, this.getUserData(INITIAL_DESCRIPTOR_FOR_SUSPEND_FUNCTION)

js/js.translator/src/org/jetbrains/kotlin/js/translate/callTranslator/CallTranslator.kt

@@ -18,9 +18,9 @@ package org.jetbrains.kotlin.js.translate.callTranslator
 
 import com.google.dart.compiler.backend.js.ast.JsExpression
 import com.google.dart.compiler.backend.js.ast.JsInvocation
-import com.google.dart.compiler.backend.js.ast.JsLiteral
 import com.google.dart.compiler.backend.js.ast.JsNameRef
 import com.google.dart.compiler.backend.js.ast.metadata.*
+import org.jetbrains.kotlin.backend.common.getBuiltInSuspendWithCurrentContinuation
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.VariableDescriptor
@@ -28,17 +28,21 @@ import org.jetbrains.kotlin.js.translate.context.TranslationContext
 import org.jetbrains.kotlin.js.translate.general.Translation
 import org.jetbrains.kotlin.js.translate.reference.CallArgumentTranslator
 import org.jetbrains.kotlin.js.translate.reference.CallExpressionTranslator
+import org.jetbrains.kotlin.js.translate.reference.ReferenceTranslator
 import org.jetbrains.kotlin.js.translate.utils.AnnotationsUtils
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
+import org.jetbrains.kotlin.js.translate.utils.TranslationUtils
 import org.jetbrains.kotlin.js.translate.utils.setInlineCallMetadata
 import org.jetbrains.kotlin.psi.Call.CallType
 import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.DescriptorEquivalenceForOverrides
 import org.jetbrains.kotlin.resolve.calls.callResolverUtil.isInvokeCallOnVariable
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.VariableAsFunctionResolvedCall
-import org.jetbrains.kotlin.resolve.calls.resolvedCallUtil.getImplicitReceiverValue
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind.NO_EXPLICIT_RECEIVER
+import org.jetbrains.kotlin.resolve.inline.InlineStrategy
 import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
 
@@ -123,11 +127,18 @@ private fun translateCall(context: TranslationContext,
     return translateFunctionCall(context, resolvedCall, resolvedCall, explicitReceivers)
 }
 
-private fun translateFunctionCall(context: TranslationContext,
-                                  resolvedCall: ResolvedCall<out FunctionDescriptor>,
-                                  inlineResolvedCall: ResolvedCall<out CallableDescriptor>,
-                                  explicitReceivers: ExplicitReceivers
+private fun translateFunctionCall(
+        context: TranslationContext,
+        resolvedCall: ResolvedCall<out FunctionDescriptor>,
+        inlineResolvedCall: ResolvedCall<out CallableDescriptor>,
+        explicitReceivers: ExplicitReceivers
 ): JsExpression {
+    val descriptorToCall = resolvedCall.resultingDescriptor
+    if (descriptorToCall is FunctionDescriptor && DescriptorEquivalenceForOverrides.areEquivalent(
+            descriptorToCall.getBuiltInSuspendWithCurrentContinuation(), descriptorToCall.original)) {
+        return translateCallWithContinuation(context, resolvedCall)
+    }
+
     val callExpression = context.getCallInfo(resolvedCall, explicitReceivers).translateFunctionCall()
 
     if (CallExpressionTranslator.shouldBeInlined(inlineResolvedCall.resultingDescriptor, context)) {
@@ -135,13 +146,12 @@ private fun translateFunctionCall(context: TranslationContext,
                               inlineResolvedCall.resultingDescriptor, context)
     }
 
-    if (resolvedCall.resultingDescriptor.isSuspend && resolvedCall.resultingDescriptor.initialSignatureDescriptor != null) {
+    if (resolvedCall.resultingDescriptor.isSuspend && !context.isInSuspendFunction) {
         context.currentBlock.statements += JsAstUtils.asSyntheticStatement((callExpression as JsInvocation).apply {
             isSuspend = true
             isPreSuspend = true
         })
-        val coroutineDescriptor = resolvedCall.getImplicitReceiverValue()!!.declarationDescriptor
-        val coroutineRef = context.getAliasForDescriptor(coroutineDescriptor) ?: JsLiteral.THIS
+        val coroutineRef = TranslationUtils.translateContinuationArgument(context, resolvedCall)
         return context.defineTemporary(JsNameRef("\$\$coroutineResult\$\$", coroutineRef).apply {
             sideEffects = SideEffectKind.DEPENDS_ON_STATE
             coroutineResult = true
@@ -150,6 +160,14 @@ private fun translateFunctionCall(context: TranslationContext,
     return callExpression
 }
 
+private fun translateCallWithContinuation(context: TranslationContext,resolvedCall: ResolvedCall<out FunctionDescriptor>): JsExpression {
+    val arguments = CallArgumentTranslator.translate(resolvedCall, null, context)
+    val coroutineArgument = TranslationUtils.getEnclosingContinuationParameter(context)
+    val invocation = JsInvocation(arguments.valueArguments[0], ReferenceTranslator.translateAsValueReference(coroutineArgument, context))
+    invocation.inlineStrategy = InlineStrategy.IN_PLACE
+    return invocation
+}
+
 fun computeExplicitReceiversForInvoke(
         context: TranslationContext,
         resolvedCall: ResolvedCall<out FunctionDescriptor>,

js/js.translator/src/org/jetbrains/kotlin/js/translate/context/TranslationContext.java

@@ -22,11 +22,14 @@ import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.descriptors.*;
+import org.jetbrains.kotlin.descriptors.annotations.Annotations;
+import org.jetbrains.kotlin.descriptors.impl.LocalVariableDescriptor;
 import org.jetbrains.kotlin.descriptors.impl.TypeAliasConstructorDescriptor;
 import org.jetbrains.kotlin.js.config.JsConfig;
 import org.jetbrains.kotlin.js.translate.intrinsic.Intrinsics;
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils;
 import org.jetbrains.kotlin.js.translate.utils.TranslationUtils;
+import org.jetbrains.kotlin.name.Name;
 import org.jetbrains.kotlin.psi.KtExpression;
 import org.jetbrains.kotlin.resolve.BindingContext;
 import org.jetbrains.kotlin.resolve.BindingTrace;
@@ -60,6 +63,8 @@ public class TranslationContext {
     private final DeclarationDescriptor declarationDescriptor;
     @Nullable
     private final ClassDescriptor classDescriptor;
+    @Nullable
+    private final VariableDescriptor continuationParameterDescriptor;
 
     @NotNull
     public static TranslationContext rootContext(@NotNull StaticContext staticContext, @NotNull JsFunction rootFunction) {
@@ -91,6 +96,30 @@ public class TranslationContext {
         else {
             this.classDescriptor = parent != null ? parent.classDescriptor : null;
         }
+
+        continuationParameterDescriptor = calculateContinuationParameter();
+    }
+
+    private VariableDescriptor calculateContinuationParameter() {
+        if (parent != null && parent.declarationDescriptor == declarationDescriptor) {
+            return parent.continuationParameterDescriptor;
+        }
+        if (declarationDescriptor instanceof FunctionDescriptor) {
+            FunctionDescriptor function = (FunctionDescriptor) declarationDescriptor;
+            if (function.isSuspend()) {
+                ClassDescriptor continuationDescriptor = getCurrentModule().getBuiltIns().getBuiltInClassByFqName(
+                        DescriptorUtils.CONTINUATION_INTERFACE_FQ_NAME);
+                return new LocalVariableDescriptor(
+                        declarationDescriptor,
+                        Annotations.Companion.getEMPTY(),
+                        Name.identifier("continuation"),
+                        continuationDescriptor.getDefaultType(),
+                        /* mutable = */ false,
+                        /* delegated = */ false,
+                        SourceElement.NO_SOURCE);
+            }
+        }
+        return null;
     }
 
     @NotNull
@@ -615,6 +644,15 @@ public class TranslationContext {
         return false;
     }
 
+    public boolean isInSuspendFunction() {
+        return declarationDescriptor instanceof FunctionDescriptor && ((FunctionDescriptor) declarationDescriptor).isSuspend();
+    }
+
+    @Nullable
+    public VariableDescriptor getContinuationParameterDescriptor() {
+        return continuationParameterDescriptor;
+    }
+
     @NotNull
     public ModuleDescriptor getCurrentModule() {
         return staticContext.getCurrentModule();

js/js.translator/src/org/jetbrains/kotlin/js/translate/context/UsageTracker.kt

@@ -163,9 +163,9 @@ class UsageTracker(
     }
 
     private fun DeclarationDescriptor.getJsNameForCapturedDescriptor(): JsName {
-        val suggestedName = when {
-            this is ReceiverParameterDescriptor -> getNameForCapturedReceiver()
-            this is TypeParameterDescriptor -> Namer.isInstanceSuggestedName(this)
+        val suggestedName = when (this) {
+            is ReceiverParameterDescriptor -> getNameForCapturedReceiver()
+            is TypeParameterDescriptor -> Namer.isInstanceSuggestedName(this)
 
             // Append 'closure$' prefix to avoid name clash between closure and member fields in case of local classes
             else -> {

js/js.translator/src/org/jetbrains/kotlin/js/translate/expression/FunctionTranslator.java

@@ -116,7 +116,15 @@ public final class FunctionTranslator extends AbstractTranslator {
     }
 
     private void generateFunctionObject() {
-        setParameters(functionObject, translateParameters());
+        List<JsParameter> parameters = translateParameters();
+
+        VariableDescriptor continuationDescriptor = functionBodyContext.getContinuationParameterDescriptor();
+        if (continuationDescriptor != null) {
+            JsParameter jsParameter = new JsParameter(functionBodyContext.getNameForDescriptor(continuationDescriptor));
+            parameters.add(jsParameter);
+        }
+
+        setParameters(functionObject, parameters);
         translateBody();
     }
 
@@ -154,6 +162,7 @@ public final class FunctionTranslator extends AbstractTranslator {
 
         mayBeAddThisParameterForExtensionFunction(jsParameters);
         addParameters(jsParameters, descriptor, context());
+
         return jsParameters;
     }
 

js/js.translator/src/org/jetbrains/kotlin/js/translate/expression/LiteralFunctionTranslator.kt

@@ -51,7 +51,15 @@ class LiteralFunctionTranslator(context: TranslationContext) : AbstractTranslato
         val descriptor = getFunctionDescriptor(invokingContext.bindingContext(), declaration)
 
         val lambda = invokingContext.getFunctionObject(descriptor)
-        val functionContext = invokingContext.newFunctionBodyWithUsageTracker(lambda, descriptor)
+
+        val aliases = mutableMapOf<DeclarationDescriptor, JsExpression>()
+        if (descriptor.isCoroutineLambda) {
+            aliases.put(descriptor, JsLiteral.THIS)
+        }
+
+        val functionContext = invokingContext
+                .newFunctionBodyWithUsageTracker(lambda, descriptor)
+                .innerContextWithDescriptorsAliased(aliases)
         FunctionTranslator.addParameters(lambda.parameters, descriptor, functionContext)
 
         descriptor.valueParameters.forEach {

js/js.translator/src/org/jetbrains/kotlin/js/translate/reference/CallArgumentTranslator.kt

@@ -39,7 +39,6 @@ import org.jetbrains.kotlin.psi.KtLambdaExpression
 import org.jetbrains.kotlin.psi.KtPsiUtil
 import org.jetbrains.kotlin.psi.ValueArgument
 import org.jetbrains.kotlin.resolve.calls.model.*
-import org.jetbrains.kotlin.resolve.calls.resolvedCallUtil.getImplicitReceiverValue
 import org.jetbrains.kotlin.types.KotlinType
 import java.util.*
 
@@ -157,11 +156,8 @@ class CallArgumentTranslator private constructor(
         }
 
         val callableDescriptor = resolvedCall.resultingDescriptor
-        if (callableDescriptor is FunctionDescriptor && callableDescriptor.isSuspend &&
-            callableDescriptor.initialSignatureDescriptor != null
-        ) {
-            val coroutineDescriptor = resolvedCall.getImplicitReceiverValue()!!.declarationDescriptor
-            result.add(context().getAliasForDescriptor(coroutineDescriptor) ?: JsLiteral.THIS)
+        if (callableDescriptor is FunctionDescriptor && callableDescriptor.isSuspend) {
+            result.add(TranslationUtils.translateContinuationArgument(context(), resolvedCall))
         }
 
         removeLastUndefinedArguments(result)
@@ -331,9 +327,7 @@ class CallArgumentTranslator private constructor(
 
 }
 
-fun Map<TypeParameterDescriptor, KotlinType>.buildReifiedTypeArgs(
-        context: TranslationContext
-): List<JsExpression> {
+fun Map<TypeParameterDescriptor, KotlinType>.buildReifiedTypeArgs(context: TranslationContext): List<JsExpression> {
 
     val reifiedTypeArguments = SmartList<JsExpression>()
     val patternTranslator = PatternTranslator.newInstance(context)

js/js.translator/src/org/jetbrains/kotlin/js/translate/utils/TranslationUtils.java

@@ -31,12 +31,14 @@ import org.jetbrains.kotlin.js.translate.context.TemporaryConstVariable;
 import org.jetbrains.kotlin.js.translate.context.TranslationContext;
 import org.jetbrains.kotlin.js.translate.expression.InlineMetadata;
 import org.jetbrains.kotlin.js.translate.general.Translation;
+import org.jetbrains.kotlin.js.translate.reference.ReferenceTranslator;
 import org.jetbrains.kotlin.psi.*;
 import org.jetbrains.kotlin.psi.psiUtil.PsiUtilsKt;
 import org.jetbrains.kotlin.resolve.BindingContext;
 import org.jetbrains.kotlin.resolve.DescriptorUtils;
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall;
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedValueArgument;
+import org.jetbrains.kotlin.resolve.coroutine.CoroutineReceiverValue;
 import org.jetbrains.kotlin.resolve.scopes.receivers.ImplicitReceiver;
 import org.jetbrains.kotlin.resolve.inline.InlineUtil;
 import org.jetbrains.kotlin.types.KotlinType;
@@ -352,4 +354,31 @@ public final class TranslationUtils {
 
         return handleResultInvocation;
     }
+
+    @NotNull
+    public static JsExpression translateContinuationArgument(@NotNull TranslationContext context, @NotNull ResolvedCall<?> resolvedCall) {
+        CoroutineReceiverValue coroutineReceiver = context.bindingContext().get(
+                BindingContext.COROUTINE_RECEIVER_FOR_SUSPENSION_POINT,
+                resolvedCall.getCall());
+
+        CallableDescriptor continuationDescriptor = null;
+        if (coroutineReceiver != null) {
+            continuationDescriptor = coroutineReceiver.getDeclarationDescriptor();
+        }
+        if (continuationDescriptor == null) {
+            continuationDescriptor = getEnclosingContinuationParameter(context);
+        }
+
+        return ReferenceTranslator.translateAsValueReference(continuationDescriptor, context);
+    }
+
+    @NotNull
+    public static VariableDescriptor getEnclosingContinuationParameter(@NotNull TranslationContext context) {
+        VariableDescriptor result = context.getContinuationParameterDescriptor();
+        if (result == null) {
+            assert context.getParent() != null;
+            result = getEnclosingContinuationParameter(context.getParent());
+        }
+        return result;
+    }
 }