JS: allow to call non-suspend controller functions from coroutine lambda

Author: Alexey Andreev

js/js.dart-ast/src/com/google/dart/compiler/backend/js/ast/metadata/metadataProperties.kt

@@ -66,6 +66,8 @@ var JsInvocation.isSuspend: Boolean by MetadataProperty(default = false)
 
 var JsInvocation.isHandleResult: Boolean by MetadataProperty(default = false)
 
+var JsNameRef.coroutineController by MetadataProperty(default = false)
+
 enum class TypeCheck {
     TYPEOF,
     INSTANCEOF,

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

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

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

@@ -17,7 +17,7 @@
 package org.jetbrains.kotlin.js.coroutine
 
 import com.google.dart.compiler.backend.js.ast.*
-import com.google.dart.compiler.backend.js.ast.metadata.isHandleResult
+import com.google.dart.compiler.backend.js.ast.metadata.coroutineController
 import com.google.dart.compiler.backend.js.ast.metadata.isSuspend
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
@@ -69,8 +69,7 @@ class CoroutineFunctionTransformer(
         val coroutineBlocks = bodyTransformer.postProcess()
 
         for (coroutineBlock in coroutineBlocks) {
-            coroutineBlock.jsBlock.replaceLocalVariables(function.scope)
-            coroutineBlock.jsBlock.replaceHandleResult(bodyTransformer)
+            coroutineBlock.jsBlock.replaceLocalVariables(function.scope, bodyTransformer)
         }
 
         val additionalStatements = mutableListOf<JsStatement>()
@@ -229,19 +228,7 @@ class CoroutineFunctionTransformer(
         return listOf(throwExceptionIfNeeded, loop)
     }
 
-    private fun JsBlock.replaceHandleResult(transformer: CoroutineBodyTransformer) {
-        accept(object : RecursiveJsVisitor() {
-            override fun visitInvocation(invocation: JsInvocation) {
-                super.visitInvocation(invocation)
-                if (invocation.isHandleResult) {
-                    val methodRef = invocation.qualifier as JsNameRef
-                    methodRef.qualifier = JsNameRef(transformer.controllerFieldName, methodRef.qualifier)
-                }
-            }
-        })
-    }
-
-    private fun JsBlock.replaceLocalVariables(scope: JsScope) {
+    private fun JsBlock.replaceLocalVariables(scope: JsScope, transformer: CoroutineBodyTransformer) {
         accept(object : RecursiveJsVisitor() {
             override fun visit(x: JsVars.JsVar) {
                 super.visit(x)
@@ -251,6 +238,9 @@ class CoroutineFunctionTransformer(
 
         val visitor = object : JsVisitorWithContextImpl() {
             override fun endVisit(x: JsNameRef, ctx: JsContext<in JsNode>) {
+                if (x.coroutineController) {
+                    ctx.replaceMe(JsNameRef(transformer.controllerFieldName, x.qualifier))
+                }
                 if (x.qualifier == null && x.name in localVariables) {
                     val fieldName = scope.getFieldName(x.name!!)
                     ctx.replaceMe(JsNameRef(fieldName, JsLiteral.THIS))

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

@@ -17,10 +17,12 @@
 package org.jetbrains.kotlin.js.translate.context;
 
 import com.google.dart.compiler.backend.js.ast.*;
+import com.google.dart.compiler.backend.js.ast.metadata.MetadataProperties;
 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.impl.AnonymousFunctionDescriptor;
 import org.jetbrains.kotlin.descriptors.impl.TypeAliasConstructorDescriptor;
 import org.jetbrains.kotlin.js.config.JsConfig;
 import org.jetbrains.kotlin.js.translate.intrinsic.Intrinsics;
@@ -356,6 +358,16 @@ 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 (DescriptorUtils.isObject(descriptor.getContainingDeclaration())) {
             if (isConstructorOrDirectScope(descriptor.getContainingDeclaration())) {
                 return JsLiteral.THIS;

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

@@ -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.
@@ -23,6 +23,7 @@
 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;
@@ -36,6 +37,7 @@
 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.scopes.receivers.ImplicitReceiver;
 import org.jetbrains.kotlin.types.KotlinType;
 
 import java.util.ArrayList;
@@ -325,7 +327,15 @@ else if (KotlinBuiltIns.isUnit(returnType)) {
         aliases.put(continuationArgument.getArgumentExpression(), JsLiteral.THIS);
 
         TranslationContext returnContext = context.innerContextWithAliasesForExpressions(aliases);
-        JsInvocation handleResultInvocation = (JsInvocation) CallTranslator.translate(returnContext, returnCall, JsLiteral.THIS);
+
+        ImplicitReceiver receiver = (ImplicitReceiver) returnCall.getDispatchReceiver();
+        assert receiver != null;
+        AnonymousFunctionDescriptor lambdaDescriptor = (AnonymousFunctionDescriptor) receiver.getDeclarationDescriptor();
+        ReceiverParameterDescriptor receiverParameter = lambdaDescriptor.getExtensionReceiverParameter();
+        assert receiverParameter != null;
+        JsExpression jsReceiver = returnContext.getDispatchReceiver(receiverParameter);
+
+        JsInvocation handleResultInvocation = (JsInvocation) CallTranslator.translate(returnContext, returnCall, jsReceiver);
         MetadataProperties.setHandleResult(handleResultInvocation, true);
 
         return handleResultInvocation;