JS: fix noInline version of inline suspend fun (KT-27611 fixed)

compiler/testData/codegen/box/coroutines/noStdLib/crossinline.kt

@@ -1,5 +1,4 @@
 // IGNORE_BACKEND: JVM_IR
-// IGNORE_BACKEND: JS
 // COMMON_COROUTINES_TEST
 // WITH_COROUTINES
 // CHECK_BYTECODE_LISTING

compiler/testData/codegen/box/coroutines/tailCallOptimizations/crossinline.kt

@@ -1,5 +1,4 @@
 // IGNORE_BACKEND: JVM_IR
-// IGNORE_BACKEND: JS
 // COMMON_COROUTINES_TEST
 // WITH_RUNTIME
 // WITH_COROUTINES

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

@@ -24,39 +24,21 @@ import org.jetbrains.kotlin.js.translate.expression.InlineMetadata
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
 
 class CoroutineTransformer : JsVisitorWithContextImpl() {
-    private val additionalStatementsByNode = mutableMapOf<JsNode, List<JsStatement>>()
 
-    override fun endVisit(x: JsExpressionStatement, ctx: JsContext<in JsStatement>) {
+    val functionName = mutableMapOf<JsFunction, String?>()
-        additionalStatementsByNode.remove(x)?.forEach { ctx.addNext(it) }
-        super.endVisit(x, ctx)
-    }
-
-    override fun endVisit(x: JsVars, ctx: JsContext<in JsStatement>) {
-        for (v in x.vars) {
-            additionalStatementsByNode.remove(v)?.forEach { ctx.addNext(it) }
-        }
-        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?.function
+            InlineMetadata.tryExtractFunction(rhs)?.let { wrapper ->
-            if (function?.coroutineMetadata != null) {
+                val function = wrapper.function
                 val name = ((lhs as? JsNameRef)?.name ?: function.name)?.ident
-                additionalStatementsByNode[x] = CoroutineFunctionTransformer(function, name).transform()
+                functionName[function] = name
-                function.coroutineMetadata = null
-                return false
-            }
-        }
-        else if (expression is JsFunction) {
-            if (expression.coroutineMetadata != null) {
-                additionalStatementsByNode[x] = CoroutineFunctionTransformer(expression, expression.name?.ident).transform()
-                expression.coroutineMetadata = null
-                return false
             }
+        } else if (expression is JsFunction) {
+            functionName[expression] = expression.name?.ident
         }
         return super.visit(x, ctx)
     }
@@ -66,18 +48,19 @@ class CoroutineTransformer : JsVisitorWithContextImpl() {
             x.body = transformCoroutineMetadataToSpecialFunctions(x.body)
             return false
         }
+        if (x.coroutineMetadata != null) {
+            lastStatementLevelContext.addPrevious(CoroutineFunctionTransformer(x, functionName[x]).transform())
+            x.coroutineMetadata = null
+            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?.function
+            InlineMetadata.tryExtractFunction(initExpression)?.let { wrapper ->
-            if (function?.coroutineMetadata != null) {
+                functionName[wrapper.function] = x.name.ident
-                val name = x.name.ident
-                additionalStatementsByNode[x] = CoroutineFunctionTransformer(function, name).transform()
-                function.coroutineMetadata = null
-                return false
             }
         }
         return super.visit(x, ctx)

js/js.inliner/src/org/jetbrains/kotlin/js/inline/JsInliner.java

@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.js.inline.context.FunctionContext;
 import org.jetbrains.kotlin.js.inline.context.InliningContext;
 import org.jetbrains.kotlin.js.inline.context.NamingContext;
 import org.jetbrains.kotlin.js.inline.util.*;
+import org.jetbrains.kotlin.js.translate.expression.InlineMetadata;
 import org.jetbrains.kotlin.resolve.inline.InlineStrategy;
 
 import java.util.*;
@@ -348,6 +349,61 @@ public class JsInliner extends JsVisitorWithContextImpl {
         }
     }
 
+    @Override
+    public void endVisit(@NotNull JsExpressionStatement x, @NotNull JsContext ctx) {
+        JsExpression e = x.getExpression();
+        if (e instanceof JsBinaryOperation) {
+            JsBinaryOperation binOp = (JsBinaryOperation) e;
+            if (binOp.getOperator() == JsBinaryOperator.ASG) {
+                JsFunction splitSuspendInlineFunction = splitExportedSuspendInlineFunctionDeclarations(binOp.getArg2());
+                if (splitSuspendInlineFunction != null) {
+                    binOp.setArg2(splitSuspendInlineFunction);
+                }
+            }
+        }
+
+        super.endVisit(x, ctx);
+    }
+
+    @Override
+    public void endVisit(@NotNull JsVars.JsVar x, @NotNull JsContext ctx) {
+        JsFunction splitSuspendInlineFunction = splitExportedSuspendInlineFunctionDeclarations(x.getInitExpression());
+        if (splitSuspendInlineFunction != null) {
+            x.setInitExpression(splitSuspendInlineFunction);
+        }
+    }
+
+    @Nullable
+    private JsFunction splitExportedSuspendInlineFunctionDeclarations(@NotNull JsExpression expression) {
+        InlineMetadata inlineMetadata = InlineMetadata.decompose(expression);
+        if (inlineMetadata != null) {
+            FunctionWithWrapper functionWithWrapper = inlineMetadata.getFunction();
+            JsFunction originalFunction = functionWithWrapper.getFunction();
+            if (MetadataProperties.getCoroutineMetadata(originalFunction) != null) {
+                JsContext<JsStatement> statementContext = getLastStatementLevelContext();
+
+                // This function will be exported to JS
+                JsFunction function = originalFunction.deepCopy();
+
+                // Original function should be not be transformed into a state machine
+                originalFunction.setName(null);
+                MetadataProperties.setCoroutineMetadata(originalFunction, null);
+                MetadataProperties.setInlineableCoroutineBody(originalFunction, true);
+                if (functionWithWrapper.getWrapperBody() != null) {
+                    // Extract local declarations
+                    applyWrapper(functionWithWrapper.getWrapperBody(), function, originalFunction, new JsInliningContext(statementContext));
+                }
+
+                // Keep the `defineInlineFunction` for the inliner to find
+                statementContext.addNext(expression.makeStmt());
+
+                // Return the function body to be used without inlining.
+                return function;
+            }
+        }
+        return null;
+    }
+
     @Override
     protected void doAcceptStatementList(List<JsStatement> statements) {
         // at top level of js ast, contexts stack can be empty,

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

@@ -907,6 +907,11 @@ public class BoxJsTestGenerated extends AbstractBoxJsTest {
         public void testNativeExceptions() throws Exception {
             runTest("js/js.translator/testData/box/coroutines/nativeExceptions.kt");
         }
+
+        @TestMetadata("onlyInlineSuspendFunction.kt")
+        public void testOnlyInlineSuspendFunction() throws Exception {
+            runTest("js/js.translator/testData/box/coroutines/onlyInlineSuspendFunction.kt");
+        }
     }
 
     @TestMetadata("js/js.translator/testData/box/crossModuleRef")

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

@@ -907,6 +907,11 @@ public class IrBoxJsTestGenerated extends AbstractIrBoxJsTest {
         public void testNativeExceptions() throws Exception {
             runTest("js/js.translator/testData/box/coroutines/nativeExceptions.kt");
         }
+
+        @TestMetadata("onlyInlineSuspendFunction.kt")
+        public void testOnlyInlineSuspendFunction() throws Exception {
+            runTest("js/js.translator/testData/box/coroutines/onlyInlineSuspendFunction.kt");
+        }
     }
 
     @TestMetadata("js/js.translator/testData/box/crossModuleRef")

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

@@ -21,14 +21,8 @@ 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.js.backend.ast.JsFunction
 import org.jetbrains.kotlin.js.backend.ast.JsName
-import org.jetbrains.kotlin.js.backend.ast.metadata.coroutineMetadata
-import org.jetbrains.kotlin.js.backend.ast.metadata.isInlineableCoroutineBody
-import org.jetbrains.kotlin.js.descriptorUtils.shouldBeExported
-import org.jetbrains.kotlin.js.translate.callTranslator.CallTranslator
 import org.jetbrains.kotlin.js.translate.context.TranslationContext
-import org.jetbrains.kotlin.js.translate.expression.InlineMetadata
 import org.jetbrains.kotlin.js.translate.expression.translateAndAliasParameters
 import org.jetbrains.kotlin.js.translate.expression.translateFunction
 import org.jetbrains.kotlin.js.translate.expression.wrapWithInlineMetadata
@@ -36,7 +30,6 @@ import org.jetbrains.kotlin.js.translate.general.TranslatorVisitor
 import org.jetbrains.kotlin.js.translate.utils.*
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.descriptorUtil.isExtensionProperty
-import org.jetbrains.kotlin.resolve.source.getPsi
 
 abstract class AbstractDeclarationVisitor : TranslatorVisitor<Unit>()  {
     override fun emptyResult(context: TranslationContext) { }
@@ -102,20 +95,7 @@ abstract class AbstractDeclarationVisitor : TranslatorVisitor<Unit>()  {
             null
         }
 
-        if (descriptor.isSuspend && descriptor.isInline && descriptor.shouldBeExported(context.config) && functionAndContext != null) {
+        addFunction(descriptor, functionAndContext?.first, expression)
-            val exportFunction = functionAndContext.first as JsFunction
-            exportFunction.fillCoroutineMetadata(context, descriptor, exportFunction.coroutineMetadata!!.hasController)
-            addFunction(descriptor, exportFunction, expression)
-
-            val innerContext = functionAndContext.second
-            val inlineFunction = functionAndContext.first.deepCopy() as JsFunction
-            inlineFunction.name = null
-            inlineFunction.coroutineMetadata = null
-            inlineFunction.isInlineableCoroutineBody = true
-            context.addDeclarationStatement(innerContext.wrapWithInlineMetadata(context, inlineFunction, descriptor).makeStmt())
-        } else {
-            addFunction(descriptor, functionAndContext?.first, expression)
-        }
     }
 
     override fun visitTypeAlias(typeAlias: KtTypeAlias, data: TranslationContext?) {}
@@ -137,14 +117,8 @@ abstract class AbstractDeclarationVisitor : TranslatorVisitor<Unit>()  {
             function.body.statements += FunctionBodyTranslator.setDefaultValueForArguments(descriptor, innerContext)
         }
         innerContext.translateFunction(expression, function)
-        val result = if (descriptor.isSuspend && descriptor.shouldBeExported(context.config)) {
-            function
-        }
-        else {
-            innerContext.wrapWithInlineMetadata(context, function, descriptor)
-        }
 
-        return Pair(result, innerContext)
+        return Pair(innerContext.wrapWithInlineMetadata(context, function, descriptor), innerContext)
     }
 
     // used from kotlinx.serialization

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

@@ -21,10 +21,7 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.incremental.js.IncrementalResultsConsumer
 import org.jetbrains.kotlin.js.backend.ast.*
-import org.jetbrains.kotlin.js.backend.ast.metadata.descriptor
+import org.jetbrains.kotlin.js.backend.ast.metadata.*
-import org.jetbrains.kotlin.js.backend.ast.metadata.functionDescriptor
-import org.jetbrains.kotlin.js.backend.ast.metadata.hasDefaultValue
-import org.jetbrains.kotlin.js.backend.ast.metadata.type
 import org.jetbrains.kotlin.js.config.JSConfigurationKeys
 import org.jetbrains.kotlin.js.descriptorUtils.shouldBeExported
 import org.jetbrains.kotlin.js.inline.util.FunctionWithWrapper

js/js.translator/testData/box/coroutines/onlyInlineSuspendFunction.kt

@@ -0,0 +1,61 @@
+// IGNORE_BACKEND: JS_IR
+// EXPECTED_REACHABLE_NODES: 1280
+
+// MODULE: lib
+// FILE: lib.kt
+import kotlin.coroutines.*
+
+var resume: () -> Unit = {}
+
+suspend fun suspendAndReturn(msg: String): String {
+    return suspendCoroutine<String> { cont ->
+        resume = {
+            cont.resume(msg)
+        }
+    }
+}
+
+suspend fun <T> runSus(block: suspend () -> T): T = block()
+
+abstract class A {
+    fun f() = o()
+
+    abstract fun o(): String
+}
+
+// MODULE: libInline(lib)
+// FILE: libInline.kt
+
+suspend inline fun foo(): String {
+    val a = object : A() {
+        override fun o(): String = "O"
+    }
+    val k = "K"
+
+    return suspendAndReturn(a.f()) + runSus {
+        val b = object : A() {
+            override fun o(): String = k
+        }
+        runSus { suspendAndReturn(b.f()) }
+    }
+}
+
+// MODULE: main(libInline, lib)
+// FILE: main.kt
+
+import kotlin.coroutines.*
+
+fun box(): String {
+    var testResult: String = "fail"
+
+    val continuation = Continuation<String>(EmptyCoroutineContext) { result ->
+        testResult = result.getOrThrow()
+    }
+
+    // Test the noinline version of public inline suspend functions
+    js("libInline").foo(continuation)
+    resume()
+    resume()
+
+    return testResult
+}

js/js.translator/testData/lineNumbers/coroutine.kt

@@ -14,4 +14,4 @@ suspend fun bar(): Unit {
     println(a + b)
 }
 
-// LINES: 39 39 4 4 4 5 5 44 44 5 88 44 5 5 6 4 4 4 15 9 9 9 9 15 9 9 9 * 15 10 10 11 11 11 2 11 11 * 11 12 12 13 13 13 2 13 13 13 13 14 14
+// LINES: 39 39 4 4 4 5 5 44 44 5 88 44 5 5 6 4 4 4 15 9 9 9 * 15 10 10 11 11 11 2 11 11 * 11 12 12 13 13 13 2 13 13 13 13 14 14 * 15 9 9 9 9

js/js.translator/testData/lineNumbers/coroutineNullAssertion.kt

@@ -8,4 +8,4 @@ suspend fun foo() {
 suspend fun delay() {
 }
 
-// LINES: 6 1 1 1 1 6 1 1 * 6 2 2 2 2 2 * 3 3 4 4 4 2 4 4 5 5 * 9
+// LINES: 6 1 1 * 6 2 2 2 2 2 * 3 3 4 4 4 2 4 4 5 5 * 6 1 1 1 1 9