JS: allow to call suspend function from any position within another suspend function

js/js.ast/src/org/jetbrains/kotlin/js/backend/ast/metadata/metadataProperties.kt

@@ -24,7 +24,6 @@ import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.js.backend.ast.*
 import org.jetbrains.kotlin.resolve.inline.InlineStrategy
-import org.jetbrains.kotlin.types.KotlinType
 
 var JsName.staticRef: JsNode? by MetadataProperty(default = null)
 
@@ -68,8 +67,6 @@ var HasMetadata.synthetic: Boolean by MetadataProperty(default = false)
 
 var HasMetadata.sideEffects: SideEffectKind by MetadataProperty(default = SideEffectKind.AFFECTS_STATE)
 
-var JsFunction.coroutineType: KotlinType? by MetadataProperty(default = null)
-
 /**
  * Denotes a suspension call-site that is to be processed by coroutine transformer.
  * More clearly, denotes invocation that should immediately return from coroutine state machine
@@ -99,6 +96,12 @@ var JsNameRef.coroutineResult by MetadataProperty(default = false)
  */
 var JsNameRef.coroutineController by MetadataProperty(default = false)
 
+/**
+ * Denotes a reference to coroutine's receiver. It's later rewritten to `this`. Required to distinguish between `this` for
+ * function's dispatch receiver and coroutine's state receiver.
+ */
+var JsNameRef.coroutineReceiver by MetadataProperty(default = false)
+
 var JsName.imported by MetadataProperty(default = false)
 
 var JsFunction.coroutineMetadata: CoroutineMetadata? by MetadataProperty(default = null)
@@ -114,7 +117,9 @@ class CoroutineMetadata(
         val facadeName: JsName,
         val baseClassRef: JsExpression,
         val suspendObjectRef: JsExpression,
-        val hasController: Boolean
+        val isLambda: Boolean,
+        val hasController: Boolean,
+        val hasReceiver: Boolean
 )
 
 enum class TypeCheck {

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

@@ -249,7 +249,8 @@ class CoroutineBodyTransformer(private val program: JsProgram, private val conte
         }
 
         if (catchNode != null) {
-            currentStatements += JsAstUtils.newVar(catchNode.parameter.name, JsNameRef(context.metadata.exceptionName, JsLiteral.THIS))
+            currentStatements += JsAstUtils.newVar(catchNode.parameter.name, JsNameRef(
+                    context.metadata.exceptionName, JsAstUtils.stateMachineReceiver()))
             catchNode.body.statements.forEach { it.accept(this) }
 
             if (finallyNode == null) {
@@ -278,8 +279,8 @@ class CoroutineBodyTransformer(private val program: JsProgram, private val conte
     // for simple `when` statement, we will need to support JsSwitch here
 
     private fun generateFinallyExit() {
-        val finallyPathRef = JsNameRef(context.metadata.finallyPathName, JsLiteral.THIS)
+        val finallyPathRef = JsNameRef(context.metadata.finallyPathName, JsAstUtils.stateMachineReceiver())
-        val stateRef = JsNameRef(context.metadata.stateName, JsLiteral.THIS)
+        val stateRef = JsNameRef(context.metadata.stateName, JsAstUtils.stateMachineReceiver())
         val nextState = JsInvocation(JsNameRef("shift", finallyPathRef))
         currentStatements += JsAstUtils.assignment(stateRef, nextState).makeStmt()
         currentStatements += jump()
@@ -307,7 +308,7 @@ class CoroutineBodyTransformer(private val program: JsProgram, private val conte
             jumpWithFinally(0, returnBlock)
             val returnExpression = x.expression
             val returnFieldRef = if (returnExpression != null) {
-                val ref = JsNameRef(context.returnValueFieldName, JsLiteral.THIS)
+                val ref = JsNameRef(context.returnValueFieldName, JsAstUtils.stateMachineReceiver())
                 currentStatements += JsAstUtils.assignment(ref, x.expression).makeStmt()
                 ref
             }
@@ -359,7 +360,9 @@ class CoroutineBodyTransformer(private val program: JsProgram, private val conte
             listOf(JsReturn(context.metadata.suspendObjectRef.deepCopy()))
         }
         else {
-            val resultRef = JsNameRef(context.metadata.resultName, JsLiteral.THIS).apply { sideEffects = SideEffectKind.DEPENDS_ON_STATE }
+            val resultRef = JsNameRef(context.metadata.resultName, JsAstUtils.stateMachineReceiver()).apply {
+                sideEffects = SideEffectKind.DEPENDS_ON_STATE
+            }
             val invocationStatement = JsAstUtils.assignment(resultRef, invokeExpression).makeStmt()
             val suspendCondition = JsAstUtils.equality(resultRef.deepCopy(), context.metadata.suspendObjectRef.deepCopy())
             val suspendIfNeeded = JsIf(suspendCondition, JsReturn(context.metadata.suspendObjectRef.deepCopy()))

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

@@ -24,12 +24,12 @@ import org.jetbrains.kotlin.js.inline.util.getInnerFunction
 import org.jetbrains.kotlin.js.translate.context.Namer
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
 
-class CoroutineFunctionTransformer(private val program: JsProgram, private val function: JsFunction) {
+class CoroutineFunctionTransformer(private val program: JsProgram, private val function: JsFunction, name: String?) {
     private val innerFunction = function.getInnerFunction()
     private val functionWithBody = innerFunction ?: function
     private val body = functionWithBody.body
     private val localVariables = (function.collectLocalVariables() + functionWithBody.collectLocalVariables()).toMutableSet()
-    private val className = function.scope.parent.declareFreshName("Coroutine\$${function.name}")
+    private val className = function.scope.parent.declareFreshName("Coroutine\$${name ?: "anonymous"}")
 
     fun transform(): List<JsStatement> {
         val context = CoroutineTransformationContext(function.scope, function)
@@ -58,11 +58,16 @@ class CoroutineFunctionTransformer(private val program: JsProgram, private val f
     ) {
         val constructor = JsFunction(function.scope.parent, JsBlock(), "Continuation")
         constructor.name = className
+        if (context.metadata.hasReceiver) {
+            constructor.parameters += JsParameter(context.receiverFieldName)
+        }
         constructor.parameters += function.parameters.map { JsParameter(it.name) }
         if (innerFunction != null) {
             constructor.parameters += innerFunction.parameters.map { JsParameter(it.name) }
         }
 
+        val lastParameter = function.parameters.lastOrNull()?.name
+
         val controllerName = if (context.metadata.hasController) {
             function.scope.declareFreshName("controller").apply { constructor.parameters += JsParameter(this) }
         }
@@ -70,18 +75,27 @@ class CoroutineFunctionTransformer(private val program: JsProgram, private val f
             null
         }
 
-        val interceptorName = function.scope.declareFreshName("interceptor")
+        val interceptorRef = if (context.metadata.isLambda) {
-        constructor.parameters += JsParameter(interceptorName)
+            val interceptorName = function.scope.declareFreshName("interceptor")
+            constructor.parameters += JsParameter(interceptorName)
+            interceptorName.makeRef()
+        }
+        else {
+            lastParameter!!.makeRef()
+        }
 
         val parameterNames = (function.parameters.map { it.name } + innerFunction?.parameters?.map { it.name }.orEmpty()).toSet()
 
         constructor.body.statements.run {
             val baseClass = context.metadata.baseClassRef.deepCopy()
-            this += JsInvocation(Namer.getFunctionCallRef(baseClass), JsLiteral.THIS, interceptorName.makeRef()).makeStmt()
+            this += JsInvocation(Namer.getFunctionCallRef(baseClass), JsLiteral.THIS, interceptorRef).makeStmt()
             if (controllerName != null) {
                 assignToField(context.controllerFieldName, controllerName.makeRef())
             }
             assignToField(context.metadata.exceptionStateName, program.getNumberLiteral(globalCatchBlockIndex))
+            if (context.metadata.hasReceiver) {
+                assignToField(context.receiverFieldName, context.receiverFieldName.makeRef())
+            }
             for (localVariable in localVariables) {
                 val value = if (localVariable !in parameterNames) Namer.getUndefinedExpression() else localVariable.makeRef()
                 assignToField(function.scope.getFieldName(localVariable), value)
@@ -139,6 +153,9 @@ class CoroutineFunctionTransformer(private val program: JsProgram, private val f
 
     private fun generateCoroutineInstantiation(context: CoroutineTransformationContext) {
         val instantiation = JsNew(className.makeRef())
+        if (context.metadata.hasReceiver) {
+            instantiation.arguments += JsLiteral.THIS
+        }
         instantiation.arguments += function.parameters.map { it.name.makeRef() }
         if (innerFunction != null) {
             instantiation.arguments += innerFunction.parameters.map { it.name.makeRef() }
@@ -148,14 +165,15 @@ class CoroutineFunctionTransformer(private val program: JsProgram, private val f
             instantiation.arguments += JsLiteral.THIS
         }
 
-        val interceptorParamName = functionWithBody.scope.declareFreshName("interceptor")
+        if (context.metadata.isLambda) {
-        functionWithBody.parameters += JsParameter(interceptorParamName)
+            val interceptorParamName = functionWithBody.scope.declareFreshName("interceptor")
+            functionWithBody.parameters += JsParameter(interceptorParamName)
+            instantiation.arguments += interceptorParamName.makeRef()
+        }
 
         val suspendedName = functionWithBody.scope.declareFreshName("suspended")
         functionWithBody.parameters += JsParameter(suspendedName)
 
-        instantiation.arguments += interceptorParamName.makeRef()
-
         val instanceName = functionWithBody.scope.declareFreshName("instance")
         functionWithBody.body.statements += JsAstUtils.newVar(instanceName, JsNameRef(context.metadata.facadeName, instantiation))
 

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

@@ -112,7 +112,7 @@ fun List<CoroutineBlock>.replaceCoroutineFlowStatements(context: CoroutineTransf
         override fun endVisit(x: JsDebugger, ctx: JsContext<in JsStatement>) {
             val target = x.targetBlock
             if (target != null) {
-                val lhs = JsNameRef(context.metadata.stateName, JsLiteral.THIS)
+                val lhs = JsNameRef(context.metadata.stateName, JsAstUtils.stateMachineReceiver())
                 val rhs = program.getNumberLiteral(blockIndexes[target]!!)
                 ctx.replaceMe(JsExpressionStatement(JsAstUtils.assignment(lhs, rhs)).apply {
                     targetBlock = true
@@ -121,7 +121,7 @@ fun List<CoroutineBlock>.replaceCoroutineFlowStatements(context: CoroutineTransf
 
             val exceptionTarget = x.targetExceptionBlock
             if (exceptionTarget != null) {
-                val lhs = JsNameRef(context.metadata.exceptionStateName, JsLiteral.THIS)
+                val lhs = JsNameRef(context.metadata.exceptionStateName, JsAstUtils.stateMachineReceiver())
                 val rhs = program.getNumberLiteral(blockIndexes[exceptionTarget]!!)
                 ctx.replaceMe(JsExpressionStatement(JsAstUtils.assignment(lhs, rhs)).apply {
                     targetExceptionBlock = true
@@ -131,7 +131,7 @@ fun List<CoroutineBlock>.replaceCoroutineFlowStatements(context: CoroutineTransf
             val finallyPath = x.finallyPath
             if (finallyPath != null) {
                 if (finallyPath.isNotEmpty()) {
-                    val lhs = JsNameRef(context.metadata.finallyPathName, JsLiteral.THIS)
+                    val lhs = JsNameRef(context.metadata.finallyPathName, JsAstUtils.stateMachineReceiver())
                     val rhs = JsArrayLiteral(finallyPath.map { program.getNumberLiteral(blockIndexes[it]!!) })
                     ctx.replaceMe(JsExpressionStatement(JsAstUtils.assignment(lhs, rhs)).apply {
                         this.finallyPath = true
@@ -200,8 +200,16 @@ private fun CoroutineBlock.collectFinallyPaths(): List<List<CoroutineBlock>> {
 
 fun JsBlock.replaceSpecialReferences(context: CoroutineTransformationContext) {
     val visitor = object : JsVisitorWithContextImpl() {
+        override fun endVisit(x: JsLiteral.JsThisRef, ctx: JsContext<in JsNode>) {
+            ctx.replaceMe(JsNameRef(context.receiverFieldName, JsLiteral.THIS))
+        }
+
         override fun endVisit(x: JsNameRef, ctx: JsContext<in JsNode>) {
             when {
+                x.coroutineReceiver -> {
+                    ctx.replaceMe(JsLiteral.THIS)
+                }
+
                 x.coroutineController -> {
                     ctx.replaceMe(JsNameRef(context.controllerFieldName, x.qualifier).apply {
                         sideEffects = SideEffectKind.PURE

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

@@ -26,4 +26,5 @@ class CoroutineTransformationContext(private val scope: JsScope, function: JsFun
     val metadata = function.coroutineMetadata!!
     val controllerFieldName by lazy { scope.declareFreshName("\$controller") }
     val returnValueFieldName by lazy { scope.declareFreshName("\$returnValue") }
+    val receiverFieldName by lazy { scope.declareFreshName("\$this") }
 }

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

@@ -17,21 +17,56 @@
 package org.jetbrains.kotlin.js.coroutine
 
 import org.jetbrains.kotlin.js.backend.ast.*
-import org.jetbrains.kotlin.js.backend.ast.metadata.coroutineType
+import org.jetbrains.kotlin.js.backend.ast.metadata.coroutineMetadata
+import org.jetbrains.kotlin.js.translate.expression.InlineMetadata
+import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
 
 class CoroutineTransformer(private val program: JsProgram) : JsVisitorWithContextImpl() {
-    private val additionalStatements = mutableListOf<JsStatement>()
+    private val additionalStatementsByNode = mutableMapOf<JsNode, List<JsStatement>>()
 
     override fun endVisit(x: JsExpressionStatement, ctx: JsContext<in JsStatement>) {
-        additionalStatements.forEach { ctx.addNext(it) }
+        additionalStatementsByNode.remove(x)?.forEach { ctx.addNext(it) }
-        additionalStatements.clear()
         super.endVisit(x, ctx)
     }
 
-    override fun endVisit(x: JsFunction, ctx: JsContext<in JsStatement>) {
+    override fun endVisit(x: JsVars, ctx: JsContext<in JsStatement>) {
-        val coroutineType = x.coroutineType
+        for (v in x.vars) {
-        if (coroutineType != null) {
+            additionalStatementsByNode.remove(v)?.forEach { ctx.addNext(it) }
-            additionalStatements += CoroutineFunctionTransformer(program, x).transform()
         }
+        super.endVisit(x, ctx)
+    }
+
+    override fun visit(x: JsExpressionStatement, ctx: JsContext<*>): Boolean {
+        val expression = x.expression
+        val assignment = JsAstUtils.decomposeAssignment(expression)
+        if (assignment != null) {
+            val (lhs, rhs) = assignment
+            val function = rhs as? JsFunction ?: InlineMetadata.decompose(rhs)?.function
+            if (function?.coroutineMetadata != null) {
+                val name = ((lhs as? JsNameRef)?.name ?: function.name)?.ident
+                additionalStatementsByNode[x] = CoroutineFunctionTransformer(program, function, name).transform()
+                return false
+            }
+        }
+        else if (expression is JsFunction) {
+            if (expression.coroutineMetadata != null) {
+                additionalStatementsByNode[x] = CoroutineFunctionTransformer(program, expression, expression.name?.ident).transform()
+                return false
+            }
+        }
+        return super.visit(x, ctx)
+    }
+
+    override fun visit(x: JsVars.JsVar, ctx: JsContext<*>): Boolean {
+        val initExpression = x.initExpression
+        if (initExpression != null) {
+            val function = initExpression as? JsFunction ?: InlineMetadata.decompose(initExpression)?.function
+            if (function?.coroutineMetadata != null) {
+                val name = x.name.ident
+                additionalStatementsByNode[x] = CoroutineFunctionTransformer(program, function, name).transform()
+                return false
+            }
+        }
+        return super.visit(x, ctx)
     }
 }

js/js.libraries/src/core/coroutines.kt

@@ -126,11 +126,13 @@ internal abstract class CoroutineImpl(private val resultContinuation: Continuati
     }
 
     override fun resume(data: Any?) {
-        this.result = data
+        result = data
         try {
-            val result = doResume()
+            result = doResume()
             if (result != SUSPENDED_MARKER) {
-                resultContinuation.resume(result)
+                val data = result
+                result = SUSPENDED_MARKER
+                resultContinuation.resume(data)
             }
         }
         catch (e: Throwable) {
@@ -142,9 +144,11 @@ internal abstract class CoroutineImpl(private val resultContinuation: Continuati
         state = exceptionState
         this.exception = exception
         try {
-            val result = doResume()
+            result = doResume()
             if (result != SUSPENDED_MARKER) {
-                resultContinuation.resume(result)
+                val data = result
+                result = SUSPENDED_MARKER
+                resultContinuation.resume(data)
             }
         }
         catch (e: Throwable) {

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

@@ -28,10 +28,7 @@ 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.*
-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.calls.callResolverUtil.isInvokeCallOnVariable
@@ -143,7 +140,7 @@ private fun translateFunctionCall(
                               inlineResolvedCall.resultingDescriptor, context)
     }
 
-    if (resolvedCall.resultingDescriptor.isSuspend && !context.isInSuspendFunction) {
+    if (resolvedCall.resultingDescriptor.isSuspend && context.isInStateMachine) {
         context.currentBlock.statements += JsAstUtils.asSyntheticStatement((callExpression as JsInvocation).apply {
             isSuspend = true
             isPreSuspend = true
@@ -157,6 +154,9 @@ private fun translateFunctionCall(
     return callExpression
 }
 
+private val TranslationContext.isInStateMachine
+    get() = (declarationDescriptor as? FunctionDescriptor)?.requiresStateMachineTransformation(this) == true
+
 private fun translateCallWithContinuation(context: TranslationContext, resolvedCall: ResolvedCall<out FunctionDescriptor>): JsExpression {
     val arguments = CallArgumentTranslator.translate(resolvedCall, null, context)
     val coroutineArgument = TranslationUtils.getEnclosingContinuationParameter(context)

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

@@ -21,6 +21,7 @@ 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.AnonymousFunctionDescriptor;
 import org.jetbrains.kotlin.descriptors.impl.LocalVariableDescriptor;
 import org.jetbrains.kotlin.descriptors.impl.TypeAliasConstructorDescriptor;
 import org.jetbrains.kotlin.js.backend.ast.*;
@@ -106,7 +107,7 @@ public class TranslationContext {
         }
         if (declarationDescriptor instanceof FunctionDescriptor) {
             FunctionDescriptor function = (FunctionDescriptor) declarationDescriptor;
-            if (function.isSuspend()) {
+            if (function.isSuspend() && !(function instanceof AnonymousFunctionDescriptor)) {
                 ClassDescriptor continuationDescriptor = getCurrentModule().getBuiltIns().getBuiltInClassByFqName(
                         DescriptorUtils.CONTINUATION_INTERFACE_FQ_NAME);
                 return new LocalVariableDescriptor(
@@ -403,7 +404,7 @@ public class TranslationContext {
             return alias;
         }
         if (isCoroutineLambda(descriptor.getContainingDeclaration())) {
-            JsNameRef result = new JsNameRef("$$controller$$", JsLiteral.THIS);
+            JsNameRef result = new JsNameRef("$$controller$$", JsAstUtils.stateMachineReceiver());
             MetadataProperties.setCoroutineController(result, true);
             return result;
         }
@@ -680,10 +681,6 @@ public class TranslationContext {
         return false;
     }
 
-    public boolean isInSuspendFunction() {
-        return declarationDescriptor instanceof FunctionDescriptor && ((FunctionDescriptor) declarationDescriptor).isSuspend();
-    }
-
     @Nullable
     public VariableDescriptor getContinuationParameterDescriptor() {
         return continuationParameterDescriptor;

js/js.translator/src/org/jetbrains/kotlin/js/translate/declaration/AbstractDeclarationVisitor.kt

@@ -16,16 +16,18 @@
 
 package org.jetbrains.kotlin.js.translate.declaration
 
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.PropertyDescriptor
+import org.jetbrains.kotlin.descriptors.isOverridable
 import org.jetbrains.kotlin.js.backend.ast.JsExpression
-import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.js.translate.context.TranslationContext
 import org.jetbrains.kotlin.js.translate.expression.translateAndAliasParameters
 import org.jetbrains.kotlin.js.translate.expression.translateFunction
 import org.jetbrains.kotlin.js.translate.expression.wrapWithInlineMetadata
 import org.jetbrains.kotlin.js.translate.general.TranslatorVisitor
-import org.jetbrains.kotlin.js.translate.utils.BindingUtils
+import org.jetbrains.kotlin.js.translate.reference.ReferenceTranslator
-import org.jetbrains.kotlin.js.translate.utils.FunctionBodyTranslator
+import org.jetbrains.kotlin.js.translate.utils.*
-import org.jetbrains.kotlin.js.translate.utils.TranslationUtils
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.descriptorUtil.isExtensionProperty
 
@@ -93,7 +95,20 @@ abstract class AbstractDeclarationVisitor : TranslatorVisitor<Unit>()  {
             context: TranslationContext
     ): JsExpression {
         val function = context.getFunctionObject(descriptor)
-        val innerContext = context.newDeclaration(descriptor).translateAndAliasParameters(descriptor, function.parameters)
+        var innerContext = context.newDeclaration(descriptor).translateAndAliasParameters(descriptor, function.parameters)
+
+        if (descriptor.isSuspend) {
+            if (descriptor.requiresStateMachineTransformation(context)) {
+                function.fillCoroutineMetadata(context, descriptor, hasController = false, isLambda = false)
+                innerContext = innerContext.innerContextWithAliased(descriptor, JsAstUtils.stateMachineReceiver())
+            }
+            else {
+                val continuationRef = ReferenceTranslator.translateAsValueReference(
+                        innerContext.continuationParameterDescriptor!!, innerContext)
+                innerContext = innerContext.innerContextWithAliased(descriptor, continuationRef)
+            }
+        }
+
         innerContext.getInnerNameForDescriptor(descriptor)
         if (!descriptor.isOverridable) {
             function.body.statements += FunctionBodyTranslator.setDefaultValueForArguments(descriptor, innerContext)

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

@@ -20,7 +20,6 @@ import org.jetbrains.kotlin.backend.common.SUSPENDED_MARKER_NAME
 import org.jetbrains.kotlin.builtins.isBuiltinExtensionFunctionalType
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.impl.ValueParameterDescriptorImpl
-import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.js.backend.ast.*
 import org.jetbrains.kotlin.js.backend.ast.metadata.*
 import org.jetbrains.kotlin.js.descriptorUtils.isCoroutineLambda
@@ -30,15 +29,12 @@ import org.jetbrains.kotlin.js.translate.context.getNameForCapturedDescriptor
 import org.jetbrains.kotlin.js.translate.context.hasCapturedExceptContaining
 import org.jetbrains.kotlin.js.translate.context.isCaptured
 import org.jetbrains.kotlin.js.translate.general.AbstractTranslator
-import org.jetbrains.kotlin.js.translate.reference.ReferenceTranslator
 import org.jetbrains.kotlin.js.translate.utils.BindingUtils.getFunctionDescriptor
 import org.jetbrains.kotlin.js.translate.utils.FunctionBodyTranslator.setDefaultValueForArguments
 import org.jetbrains.kotlin.js.translate.utils.FunctionBodyTranslator.translateFunctionBody
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
-import org.jetbrains.kotlin.js.translate.utils.TranslationUtils
 import org.jetbrains.kotlin.js.translate.utils.TranslationUtils.simpleReturnFunction
-import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.js.translate.utils.fillCoroutineMetadata
-import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtDeclarationWithBody
 import org.jetbrains.kotlin.psi.KtParameter
 import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
@@ -46,10 +42,6 @@ import org.jetbrains.kotlin.resolve.inline.InlineUtil
 import org.jetbrains.kotlin.types.KotlinType
 
 class LiteralFunctionTranslator(context: TranslationContext) : AbstractTranslator(context) {
-    companion object {
-        private val COROUTINE_INTRINSICS_PACKAGE_FQ_NAME = FqName("kotlin.coroutines.intrinsics")
-    }
-
     fun translate(
             declaration: KtDeclarationWithBody,
             continuationType: KotlinType? = null
@@ -61,7 +53,7 @@ class LiteralFunctionTranslator(context: TranslationContext) : AbstractTranslato
 
         val aliases = mutableMapOf<DeclarationDescriptor, JsExpression>()
         if (descriptor.isCoroutineLambda) {
-            aliases.put(descriptor, JsLiteral.THIS)
+            aliases.put(descriptor, JsAstUtils.stateMachineReceiver())
         }
 
         val functionContext = invokingContext
@@ -98,7 +90,7 @@ class LiteralFunctionTranslator(context: TranslationContext) : AbstractTranslato
                 lambda.name = null
             }
             lambdaCreator.name.staticRef = lambdaCreator
-            lambdaCreator.fillCoroutineMetadata(invokingContext, continuationType)
+            lambdaCreator.fillCoroutineMetadata(invokingContext, descriptor, continuationType)
             return lambdaCreator.withCapturedParameters(descriptor, descriptor.wrapContextForCoroutineIfNecessary(functionContext),
                                                         invokingContext)
         }
@@ -106,36 +98,15 @@ class LiteralFunctionTranslator(context: TranslationContext) : AbstractTranslato
         lambda.isLocal = true
 
         invokingContext.addDeclarationStatement(lambda.makeStmt())
-        lambda.fillCoroutineMetadata(invokingContext, continuationType)
+        lambda.fillCoroutineMetadata(invokingContext, descriptor, continuationType)
         lambda.name.staticRef = lambda
         return getReferenceToLambda(invokingContext, descriptor, lambda.name)
     }
 
-    fun JsFunction.fillCoroutineMetadata(context: TranslationContext, continuationType: KotlinType?) {
+    fun JsFunction.fillCoroutineMetadata(context: TranslationContext, descriptor: FunctionDescriptor, continuationType: KotlinType?) {
         if (continuationType == null) return
 
-        val suspendPropertyDescriptor = context().currentModule.getPackage(COROUTINE_INTRINSICS_PACKAGE_FQ_NAME)
+        fillCoroutineMetadata(context, descriptor, hasController = continuationType.isBuiltinExtensionFunctionalType, isLambda = true)
-                .memberScope.getContributedVariables(SUSPENDED_MARKER_NAME, NoLookupLocation.FROM_BACKEND).first()
-
-        val coroutineBaseClassRef = ReferenceTranslator.translateAsTypeReference(TranslationUtils.getCoroutineBaseClass(context), context)
-
-        fun getCoroutinePropertyName(id: String) =
-                context.getNameForDescriptor(TranslationUtils.getCoroutineProperty(context, id))
-
-        coroutineMetadata = CoroutineMetadata(
-                doResumeName = context.getNameForDescriptor(TranslationUtils.getCoroutineDoResumeFunction(context)),
-                resumeName = context.getNameForDescriptor(TranslationUtils.getCoroutineResumeFunction(context)),
-                suspendObjectRef = ReferenceTranslator.translateAsValueReference(suspendPropertyDescriptor, context()),
-                baseClassRef = coroutineBaseClassRef,
-                stateName = getCoroutinePropertyName("state"),
-                exceptionStateName = getCoroutinePropertyName("exceptionState"),
-                finallyPathName = getCoroutinePropertyName("finallyPath"),
-                resultName = getCoroutinePropertyName("result"),
-                exceptionName = getCoroutinePropertyName("exception"),
-                facadeName = getCoroutinePropertyName("facade"),
-                hasController = continuationType.isBuiltinExtensionFunctionalType
-        )
-        coroutineType = continuationType
     }
 
     fun ValueParameterDescriptorImpl.WithDestructuringDeclaration.translate(context: TranslationContext): JsVars {
@@ -149,7 +120,12 @@ class LiteralFunctionTranslator(context: TranslationContext) : AbstractTranslato
 }
 
 private fun CallableMemberDescriptor.wrapContextForCoroutineIfNecessary(context: TranslationContext): TranslationContext {
-    return if (isCoroutineLambda) context.innerContextWithDescriptorsAliased(mapOf(this to JsLiteral.THIS)) else context
+    return if (isCoroutineLambda) {
+        context.innerContextWithDescriptorsAliased(mapOf(this to JsAstUtils.stateMachineReceiver()))
+    }
+    else {
+        context
+    }
 }
 
 fun JsFunction.withCapturedParameters(

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

@@ -539,4 +539,11 @@ public final class JsAstUtils {
     public static JsExpression prototypeOf(@NotNull JsExpression expression) {
         return pureFqn("prototype", expression);
     }
+
+    @NotNull
+    public static JsExpression stateMachineReceiver() {
+        JsNameRef result = new JsNameRef("$this$");
+        MetadataProperties.setCoroutineReceiver(result, true);
+        return result;
+    }
 }

js/js.translator/src/org/jetbrains/kotlin/js/translate/utils/utils.kt

@@ -17,15 +17,22 @@
 package org.jetbrains.kotlin.js.translate.utils
 
 import com.intellij.util.SmartList
+import org.jetbrains.kotlin.backend.common.SUSPENDED_MARKER_NAME
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.PropertyDescriptor
 import org.jetbrains.kotlin.descriptors.TypeParameterDescriptor
+import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor
+import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.js.backend.ast.*
+import org.jetbrains.kotlin.js.backend.ast.metadata.CoroutineMetadata
+import org.jetbrains.kotlin.js.backend.ast.metadata.coroutineMetadata
 import org.jetbrains.kotlin.js.translate.context.Namer
 import org.jetbrains.kotlin.js.translate.context.TranslationContext
 import org.jetbrains.kotlin.js.translate.reference.ReferenceTranslator
 import org.jetbrains.kotlin.js.translate.utils.TranslationUtils.simpleReturnFunction
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.types.KotlinType
 
@@ -119,4 +126,43 @@ fun TranslationContext.addAccessorsToPrototype(
     val propertyName = getNameForDescriptor(propertyDescriptor)
     val defineProperty = JsAstUtils.defineProperty(prototypeRef, propertyName.ident, literal, program())
     addDeclarationStatement(defineProperty.makeStmt())
-}
+}
+
+fun FunctionDescriptor.requiresStateMachineTransformation(context: TranslationContext): Boolean =
+        this is AnonymousFunctionDescriptor ||
+        context.bindingContext()[BindingContext.CONTAINS_NON_TAIL_SUSPEND_CALLS, this] == true
+
+fun JsFunction.fillCoroutineMetadata(
+        context: TranslationContext,
+        descriptor: FunctionDescriptor,
+        hasController: Boolean, isLambda: Boolean
+) {
+    if (!descriptor.requiresStateMachineTransformation(context)) return
+
+    val suspendPropertyDescriptor = context.currentModule.getPackage(COROUTINES_INTRINSICS_PACKAGE_FQ_NAME)
+            .memberScope
+            .getContributedVariables(SUSPENDED_MARKER_NAME, NoLookupLocation.FROM_BACKEND).first()
+
+    val coroutineBaseClassRef = ReferenceTranslator.translateAsTypeReference(TranslationUtils.getCoroutineBaseClass(context), context)
+
+    fun getCoroutinePropertyName(id: String) =
+            context.getNameForDescriptor(TranslationUtils.getCoroutineProperty(context, id))
+
+    coroutineMetadata = CoroutineMetadata(
+            doResumeName = context.getNameForDescriptor(TranslationUtils.getCoroutineDoResumeFunction(context)),
+            resumeName = context.getNameForDescriptor(TranslationUtils.getCoroutineResumeFunction(context)),
+            suspendObjectRef = ReferenceTranslator.translateAsValueReference(suspendPropertyDescriptor, context),
+            baseClassRef = coroutineBaseClassRef,
+            stateName = getCoroutinePropertyName("state"),
+            exceptionStateName = getCoroutinePropertyName("exceptionState"),
+            finallyPathName = getCoroutinePropertyName("finallyPath"),
+            resultName = getCoroutinePropertyName("result"),
+            exceptionName = getCoroutinePropertyName("exception"),
+            facadeName = getCoroutinePropertyName("facade"),
+            hasController = hasController,
+            isLambda = isLambda,
+            hasReceiver = descriptor.dispatchReceiverParameter != null
+    )
+}
+
+private val COROUTINES_INTRINSICS_PACKAGE_FQ_NAME = FqName("kotlin.coroutines.intrinsics")