JS: support stack unwinding convention in coroutines

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

@@ -1,8 +1,8 @@
 // WITH_RUNTIME
 // WITH_REFLECT
-// CHECK_NOT_CALLED: suspendInline_die06n$
+// CHECK_NOT_CALLED: suspendInline_61zpoe$
-// CHECK_NOT_CALLED: suspendInline_nesahw$
+// CHECK_NOT_CALLED: suspendInline_6r51u9$
-// CHECK_NOT_CALLED: suspendInline_grpnnl$
+// CHECK_NOT_CALLED: suspendInline
 class Controller {
     fun withValue(v: String, x: Continuation<String>) {
         x.resume(v)

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

@@ -1,8 +1,8 @@
 // WITH_RUNTIME
 // WITH_REFLECT
-// CHECK_NOT_CALLED: suspendInline_die06n$
+// CHECK_NOT_CALLED: suspendInline_61zpoe$
-// CHECK_NOT_CALLED: suspendInline_nesahw$
+// CHECK_NOT_CALLED: suspendInline_6r51u9$
-// CHECK_NOT_CALLED: suspendInline_grpnnl$
+// CHECK_NOT_CALLED: suspendInline
 class Controller {
     suspend inline fun suspendInline(v: String): String = v
 

compiler/testData/codegen/box/coroutines/stackUnwinding/suspendInCycle.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JS
 class Controller {
     suspend fun suspendHere(): Int = suspendWithCurrentContinuation { x ->
         1

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

@@ -103,6 +103,8 @@ var JsNameRef.coroutineResult by MetadataProperty(default = false)
  */
 var JsNameRef.coroutineController by MetadataProperty(default = false)
 
+var JsFunction.suspendObjectRef: JsExpression? by MetadataProperty(default = null)
+
 var JsFunction.continuationInterfaceRef: JsExpression? by MetadataProperty(default = null)
 
 var JsName.imported by MetadataProperty(default = false)

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

@@ -361,7 +361,18 @@ class CoroutineBodyTransformer(
 
     private fun handleSuspend(invocation: JsInvocation) {
         val invokeExpression = if (invocation.isFakeSuspend) invocation.arguments.getOrNull(0) else invocation
-        currentStatements += JsReturn(invokeExpression)
+        val suspendObjectVar = context.suspendObjectVar
+        val statements = if (invokeExpression == null || suspendObjectVar == null) {
+            listOf(JsReturn(invokeExpression))
+        }
+        else {
+            val resultRef = JsNameRef(context.resultFieldName, JsLiteral.THIS).apply { sideEffects = SideEffectKind.DEPENDS_ON_STATE }
+            val invocationStatement = JsAstUtils.assignment(resultRef, invokeExpression).makeStmt()
+            val suspendCondition = JsAstUtils.equality(resultRef.deepCopy(), JsAstUtils.pureFqn(suspendObjectVar, null))
+            val suspendIfNeeded = JsIf(suspendCondition, JsReturn())
+            listOf(invocationStatement, suspendIfNeeded, JsBreak())
+        }
+        currentStatements += statements
         currentBlock = suspendTarget!!
     }
 

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

@@ -48,7 +48,7 @@ class CoroutineFunctionTransformer(
             function.scope.declareName(throwId)
         }
 
-        val context = CoroutineTransformationContext(function.scope)
+        val context = CoroutineTransformationContext(function.scope, function.suspendObjectRef != null)
         val bodyTransformer = CoroutineBodyTransformer(program, context, throwName)
         bodyTransformer.preProcess(body)
         body.statements.forEach { it.accept(bodyTransformer) }
@@ -87,15 +87,15 @@ class CoroutineFunctionTransformer(
         val parameterNames = (function.parameters.map { it.name } + innerFunction?.parameters?.map { it.name }.orEmpty()).toSet()
 
         constructor.body.statements.run {
-            assign(context.stateFieldName, program.getNumberLiteral(0))
+            assignToField(context.stateFieldName, program.getNumberLiteral(0))
-            assign(context.exceptionStateName, program.getNumberLiteral(globalCatchBlockIndex))
+            assignToField(context.exceptionStateName, program.getNumberLiteral(globalCatchBlockIndex))
             if (hasFinallyBlocks) {
-                assign(context.finallyPathFieldName, JsLiteral.NULL)
+                assignToField(context.finallyPathFieldName, JsLiteral.NULL)
             }
-            assign(context.controllerFieldName, controllerName.makeRef())
+            assignToField(context.controllerFieldName, controllerName.makeRef())
             for (localVariable in localVariables) {
                 val value = if (localVariable !in parameterNames) JsLiteral.NULL else localVariable.makeRef()
-                assign(function.scope.getFieldName(localVariable), value)
+                assignToField(function.scope.getFieldName(localVariable), value)
             }
         }
 
@@ -161,7 +161,12 @@ class CoroutineFunctionTransformer(
         functionWithBody.body.statements.clear()
 
         resumeFunction.body.statements.apply {
-            assign(context.resultFieldName, resumeParameter.makeRef())
+            assignToField(context.resultFieldName, resumeParameter.makeRef())
+            if (context.suspendObjectVar != null) {
+                add(JsAstUtils.newVar(context.suspendObjectVar!!, function.suspendObjectRef!!.deepCopy()).apply {
+                    synthetic = true
+                })
+            }
             this += coroutineBody
         }
 
@@ -249,7 +254,7 @@ class CoroutineFunctionTransformer(
         return functions.mapNotNull { it as? FunctionDescriptor }.firstOrNull { it.kind.isReal }
     }
 
-    private fun MutableList<JsStatement>.assign(fieldName: JsName, value: JsExpression) {
+    private fun MutableList<JsStatement>.assignToField(fieldName: JsName, value: JsExpression) {
         this += JsAstUtils.assignment(JsNameRef(fieldName, JsLiteral.THIS), value).makeStmt()
     }
 

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

@@ -18,7 +18,7 @@ package org.jetbrains.kotlin.js.coroutine
 
 import com.google.dart.compiler.backend.js.ast.JsScope
 
-class CoroutineTransformationContext(private val scope: JsScope) {
+class CoroutineTransformationContext(private val scope: JsScope, private val stackUnwinding: Boolean) {
     val entryBlock = CoroutineBlock()
     val globalCatchBlock = CoroutineBlock()
     val resultFieldName by lazy { scope.declareFreshName("\$result") }
@@ -27,4 +27,5 @@ class CoroutineTransformationContext(private val scope: JsScope) {
     val controllerFieldName by lazy { scope.declareFreshName("\$controller") }
     val exceptionStateName by lazy { scope.declareFreshName("\$exceptionState") }
     val finallyPathFieldName by lazy { scope.declareFreshName("\$finallyPath") }
+    val suspendObjectVar by lazy { if (stackUnwinding) scope.declareFreshName("\$suspendObject") else null }
 }

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -5786,13 +5786,7 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             @TestMetadata("suspendInCycle.kt")
             public void testSuspendInCycle() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/stackUnwinding/suspendInCycle.kt");
-                try {
+                doTest(fileName);
-                    doTest(fileName);
-                }
-                catch (Throwable ignore) {
-                    return;
-                }
-                throw new AssertionError("Looks like this test can be unmuted. Remove IGNORE_BACKEND directive for that.");
             }
         }
     }

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

@@ -16,9 +16,7 @@
 
 package org.jetbrains.kotlin.js.translate.callTranslator
 
-import com.google.dart.compiler.backend.js.ast.JsExpression
+import com.google.dart.compiler.backend.js.ast.*
-import com.google.dart.compiler.backend.js.ast.JsInvocation
-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
@@ -35,7 +33,6 @@ 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
@@ -160,12 +157,13 @@ private fun translateFunctionCall(
     return callExpression
 }
 
-private fun translateCallWithContinuation(context: TranslationContext,resolvedCall: ResolvedCall<out FunctionDescriptor>): JsExpression {
+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
+    context.currentBlock.statements += JsReturn(invocation)
+    return JsLiteral.NULL
 }
 
 fun computeExplicitReceiversForInvoke(

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

@@ -34,6 +34,7 @@ import org.jetbrains.kotlin.js.translate.utils.FunctionBodyTranslator.translateF
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
 import org.jetbrains.kotlin.js.translate.utils.TranslationUtils.simpleReturnFunction
 import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtDeclarationWithBody
 import org.jetbrains.kotlin.psi.KtParameter
@@ -42,6 +43,10 @@ import org.jetbrains.kotlin.resolve.inline.InlineUtil
 import org.jetbrains.kotlin.serialization.deserialization.findClassAcrossModuleDependencies
 
 class LiteralFunctionTranslator(context: TranslationContext) : AbstractTranslator(context) {
+    companion object {
+        private val SUSPEND_FQ_NAME = "kotlin.coroutines.Suspend"
+    }
+
     fun translate(
             declaration: KtDeclarationWithBody,
             continuationType: ClassDescriptor? = null,
@@ -83,6 +88,14 @@ class LiteralFunctionTranslator(context: TranslationContext) : AbstractTranslato
             invokingContext.getInnerNameForDescriptor(descriptor)
         }
 
+        val suspendObjectRef = if (descriptor.isCoroutineLambda && KotlinBuiltIns.isUnit(descriptor.returnType!!)) {
+            val suspendObjectDescriptor = context().currentModule.builtIns.getBuiltInClassByFqName(FqName(SUSPEND_FQ_NAME))
+            ReferenceTranslator.translateAsValueReference(suspendObjectDescriptor, context())
+        }
+        else {
+            null
+        }
+
         if (tracker.hasCapturedExceptContaining()) {
             val lambdaCreator = simpleReturnFunction(invokingContext.scope(), lambda)
             lambdaCreator.name = invokingContext.getInnerNameForDescriptor(descriptor)
@@ -94,6 +107,7 @@ class LiteralFunctionTranslator(context: TranslationContext) : AbstractTranslato
             }
             lambdaCreator.name.staticRef = lambdaCreator
             lambdaCreator.continuationInterfaceRef = invokingContext.getContinuationInterfaceReference()
+            lambdaCreator.suspendObjectRef = suspendObjectRef
             return lambdaCreator.withCapturedParameters(descriptor, descriptor.wrapContextForCoroutineIfNecessary(functionContext),
                                                         invokingContext)
         }
@@ -105,6 +119,7 @@ class LiteralFunctionTranslator(context: TranslationContext) : AbstractTranslato
         invokingContext.addDeclarationStatement(lambda.makeStmt())
         lambda.name.staticRef = lambda
         lambda.continuationInterfaceRef = invokingContext.getContinuationInterfaceReference()
+        lambda.suspendObjectRef = suspendObjectRef
         return getReferenceToLambda(invokingContext, descriptor, lambda.name)
     }
 

js/js.translator/src/org/jetbrains/kotlin/js/translate/reference/ReferenceTranslator.java

@@ -112,7 +112,7 @@ public final class ReferenceTranslator {
     private static JsExpression getLazyReferenceToObject(@NotNull ClassDescriptor descriptor, @NotNull TranslationContext context) {
         DeclarationDescriptor container = descriptor.getContainingDeclaration();
         JsExpression qualifier = context.getInnerReference(container);
-        return JsAstUtils.pureFqn(context.getNameForDescriptor(descriptor), qualifier);
+        return new JsNameRef(context.getNameForDescriptor(descriptor), qualifier);
     }
 
     private static boolean shouldTranslateAsFQN(@NotNull DeclarationDescriptor descriptor, @NotNull TranslationContext context) {