JS: add partial tail-call optimization for suspend functions

2017-11-01 13:10:40 +03:00
parent 7c621488ad
commit f8e7861ce6
9 changed files with 119 additions and 16 deletions
@@ -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.*

@@ -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.*

@@ -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.*
@@ -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)
@@ -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 assignConstructor = JsAstUtils.assignment(JsNameRef("constructor", prototype.deepCopy()), JsNameRef(constructorName))
+        val assignPrototype = assignment(prototype, baseClass)
+        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(),
-                JsAstUtils.assignment(JsNameRef(context.metadata.exceptionName, JsThisRef()),
+                assignment(stateRef.deepCopy(), exceptionStateRef.deepCopy()).makeStmt(),
+                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()
    }
 }
@@ -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>() }
@@ -377,4 +403,7 @@ 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 }
@@ -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())
    }
@@ -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,
@@ -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 {