From f8e7861ce6d49712b25b7fd36a03e2754dddba00 Mon Sep 17 00:00:00 2001 From: Alexey Andreev Date: Wed, 1 Nov 2017 13:10:40 +0300 Subject: [PATCH] JS: add partial tail-call optimization for suspend functions --- .../inlineWithStateMachine.kt | 2 + .../inlineWithoutStateMachine.kt | 2 + .../tailCallOptimizations/simple.kt | 1 + .../ast/metadata/metadataProperties.kt | 4 +- .../coroutine/CoroutineFunctionTransformer.kt | 90 ++++++++++++++++--- .../kotlin/js/coroutine/CoroutinePasses.kt | 31 ++++++- .../util/rewriters/ReturnReplacingVisitor.kt | 1 - .../js/test/utils/DirectiveTestUtils.java | 3 + .../expression/LiteralFunctionTranslator.kt | 1 + 9 files changed, 119 insertions(+), 16 deletions(-) diff --git a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithStateMachine.kt b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithStateMachine.kt index f8cc94527ab..daa5b93de44 100644 --- a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithStateMachine.kt +++ b/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.* diff --git a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithoutStateMachine.kt b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithoutStateMachine.kt index 20416637b3d..809d3f3a905 100644 --- a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithoutStateMachine.kt +++ b/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.* diff --git a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/simple.kt b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/simple.kt index 353f5f3f5a4..c073ec293aa 100644 --- a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/simple.kt +++ b/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.* diff --git a/js/js.ast/src/org/jetbrains/kotlin/js/backend/ast/metadata/metadataProperties.kt b/js/js.ast/src/org/jetbrains/kotlin/js/backend/ast/metadata/metadataProperties.kt index 71d4f4748c4..39e5cb6e558 100644 --- a/js/js.ast/src/org/jetbrains/kotlin/js/backend/ast/metadata/metadataProperties.kt +++ b/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) diff --git a/js/js.inliner/src/org/jetbrains/kotlin/js/coroutine/CoroutineFunctionTransformer.kt b/js/js.inliner/src/org/jetbrains/kotlin/js/coroutine/CoroutineFunctionTransformer.kt index be114ba7475..2b236d26a5a 100644 --- a/js/js.inliner/src/org/jetbrains/kotlin/js/coroutine/CoroutineFunctionTransformer.kt +++ b/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 { + 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() + 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) { + 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, @@ -108,11 +174,11 @@ class CoroutineFunctionTransformer(private val function: JsFunction, name: Strin } private fun generateCoroutinePrototype(constructorName: JsName): List { - 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.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.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() } } \ No newline at end of file diff --git a/js/js.inliner/src/org/jetbrains/kotlin/js/coroutine/CoroutinePasses.kt b/js/js.inliner/src/org/jetbrains/kotlin/js/coroutine/CoroutinePasses.kt index 8fbba77946e..2862daa2159 100644 --- a/js/js.inliner/src/org/jetbrains/kotlin/js/coroutine/CoroutinePasses.kt +++ b/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): Set { @@ -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) { + 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.collectVariablesSurvivingBetweenBlocks(localVariables: Set, parameters: Set): Set { val varDefinedIn = localVariables.associate { it to mutableSetOf() } val varDeclaredIn = localVariables.associate { it to mutableSetOf() } @@ -377,4 +403,7 @@ fun JsBlock.replaceLocalVariables(context: CoroutineTransformationContext, local } } visitor.accept(this) -} \ No newline at end of file +} + +internal fun JsExpression?.isStateMachineResult() = + this is JsNameRef && this.coroutineResult && qualifier.let { it is JsNameRef && it.coroutineReceiver && it.qualifier == null } \ No newline at end of file diff --git a/js/js.inliner/src/org/jetbrains/kotlin/js/inline/util/rewriters/ReturnReplacingVisitor.kt b/js/js.inliner/src/org/jetbrains/kotlin/js/inline/util/rewriters/ReturnReplacingVisitor.kt index 548a6fe9b54..b34e7fe59b8 100644 --- a/js/js.inliner/src/org/jetbrains/kotlin/js/inline/util/rewriters/ReturnReplacingVisitor.kt +++ b/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()) } diff --git a/js/js.tests/test/org/jetbrains/kotlin/js/test/utils/DirectiveTestUtils.java b/js/js.tests/test/org/jetbrains/kotlin/js/test/utils/DirectiveTestUtils.java index 11314b12193..2251a1a87f0 100644 --- a/js/js.tests/test/org/jetbrains/kotlin/js/test/utils/DirectiveTestUtils.java +++ b/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, diff --git a/js/js.translator/src/org/jetbrains/kotlin/js/translate/expression/LiteralFunctionTranslator.kt b/js/js.translator/src/org/jetbrains/kotlin/js/translate/expression/LiteralFunctionTranslator.kt index 2a7e4797644..ae980d055ff 100644 --- a/js/js.translator/src/org/jetbrains/kotlin/js/translate/expression/LiteralFunctionTranslator.kt +++ b/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 {