JS: add partial tail-call optimization for suspend functions

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

@@ -2,6 +2,8 @@
 // WITH_COROUTINES
 import helpers.*
 // CHECK_BYTECODE_LISTING
+// CHECK_NEW_COUNT: function=suspendHere count=1
+// CHECK_NEW_COUNT: function=mainSuspend count=1
 import kotlin.coroutines.experimental.*
 import kotlin.coroutines.experimental.intrinsics.*
 

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

@@ -2,6 +2,8 @@
 // WITH_COROUTINES
 import helpers.*
 // CHECK_BYTECODE_LISTING
+// CHECK_NEW_COUNT: function=suspendHere count=0
+// CHECK_NEW_COUNT: function=complexSuspend count=0
 import kotlin.coroutines.experimental.*
 import kotlin.coroutines.experimental.intrinsics.*
 

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

@@ -1,6 +1,7 @@
 // WITH_RUNTIME
 // WITH_COROUTINES
 // CHECK_BYTECODE_LISTING
+// CHECK_NEW_COUNT: function=suspendHere count=0
 import helpers.*
 import kotlin.coroutines.experimental.*
 import kotlin.coroutines.experimental.intrinsics.*

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

@@ -96,8 +96,6 @@ var HasMetadata.sideEffects: SideEffectKind by MetadataProperty(default = SideEf
  */
 var JsExpression.isSuspend: Boolean by MetadataProperty(default = false)
 
-var JsExpression.isTailCallSuspend: Boolean by MetadataProperty(default = false)
-
 /**
  * Denotes a reference to coroutine's `result` field that contains result of
  * last suspended invocation.
@@ -115,6 +113,8 @@ var JsNameRef.coroutineController by MetadataProperty(default = false)
  */
 var JsNameRef.coroutineReceiver by MetadataProperty(default = false)
 
+var JsFunction.forceStateMachine by MetadataProperty(default = false)
+
 var JsName.imported by MetadataProperty(default = false)
 
 var JsFunction.coroutineMetadata: CoroutineMetadata? by MetadataProperty(default = null)

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

@@ -19,11 +19,14 @@ package org.jetbrains.kotlin.js.coroutine
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.js.backend.ast.*
 import org.jetbrains.kotlin.js.backend.ast.metadata.coroutineMetadata
+import org.jetbrains.kotlin.js.backend.ast.metadata.forceStateMachine
+import org.jetbrains.kotlin.js.backend.ast.metadata.isSuspend
+import org.jetbrains.kotlin.js.backend.ast.metadata.synthetic
 import org.jetbrains.kotlin.js.inline.clean.FunctionPostProcessor
 import org.jetbrains.kotlin.js.inline.util.collectLocalVariables
 import org.jetbrains.kotlin.js.inline.util.getInnerFunction
 import org.jetbrains.kotlin.js.translate.context.Namer
-import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
+import org.jetbrains.kotlin.js.translate.utils.JsAstUtils.*
 import org.jetbrains.kotlin.js.translate.utils.finalElement
 
 class CoroutineFunctionTransformer(private val function: JsFunction, name: String?) {
@@ -35,6 +38,11 @@ class CoroutineFunctionTransformer(private val function: JsFunction, name: Strin
     private val className = JsScope.declareTemporaryName("Coroutine\$${name ?: "anonymous"}")
 
     fun transform(): List<JsStatement> {
+        if (isTailCall() && !function.forceStateMachine) {
+            transformSimple()
+            return emptyList()
+        }
+
         val context = CoroutineTransformationContext(function.scope, function)
         val bodyTransformer = CoroutineBodyTransformer(context)
         bodyTransformer.preProcess(body)
@@ -57,6 +65,64 @@ class CoroutineFunctionTransformer(private val function: JsFunction, name: Strin
         return additionalStatements
     }
 
+    private fun isTailCall(): Boolean {
+        val suspendCalls = hashSetOf<JsExpression>()
+        body.accept(object : RecursiveJsVisitor() {
+            override fun visitElement(node: JsNode) {
+                if (node is JsExpression && node.isSuspend) {
+                    suspendCalls += node
+                }
+                super.visitElement(node)
+            }
+        })
+
+        if (suspendCalls.isEmpty()) return true
+
+        body.accept(object : RecursiveJsVisitor() {
+            override fun visitBlock(x: JsBlock) {
+                super.visitBlock(x)
+
+                if (body.statements.size < 2) return
+
+                val lastStatement = body.statements.last() as? JsReturn ?: return
+                if (!lastStatement.expression.isStateMachineResult()) return
+
+                val statementBeforeLast = body.statements[body.statements.lastIndex - 1] as? JsExpressionStatement ?: return
+                val suspendExpression = statementBeforeLast.expression
+                if (suspendExpression in suspendCalls) {
+                    suspendCalls -= suspendExpression
+                }
+                else {
+                    decomposeAssignment(suspendExpression)?.let { (lhs, rhs) ->
+                        if (rhs in suspendCalls && lhs.isStateMachineResult()) {
+                            suspendCalls -= rhs
+                        }
+                    }
+                }
+            }
+        })
+
+        return suspendCalls.isEmpty()
+    }
+
+    private fun transformSimple() {
+        val continuationParam = function.parameters.last()
+        val resultVar = JsScope.declareTemporaryName("\$result")
+        body.replaceSpecialReferencesInSimpleFunction(continuationParam, resultVar)
+        body.statements.add(0, newVar(resultVar, null).apply { synthetic = true })
+
+        object : JsVisitorWithContextImpl() {
+            override fun endVisit(x: JsExpressionStatement, ctx: JsContext<in JsStatement>) {
+                if (x.expression.isSuspend) {
+                    ctx.replaceMe(assignment(pureFqn(resultVar, null), x.expression).source(x.source).makeStmt())
+                }
+                super.endVisit(x, ctx)
+            }
+        }.accept(body)
+
+        FunctionPostProcessor(functionWithBody).apply()
+    }
+
     private fun generateContinuationConstructor(
             context: CoroutineTransformationContext,
             statements: MutableList<JsStatement>,
@@ -108,11 +174,11 @@ class CoroutineFunctionTransformer(private val function: JsFunction, name: Strin
     }
 
     private fun generateCoroutinePrototype(constructorName: JsName): List<JsStatement> {
-        val prototype = JsAstUtils.prototypeOf(JsNameRef(constructorName))
+        val prototype = prototypeOf(JsNameRef(constructorName))
 
         val baseClass = Namer.createObjectWithPrototypeFrom(function.coroutineMetadata!!.baseClassRef.deepCopy())
-        val assignPrototype = JsAstUtils.assignment(prototype, baseClass)
+        val assignPrototype = assignment(prototype, baseClass)
-        val assignConstructor = JsAstUtils.assignment(JsNameRef("constructor", prototype.deepCopy()), JsNameRef(constructorName))
+        val assignConstructor = assignment(JsNameRef("constructor", prototype.deepCopy()), JsNameRef(constructorName))
         return listOf(assignPrototype.makeStmt(), assignConstructor.makeStmt())
     }
 
@@ -127,7 +193,7 @@ class CoroutineFunctionTransformer(private val function: JsFunction, name: Strin
             propertyInitializers += JsPropertyInitializer(JsNameRef(Namer.METADATA_SUPERTYPES), JsArrayLiteral(listOf(baseClassRefRef)))
         }
 
-        return JsAstUtils.assignment(JsNameRef(Namer.METADATA, constructorName.makeRef()), metadataObject).makeStmt()
+        return assignment(JsNameRef(Namer.METADATA, constructorName.makeRef()), metadataObject).makeStmt()
     }
 
     private fun generateDoResume(
@@ -172,7 +238,7 @@ class CoroutineFunctionTransformer(private val function: JsFunction, name: Strin
         functionWithBody.parameters += JsParameter(suspendedName)
 
         val instanceName = JsScope.declareTemporaryName("instance")
-        functionWithBody.body.statements += JsAstUtils.newVar(instanceName, instantiation)
+        functionWithBody.body.statements += newVar(instanceName, instantiation)
 
         val invokeResume = JsReturn(JsInvocation(JsNameRef(context.metadata.doResumeName, instanceName.makeRef()), JsNullLiteral())
                                             .source(psiElement))
@@ -191,14 +257,14 @@ class CoroutineFunctionTransformer(private val function: JsFunction, name: Strin
         val stateRef = JsNameRef(context.metadata.stateName, JsThisRef())
         val exceptionStateRef = JsNameRef(context.metadata.exceptionStateName, JsThisRef())
 
-        val isFromGlobalCatch = JsAstUtils.equality(stateRef, JsIntLiteral(indexOfGlobalCatch))
+        val isFromGlobalCatch = equality(stateRef, JsIntLiteral(indexOfGlobalCatch))
         val catch = JsCatch(functionWithBody.scope, "e")
         val continueWithException = JsBlock(
-                JsAstUtils.assignment(stateRef.deepCopy(), exceptionStateRef.deepCopy()).makeStmt(),
+                assignment(stateRef.deepCopy(), exceptionStateRef.deepCopy()).makeStmt(),
-                JsAstUtils.assignment(JsNameRef(context.metadata.exceptionName, JsThisRef()),
+                assignment(JsNameRef(context.metadata.exceptionName, JsThisRef()),
                                       catch.parameter.name.makeRef()).makeStmt()
         )
-        val adjustExceptionState = JsAstUtils.assignment(exceptionStateRef.deepCopy(), stateRef.deepCopy()).makeStmt()
+        val adjustExceptionState = assignment(exceptionStateRef.deepCopy(), stateRef.deepCopy()).makeStmt()
         catch.body = JsBlock(JsIf(
                 isFromGlobalCatch,
                 JsBlock(adjustExceptionState, JsThrow(catch.parameter.name.makeRef())),
@@ -230,10 +296,10 @@ class CoroutineFunctionTransformer(private val function: JsFunction, name: Strin
     }
 
     private fun MutableList<JsStatement>.assignToField(fieldName: JsName, value: JsExpression, psiElement: PsiElement?) {
-        this += JsAstUtils.assignment(JsNameRef(fieldName, JsThisRef()), value).source(psiElement).makeStmt()
+        this += assignment(JsNameRef(fieldName, JsThisRef()), value).source(psiElement).makeStmt()
     }
 
     private fun MutableList<JsStatement>.assignToPrototype(fieldName: JsName, value: JsExpression) {
-        this += JsAstUtils.assignment(JsNameRef(fieldName, JsAstUtils.prototypeOf(className.makeRef())), value).makeStmt()
+        this += assignment(JsNameRef(fieldName, prototypeOf(className.makeRef())), value).makeStmt()
     }
 }

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

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.js.backend.ast.metadata.*
 import org.jetbrains.kotlin.js.inline.util.collectFreeVariables
 import org.jetbrains.kotlin.js.inline.util.replaceNames
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
+import org.jetbrains.kotlin.js.translate.utils.JsAstUtils.pureFqn
 import org.jetbrains.kotlin.js.translate.utils.splitToRanges
 
 fun JsNode.collectNodesToSplit(breakContinueTargets: Map<JsContinue, JsStatement>): Set<JsNode> {
@@ -240,6 +241,31 @@ fun JsBlock.replaceSpecialReferences(context: CoroutineTransformationContext) {
     visitor.accept(this)
 }
 
+fun JsBlock.replaceSpecialReferencesInSimpleFunction(continuationParam: JsParameter, resultVar: JsName) {
+    val visitor = object : JsVisitorWithContextImpl() {
+        override fun visit(x: JsFunction, ctx: JsContext<*>) = false
+
+        override fun endVisit(x: JsNameRef, ctx: JsContext<in JsNode>) {
+            when {
+                x.coroutineReceiver -> {
+                    ctx.replaceMe(pureFqn(continuationParam.name, null).source(x.source))
+                }
+
+                x.coroutineController -> {
+                    ctx.replaceMe(JsThisRef().apply {
+                        source = x.source
+                    })
+                }
+
+                x.coroutineResult && x.qualifier.let { it is JsNameRef && it.name == continuationParam.name } -> {
+                    ctx.replaceMe(pureFqn(resultVar, null).source(x.source))
+                }
+            }
+        }
+    }
+    visitor.accept(this)
+}
+
 fun List<CoroutineBlock>.collectVariablesSurvivingBetweenBlocks(localVariables: Set<JsName>, parameters: Set<JsName>): Set<JsName> {
     val varDefinedIn = localVariables.associate { it to mutableSetOf<Int>() }
     val varDeclaredIn = localVariables.associate { it to mutableSetOf<Int>() }
@@ -378,3 +404,6 @@ fun JsBlock.replaceLocalVariables(context: CoroutineTransformationContext, local
     }
     visitor.accept(this)
 }
+
+internal fun JsExpression?.isStateMachineResult() =
+        this is JsNameRef && this.coroutineResult && qualifier.let { it is JsNameRef && it.coroutineReceiver && it.qualifier == null }

js/js.inliner/src/org/jetbrains/kotlin/js/inline/util/rewriters/ReturnReplacingVisitor.kt

@@ -75,7 +75,6 @@ class ReturnReplacingVisitor(
 
     private fun processCoroutineResult(expression: JsExpression?): JsExpression? {
         if (!isSuspend) return expression
-        if (expression != null && expression.isTailCallSuspend) return expression
         val lhs = JsNameRef("\$\$coroutineResult\$\$", JsAstUtils.stateMachineReceiver()).apply { coroutineResult = true }
         return JsAstUtils.assignment(lhs, expression ?: Namer.getUndefinedExpression())
     }

js/js.tests/test/org/jetbrains/kotlin/js/test/utils/DirectiveTestUtils.java

@@ -205,6 +205,8 @@ public class DirectiveTestUtils {
 
     private static final DirectiveHandler COUNT_NULLS = new CountNodesDirective<>("CHECK_NULLS_COUNT", JsNullLiteral.class);
 
+    private static final DirectiveHandler COUNT_NEW = new CountNodesDirective<>("CHECK_NEW_COUNT", JsNew.class);
+
     private static final DirectiveHandler COUNT_CASES = new CountNodesDirective<>("CHECK_CASES_COUNT", JsCase.class);
 
     private static final DirectiveHandler COUNT_IF = new CountNodesDirective<>("CHECK_IF_COUNT", JsIf.class);
@@ -343,6 +345,7 @@ public class DirectiveTestUtils {
             COUNT_VARS,
             COUNT_BREAKS,
             COUNT_NULLS,
+            COUNT_NEW,
             COUNT_CASES,
             COUNT_IF,
             COUNT_DEBUGGER,

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

@@ -100,6 +100,7 @@ class LiteralFunctionTranslator(context: TranslationContext) : AbstractTranslato
         if (!descriptor.isSuspend) return
 
         fillCoroutineMetadata(context, descriptor, hasController = descriptor.extensionReceiverParameter != null)
+        forceStateMachine = true
     }
 
     fun ValueParameterDescriptorImpl.WithDestructuringDeclaration.translate(context: TranslationContext): JsVars {