JS: improve inline suspend function declaration splitter

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

@@ -73,7 +73,6 @@ fun transformCoroutines(fragments: Iterable<JsProgramFragment>) {
     val coroutineTransformer = CoroutineTransformer()
     for (fragment in fragments) {
         ImportInfoFragmentInliningScope.process(fragment) { scope ->
-            InlineSuspendFunctionSplitter(scope).accept(scope.allCode)
             coroutineTransformer.accept(scope.allCode)
         }
     }

js/js.inliner/src/org/jetbrains/kotlin/js/inline/InlineSuspendFunctionSplitter.kt

@@ -8,64 +8,45 @@ package org.jetbrains.kotlin.js.inline
 import org.jetbrains.kotlin.js.backend.ast.*
 import org.jetbrains.kotlin.js.backend.ast.metadata.coroutineMetadata
 import org.jetbrains.kotlin.js.backend.ast.metadata.isInlineableCoroutineBody
+import org.jetbrains.kotlin.js.inline.util.FunctionWithWrapper
 import org.jetbrains.kotlin.js.translate.expression.InlineMetadata
 
-// Goes through `suspend inline` function definitions and transforms into two entites:
+// Goes through `suspend inline` function definitions and transforms into two entities:
 // 1. A `suspend` function definition, which could be invoked at run time
 // 2. A `suspend inline` function definition, which doesn't work at run time, but is understood by the inliner.
-// TODO remove the wrapped version for `private inline suspend fun`'s
 class InlineSuspendFunctionSplitter(
     val scope: ImportInfoFragmentInliningScope
 ) : JsVisitorWithContextImpl() {
 
-    override fun endVisit(x: JsExpressionStatement, ctx: JsContext<*>) {
-        val e = x.expression
-        if (e is JsBinaryOperation) {
-            if (e.operator == JsBinaryOperator.ASG) {
-                e.arg2?.let { argument2 ->
-                    val splitSuspendInlineFunction = splitExportedSuspendInlineFunctionDeclarations(argument2)
-                    if (splitSuspendInlineFunction != null) {
-                        e.arg2 = splitSuspendInlineFunction
-                    }
-                }
+    override fun visit(x: JsInvocation, ctx: JsContext<JsNode>): Boolean {
+        // Is it `defineInlineFunction('tag', ...)`?
+        InlineMetadata.decompose(x)?.let { metadata ->
+            val fn = metadata.function
+            if (fn.function.coroutineMetadata != null) {
+                // This function will be exported to JS
+                ctx.replaceMe(scope.importFunctionDefinition(InlineFunctionDefinition(fn, metadata.tag.value)))
+
+                // Original function should be not be transformed into a state machine
+                fn.function.name = null
+                fn.function.coroutineMetadata = null
+                fn.function.isInlineableCoroutineBody = true
+
+                // Keep the `defineInlineFunction` for the inliner to find
+                lastStatementLevelContext.addNext(x.makeStmt())
+
             }
+            return false
         }
 
-        super.endVisit(x, ctx)
-    }
-
-    override fun endVisit(x: JsVars.JsVar, ctx: JsContext<*>) {
-        val initExpression = x.initExpression ?: return
-
-        val splitSuspendInlineFunction = splitExportedSuspendInlineFunctionDeclarations(initExpression)
-        if (splitSuspendInlineFunction != null) {
-            x.initExpression = splitSuspendInlineFunction
-        }
-    }
-
-    private fun splitExportedSuspendInlineFunctionDeclarations(expression: JsExpression): JsFunction? {
-        val inlineMetadata = InlineMetadata.decompose(expression)
-        if (inlineMetadata != null) {
-            inlineMetadata.function.let { f ->
-                if (f.function.coroutineMetadata != null) {
-                    val statementContext = lastStatementLevelContext
-
-                    // This function will be exported to JS
-                    val function = scope.importFunctionDefinition(InlineFunctionDefinition(inlineMetadata.function, inlineMetadata.tag.value))
-
-                    // Original function should be not be transformed into a state machine
-                    f.function.name = null
-                    f.function.coroutineMetadata = null
-                    f.function.isInlineableCoroutineBody = true
-
-                    // Keep the `defineInlineFunction` for the inliner to find
-                    statementContext.addNext(expression.makeStmt())
-
-                    // Return the function body to be used without inlining.
-                    return function
-                }
+        // Is it `wrapFunction(...)`? Which means it'a a private inline function
+        InlineMetadata.tryExtractFunction(x)?.let { fn ->
+            if (fn.function.coroutineMetadata != null) {
+                // This function will be exported to JS
+                ctx.replaceMe(scope.importFunctionDefinition(InlineFunctionDefinition(fn, null)))
             }
+            return false
         }
-        return null
+
+        return super.visit(x, ctx)
     }
 }

js/js.inliner/src/org/jetbrains/kotlin/js/inline/InliningScope.kt

@@ -172,6 +172,10 @@ class ImportInfoFragmentInliningScope private constructor(
         fragment.declarationBlock.statements.addAll(0, additionalDeclarations)
 
         // post-processing
+
+        // If run separately `private inline suspend fun`'s local declarations get inlined twice.
+        InlineSuspendFunctionSplitter(this).accept(allCode)
+
         simplifyWrappedFunctions(allCode)
         removeUnusedFunctionDefinitions(allCode, collectNamedFunctions(allCode))
         removeUnusedImports(fragment, allCode)