JS: refactor how last argument to suspend call is substituted

Author: Alexey Andreev

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

@@ -1,5 +1,6 @@
 // WITH_RUNTIME
 // WITH_REFLECT
+// TARGET_BACKEND: JVM
 class Controller {
     fun withValue(v: String, x: Continuation<String>) {
         x.resume(v)

js/js.frontend/src/org/jetbrains/kotlin/js/descriptorUtils/descriptorUtils.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2015 JetBrains s.r.o.
+ * 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.
@@ -19,7 +19,9 @@ package org.jetbrains.kotlin.js.descriptorUtils
 import com.intellij.openapi.util.text.StringUtil
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.TypeParameterDescriptor
+import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.types.KotlinType
@@ -50,3 +52,6 @@ fun KotlinType.getJetTypeFqName(printTypeArguments: Boolean): String {
 }
 
 fun ClassDescriptor.hasPrimaryConstructor(): Boolean = unsubstitutedPrimaryConstructor != null
+
+val DeclarationDescriptor.isCoroutineLambda: Boolean
+    get() = this is AnonymousFunctionDescriptor && isCoroutine

js/js.inliner/src/org/jetbrains/kotlin/js/coroutine/CoroutineBodyTransformer.kt

@@ -507,7 +507,6 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
     }
 
     private fun handleSuspend(invocation: JsInvocation): JsExpression {
-        invocation.arguments += JsLiteral.THIS
         val nextBlock = CoroutineBlock()
         currentStatements += state(nextBlock)
         currentStatements += JsReturn(invocation)

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

@@ -17,7 +17,6 @@
 package org.jetbrains.kotlin.js.translate.callTranslator
 
 import com.google.dart.compiler.backend.js.ast.JsExpression
-import com.google.dart.compiler.backend.js.ast.JsFunction
 import com.google.dart.compiler.backend.js.ast.JsInvocation
 import com.google.dart.compiler.backend.js.ast.metadata.isSuspend
 import org.jetbrains.kotlin.descriptors.CallableDescriptor

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

@@ -22,7 +22,6 @@
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.descriptors.*;
-import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor;
 import org.jetbrains.kotlin.descriptors.impl.TypeAliasConstructorDescriptor;
 import org.jetbrains.kotlin.js.config.JsConfig;
 import org.jetbrains.kotlin.js.translate.intrinsic.Intrinsics;
@@ -37,6 +36,7 @@
 
 import java.util.*;
 
+import static org.jetbrains.kotlin.js.descriptorUtils.DescriptorUtilsKt.isCoroutineLambda;
 import static org.jetbrains.kotlin.js.translate.context.UsageTrackerKt.getNameForCapturedDescriptor;
 import static org.jetbrains.kotlin.js.translate.utils.BindingUtils.getDescriptorForElement;
 
@@ -358,14 +358,10 @@ public JsExpression getDispatchReceiver(@NotNull ReceiverParameterDescriptor des
         if (alias != null) {
             return alias;
         }
-        if (descriptor.getContainingDeclaration() instanceof AnonymousFunctionDescriptor) {
-            AnonymousFunctionDescriptor function = (AnonymousFunctionDescriptor) descriptor.getContainingDeclaration();
-            if (function.isCoroutine()) {
-                assert function.getExtensionReceiverParameter() != null;
-                JsNameRef result = new JsNameRef("$$controller$$", JsLiteral.THIS);
-                MetadataProperties.setCoroutineController(result, true);
-                return result;
-            }
+        if (isCoroutineLambda(descriptor.getContainingDeclaration())) {
+            JsNameRef result = new JsNameRef("$$controller$$", JsLiteral.THIS);
+            MetadataProperties.setCoroutineController(result, true);
+            return result;
         }
 
         if (DescriptorUtils.isObject(descriptor.getContainingDeclaration())) {
@@ -513,6 +509,9 @@ public JsExpression getArgumentForClosureConstructor(@NotNull DeclarationDescrip
         if (descriptor instanceof ReceiverParameterDescriptor) {
             return getDispatchReceiver((ReceiverParameterDescriptor) descriptor);
         }
+        if (isCoroutineLambda(descriptor)) {
+            return JsLiteral.THIS;
+        }
         return getNameForDescriptor(descriptor).makeRef();
     }
 

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

@@ -16,9 +16,10 @@
 
 package org.jetbrains.kotlin.js.translate.context
 
-import org.jetbrains.kotlin.descriptors.*
 import com.google.dart.compiler.backend.js.ast.JsName
 import com.google.dart.compiler.backend.js.ast.JsScope
+import org.jetbrains.kotlin.descriptors.*
+import org.jetbrains.kotlin.js.descriptorUtils.isCoroutineLambda
 import org.jetbrains.kotlin.js.naming.NameSuggestion
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.DescriptorUtils.*
@@ -48,8 +49,13 @@ class UsageTracker(
 
         // local named function
         if (descriptor is FunctionDescriptor && descriptor.visibility == Visibilities.LOCAL) {
-            assert(!descriptor.getName().isSpecial) { "Function with special name can not be captured, descriptor: $descriptor" }
-            captureIfNeed(descriptor)
+            if (descriptor.isCoroutineLambda) {
+                captureIfNeed(descriptor)
+            }
+            else {
+                assert(!descriptor.getName().isSpecial) { "Function with special name can not be captured, descriptor: $descriptor" }
+                captureIfNeed(descriptor)
+            }
         }
         // local variable
         else if (descriptor is VariableDescriptor && descriptor !is PropertyDescriptor) {
@@ -65,14 +71,15 @@ class UsageTracker(
     }
 
     private fun captureIfNeed(descriptor: DeclarationDescriptor?) {
-
         if (descriptor == null || isCaptured(descriptor) || !isInLocalDeclaration() ||
             isAncestor(containingDescriptor, descriptor, /* strict = */ true) ||
             isReceiverAncestor(descriptor) || isSingletonReceiver(descriptor)
         ) {
             return
         }
 
+        if (descriptor.isCoroutineLambda && descriptor == containingDescriptor) return
+
         parent?.captureIfNeed(descriptor)
 
         captured[descriptor] = descriptor.getJsNameForCapturedDescriptor()
@@ -164,9 +171,9 @@ class UsageTracker(
     }
 
     private fun DeclarationDescriptor.getJsNameForCapturedDescriptor(): JsName {
-        val suggestedName = when (this) {
-            is ReceiverParameterDescriptor -> getNameForCapturedReceiver()
-            is TypeParameterDescriptor -> Namer.isInstanceSuggestedName(this)
+        val suggestedName = when {
+            this is ReceiverParameterDescriptor -> getNameForCapturedReceiver()
+            this 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/LiteralFunctionTranslator.kt

@@ -19,7 +19,9 @@ package org.jetbrains.kotlin.js.translate.expression
 import com.google.dart.compiler.backend.js.ast.*
 import com.google.dart.compiler.backend.js.ast.metadata.*
 import org.jetbrains.kotlin.descriptors.*
+import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor
 import org.jetbrains.kotlin.descriptors.impl.ValueParameterDescriptorImpl
+import org.jetbrains.kotlin.js.descriptorUtils.isCoroutineLambda
 import org.jetbrains.kotlin.js.inline.util.getInnerFunction
 import org.jetbrains.kotlin.js.translate.context.TranslationContext
 import org.jetbrains.kotlin.js.translate.context.getNameForCapturedDescriptor
@@ -75,11 +77,12 @@ class LiteralFunctionTranslator(context: TranslationContext) : AbstractTranslato
             lambdaCreator.isLocal = true
             lambdaCreator.coroutineType = continuationType
             lambdaCreator.controllerType = controllerType
-            if (!isRecursive) {
+            if (!isRecursive || descriptor.isCoroutineLambda) {
                 lambda.name = null
             }
             lambdaCreator.name.staticRef = lambdaCreator
-            return lambdaCreator.withCapturedParameters(descriptor, functionContext, invokingContext)
+            return lambdaCreator.withCapturedParameters(descriptor, descriptor.wrapContextForCoroutineIfNecessary(functionContext),
+                                                        invokingContext)
         }
 
         lambda.isLocal = true
@@ -101,6 +104,10 @@ class LiteralFunctionTranslator(context: TranslationContext) : AbstractTranslato
     }
 }
 
+private fun CallableMemberDescriptor.wrapContextForCoroutineIfNecessary(context: TranslationContext): TranslationContext {
+    return if (isCoroutineLambda) context.innerContextWithDescriptorsAliased(mapOf(this to JsLiteral.THIS)) else context
+}
+
 fun JsFunction.withCapturedParameters(
         descriptor: CallableMemberDescriptor,
         context: TranslationContext,
@@ -116,7 +123,7 @@ fun JsFunction.withCapturedParameters(
     val tracker = context.usageTracker()!!
 
     for ((capturedDescriptor, name) in tracker.capturedDescriptorToJsName) {
-        if (capturedDescriptor == tracker.containingDescriptor) continue
+        if (capturedDescriptor == tracker.containingDescriptor && !capturedDescriptor.isCoroutineLambda) continue
 
         val capturedRef = invokingContext.getArgumentForClosureConstructor(capturedDescriptor)
         var additionalArgs = listOf(capturedRef)

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

@@ -39,6 +39,7 @@ 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.*
 
@@ -75,9 +76,6 @@ class CallArgumentTranslator private constructor(
 
     private fun translate(): ArgumentsInfo {
         val valueParameters = resolvedCall.resultingDescriptor.valueParameters
-        if (valueParameters.isEmpty()) {
-            return ArgumentsInfo(listOf<JsExpression>(), false, null)
-        }
         var hasSpreadOperator = false
         var cachedReceiver: TemporaryConstVariable? = null
 
@@ -159,10 +157,15 @@ class CallArgumentTranslator private constructor(
         }
 
         val callableDescriptor = resolvedCall.resultingDescriptor
-        if (callableDescriptor !is FunctionDescriptor || !callableDescriptor.isSuspend) {
-            removeLastUndefinedArguments(result)
+        if (callableDescriptor is FunctionDescriptor && callableDescriptor.isSuspend &&
+            callableDescriptor.initialSignatureDescriptor != null
+        ) {
+            val coroutineDescriptor = resolvedCall.getImplicitReceiverValue()!!.declarationDescriptor
+            result.add(context().getAliasForDescriptor(coroutineDescriptor) ?: JsLiteral.THIS)
         }
 
+        removeLastUndefinedArguments(result)
+
         return ArgumentsInfo(result, hasSpreadOperator, cachedReceiver)
     }
 

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

@@ -23,7 +23,6 @@
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns;
 import org.jetbrains.kotlin.descriptors.*;
-import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor;
 import org.jetbrains.kotlin.descriptors.impl.LocalVariableAccessorDescriptor;
 import org.jetbrains.kotlin.descriptors.impl.LocalVariableDescriptor;
 import org.jetbrains.kotlin.js.translate.callTranslator.CallTranslator;
@@ -330,7 +329,7 @@ else if (KotlinBuiltIns.isUnit(returnType)) {
 
         ImplicitReceiver receiver = (ImplicitReceiver) returnCall.getDispatchReceiver();
         assert receiver != null;
-        AnonymousFunctionDescriptor lambdaDescriptor = (AnonymousFunctionDescriptor) receiver.getDeclarationDescriptor();
+        FunctionDescriptor lambdaDescriptor = (FunctionDescriptor) receiver.getDeclarationDescriptor();
         ReceiverParameterDescriptor receiverParameter = lambdaDescriptor.getExtensionReceiverParameter();
         assert receiverParameter != null;
         JsExpression jsReceiver = returnContext.getDispatchReceiver(receiverParameter);