From a07e21d9137f1f36ddbe8de9c7dea286c3f05786 Mon Sep 17 00:00:00 2001 From: pyos Date: Fri, 8 Apr 2022 12:31:07 +0200 Subject: [PATCH] JVM: streamline handling of tail-call suspend functions Just see if every suspend call is followed only by safe instructions or returns, then insert a suspension point check if there isn't a direct return. The first part of this is equivalent to the old implementation, just refactored. The second part generates strictly more checks; see, for example, the fixed test, in which the previous implementation failed to insert a check before a CHECKCAST. ^KT-51818 Fixed --- .../CoroutineTransformerMethodVisitor.kt | 11 +- .../coroutines/TailCallOptimization.kt | 263 +++++------------- .../FirBlackBoxCodegenTestGenerated.java | 6 + .../tailCallOptimizations/checkcast2.kt | 29 ++ .../codegen/BlackBoxCodegenTestGenerated.java | 6 + .../IrBlackBoxCodegenTestGenerated.java | 6 + .../LightAnalysisModeTestGenerated.java | 5 + 7 files changed, 124 insertions(+), 202 deletions(-) create mode 100644 compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast2.kt diff --git a/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/CoroutineTransformerMethodVisitor.kt b/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/CoroutineTransformerMethodVisitor.kt index 6a688a3d880..46917570834 100644 --- a/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/CoroutineTransformerMethodVisitor.kt +++ b/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/CoroutineTransformerMethodVisitor.kt @@ -96,15 +96,8 @@ class CoroutineTransformerMethodVisitor( addCompletionParameterToLVT(methodNode) } - val examiner = MethodNodeExaminer( - containingClassInternalName, - methodNode, - suspensionPoints, - disableTailCallOptimizationForFunctionReturningUnit - ) - if (examiner.allSuspensionPointsAreTailCalls(suspensionPoints)) { - examiner.replacePopsBeforeSafeUnitInstancesWithCoroutineSuspendedChecks() - examiner.addCoroutineSuspendedChecksBeforeSafeCheckcasts() + if (methodNode.allSuspensionPointsAreTailCalls(suspensionPoints, !disableTailCallOptimizationForFunctionReturningUnit)) { + methodNode.addCoroutineSuspendedChecks(suspensionPoints) dropSuspensionMarkers(methodNode) dropUnboxInlineClassMarkers(methodNode, suspensionPoints) return diff --git a/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/TailCallOptimization.kt b/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/TailCallOptimization.kt index e63dbee14f6..48168c038ba 100644 --- a/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/TailCallOptimization.kt +++ b/compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/TailCallOptimization.kt @@ -8,220 +8,90 @@ package org.jetbrains.kotlin.codegen.coroutines import org.jetbrains.kotlin.codegen.inline.isInlineMarker import org.jetbrains.kotlin.codegen.optimization.boxing.isUnitInstance import org.jetbrains.kotlin.codegen.optimization.common.ControlFlowGraph -import org.jetbrains.kotlin.codegen.optimization.common.asSequence import org.jetbrains.kotlin.codegen.optimization.common.isMeaningful import org.jetbrains.kotlin.codegen.optimization.fixStack.top import org.jetbrains.kotlin.codegen.optimization.transformer.MethodTransformer import org.jetbrains.kotlin.resolve.jvm.AsmTypes -import org.jetbrains.kotlin.utils.sure +import org.jetbrains.kotlin.utils.addToStdlib.popLast import org.jetbrains.org.objectweb.asm.Label import org.jetbrains.org.objectweb.asm.Opcodes import org.jetbrains.org.objectweb.asm.Type import org.jetbrains.org.objectweb.asm.tree.* import org.jetbrains.org.objectweb.asm.tree.analysis.BasicInterpreter import org.jetbrains.org.objectweb.asm.tree.analysis.BasicValue -import org.jetbrains.org.objectweb.asm.tree.analysis.Frame -internal class MethodNodeExaminer( - containingClassInternalName: String, - val methodNode: MethodNode, - suspensionPoints: List, - disableTailCallOptimizationForFunctionReturningUnit: Boolean -) { - private val frames: Array?> = - MethodTransformer.analyze(containingClassInternalName, methodNode, TcoInterpreter(suspensionPoints)) - private val controlFlowGraph = ControlFlowGraph.build(methodNode) +internal fun MethodNode.allSuspensionPointsAreTailCalls(suspensionPoints: List, optimizeReturnUnit: Boolean): Boolean { + val frames = MethodTransformer.analyze("fake", this, TcoInterpreter(suspensionPoints)) + val controlFlowGraph = ControlFlowGraph.build(this) - private val safeUnitInstances = mutableSetOf() - private val popsBeforeSafeUnitInstances = mutableSetOf() - private val areturnsAfterSafeUnitInstances = mutableSetOf() - private val meaningfulSuccessorsCache = hashMapOf>() + fun AbstractInsnNode.isSafe(): Boolean = + !isMeaningful || opcode in SAFE_OPCODES || isInvisibleInDebugVarInsn(this@allSuspensionPointsAreTailCalls) || isInlineMarker(this) - // CHECKCAST is considered safe if it is right before ARETURN and right after suspension point - // In this case, we can add check for COROUTINE_SUSPENDED, the same as we did for functions, returning Unit. - private val safeCheckcasts = mutableSetOf() - - init { - if (!disableTailCallOptimizationForFunctionReturningUnit) { - // retrieve all POP insns - val pops = methodNode.instructions.asSequence().filter { it.opcode == Opcodes.POP } - // for each of them check that all successors are PUSH Unit - val popsBeforeUnitInstances = pops.map { it to it.meaningfulSuccessors() } - .filter { (_, succs) -> succs.all { it.isUnitInstance() } } - .map { it.first }.toList() - for (pop in popsBeforeUnitInstances) { - val units = pop.meaningfulSuccessors() - val allUnitsAreSafe = units.all { unit -> - // check they have only returns among successors - unit.meaningfulSuccessors().all { it.opcode == Opcodes.ARETURN } + fun AbstractInsnNode.transitiveSuccessorsAreSafeOrReturns(): Boolean { + val visited = mutableSetOf(this) + val stack = mutableListOf(this) + while (stack.isNotEmpty()) { + val insn = stack.popLast() + // In Unit-returning functions, the last statement is followed by POP + GETSTATIC Unit.INSTANCE + // if it is itself not Unit-returning. + if (insn.opcode == Opcodes.ARETURN || (optimizeReturnUnit && insn.isPopBeforeReturnUnit)) { + if (frames[instructions.indexOf(insn)]?.top() !is FromSuspensionPointValue?) { + return false } - if (!allUnitsAreSafe) continue - // save them all to the properties - popsBeforeSafeUnitInstances += pop - safeUnitInstances += units - units.flatMapTo(areturnsAfterSafeUnitInstances) { it.meaningfulSuccessors() } - } - } - - fun AbstractInsnNode.isPartOfCheckcastChainBeforeAreturn(): Boolean { - for (succ in meaningfulSuccessors()) { - when (succ.opcode) { - Opcodes.CHECKCAST -> - if (!succ.isPartOfCheckcastChainBeforeAreturn()) return false - - Opcodes.ARETURN -> { - // do nothing + } else if (insn !== this && !insn.isSafe()) { + return false + } else { + for (nextIndex in controlFlowGraph.getSuccessorsIndices(insn)) { + val nextInsn = instructions.get(nextIndex) + if (visited.add(nextInsn)) { + stack.add(nextInsn) } - else -> return false - } - } - return true - } - - val checkcasts = methodNode.instructions.filter { it.opcode == Opcodes.CHECKCAST } - for (checkcast in checkcasts) { - if (!checkcast.isPartOfCheckcastChainBeforeAreturn()) continue - if (frames[methodNode.instructions.indexOf(checkcast)]?.top() !is FromSuspensionPointValue) continue - safeCheckcasts += checkcast - } - } - - // GETSTATIC kotlin/Unit.INSTANCE is considered safe iff - // it is part of POP, PUSH Unit, ARETURN sequence. - private fun AbstractInsnNode.isSafeUnitInstance(): Boolean = this in safeUnitInstances - - private fun AbstractInsnNode.isPopBeforeSafeUnitInstance(): Boolean = this in popsBeforeSafeUnitInstances - private fun AbstractInsnNode.isAreturnAfterSafeUnitInstance(): Boolean = this in areturnsAfterSafeUnitInstances - - private fun AbstractInsnNode.meaningfulSuccessors(): List = meaningfulSuccessorsCache.getOrPut(this) { - fun AbstractInsnNode.isMeaningful() = isMeaningful && opcode != Opcodes.NOP && opcode != Opcodes.GOTO && this !is LineNumberNode - - val visited = mutableSetOf() - fun dfs(insn: AbstractInsnNode) { - if (insn in visited) return - visited += insn - if (!insn.isMeaningful()) { - for (succIndex in controlFlowGraph.getSuccessorsIndices(insn)) { - dfs(methodNode.instructions[succIndex]) } } } - for (succIndex in controlFlowGraph.getSuccessorsIndices(this)) { - dfs(methodNode.instructions[succIndex]) - } - visited.filter { it.isMeaningful() } + return true } - fun replacePopsBeforeSafeUnitInstancesWithCoroutineSuspendedChecks() { - val isReferenceMap = - popsBeforeSafeUnitInstances.associateWith { (frames[methodNode.instructions.indexOf(it)]?.top()?.isReference == true) } - - for (pop in popsBeforeSafeUnitInstances) { - if (isReferenceMap[pop] == true) { - val label = Label() - methodNode.instructions.insertBefore(pop, withInstructionAdapter { - dup() - loadCoroutineSuspendedMarker() - ifacmpne(label) - areturn(AsmTypes.OBJECT_TYPE) - mark(label) - }) - } - } - } - - fun addCoroutineSuspendedChecksBeforeSafeCheckcasts() { - for (checkcast in safeCheckcasts) { - val label = Label() - methodNode.instructions.insertBefore(checkcast, withInstructionAdapter { - dup() - loadCoroutineSuspendedMarker() - ifacmpne(label) - areturn(AsmTypes.OBJECT_TYPE) - mark(label) - }) - } - } - - fun allSuspensionPointsAreTailCalls(suspensionPoints: List): Boolean { - val safelyReachableReturns = findSafelyReachableReturns() - - val instructions = methodNode.instructions - return suspensionPoints.all { suspensionPoint -> - val beginIndex = instructions.indexOf(suspensionPoint.suspensionCallBegin) - val endIndex = instructions.indexOf(suspensionPoint.suspensionCallEnd) - - val insideTryBlock = methodNode.tryCatchBlocks.any { block -> - val tryBlockStartIndex = instructions.indexOf(block.start) - val tryBlockEndIndex = instructions.indexOf(block.end) - - beginIndex in tryBlockStartIndex until tryBlockEndIndex - } - if (insideTryBlock) return@all false - - safelyReachableReturns[endIndex + 1]?.all { returnIndex -> - frames[returnIndex]?.top().sure { - "There must be some value on stack to return" - } is FromSuspensionPointValue - } ?: false - } - } - - /** - * Let's call an instruction safe if its execution is always invisible: stack modifications, branching, variable insns (invisible in debug) - * - * For some instruction `insn` define the result as following: - * - if there is a path leading to the non-safe instruction then result is `null` - * - Otherwise result contains all the reachable ARETURN indices - * - * @return indices of safely reachable returns for each instruction in the method node - */ - private fun findSafelyReachableReturns(): Array?> { - val insns = methodNode.instructions.toArray() - val reachableReturnsIndices = Array(insns.size) init@{ index -> - val insn = insns[index] - - if (insn.opcode == Opcodes.ARETURN && !insn.isAreturnAfterSafeUnitInstance()) { - return@init setOf(index) - } - - // Since POP, PUSH Unit, ARETURN behaves like normal return in terms of tail-call optimization, set return index to POP - if (insn.isPopBeforeSafeUnitInstance()) { - return@init setOf(index) - } - - if (!insn.isMeaningful || insn.opcode in SAFE_OPCODES || insn.isInvisibleInDebugVarInsn(methodNode) || isInlineMarker(insn) - || insn.isSafeUnitInstance() || insn.isAreturnAfterSafeUnitInstance() - ) { - setOf() - } else null - } - - var changed: Boolean - do { - changed = false - for (index in insns.indices.reversed()) { - if (insns[index].opcode == Opcodes.ARETURN) continue - - val newResult = - controlFlowGraph - .getSuccessorsIndices(index).plus(index) - .map(reachableReturnsIndices::get) - .fold?, Set?>(mutableSetOf()) { acc, successorsResult -> - if (acc != null && successorsResult != null) acc + successorsResult else null - } - - if (newResult != reachableReturnsIndices[index]) { - reachableReturnsIndices[index] = newResult - changed = true - } - } - } while (changed) - - return reachableReturnsIndices + return suspensionPoints.all { suspensionPoint -> + val index = instructions.indexOf(suspensionPoint.suspensionCallBegin) + tryCatchBlocks.all { index < instructions.indexOf(it.start) || instructions.indexOf(it.end) <= index } && + suspensionPoint.suspensionCallEnd.transitiveSuccessorsAreSafeOrReturns() } } +internal fun MethodNode.addCoroutineSuspendedChecks(suspensionPoints: List) { + for (suspensionPoint in suspensionPoints) { + if (suspensionPoint.suspensionCallEnd.nextMeaningful?.opcode == Opcodes.ARETURN) { + // `if (x == COROUTINE_SUSPENDED) return x else return x` == `return x` + continue + } + instructions.insert(suspensionPoint.suspensionCallEnd, withInstructionAdapter { + val label = Label() + dup() + loadCoroutineSuspendedMarker() + ifacmpne(label) + areturn(AsmTypes.OBJECT_TYPE) + mark(label) + }) + } +} + +private tailrec fun AbstractInsnNode?.skipUntilMeaningful(): AbstractInsnNode? = when { + this == null -> null + opcode == Opcodes.NOP || !isMeaningful -> next.skipUntilMeaningful() + opcode == Opcodes.GOTO -> (this as JumpInsnNode).label.skipUntilMeaningful() + else -> this +} + +private val AbstractInsnNode.nextMeaningful: AbstractInsnNode? + get() = next.skipUntilMeaningful() + +private val AbstractInsnNode.isReturnUnit: Boolean + get() = isUnitInstance() && nextMeaningful?.opcode == Opcodes.ARETURN + +private val AbstractInsnNode.isPopBeforeReturnUnit: Boolean + get() = opcode == Opcodes.POP && nextMeaningful?.isReturnUnit == true + private fun AbstractInsnNode?.isInvisibleInDebugVarInsn(methodNode: MethodNode): Boolean { val insns = methodNode.instructions val index = insns.indexOf(this) @@ -230,8 +100,15 @@ private fun AbstractInsnNode?.isInvisibleInDebugVarInsn(methodNode: MethodNode): }) } -private val SAFE_OPCODES = - ((Opcodes.DUP..Opcodes.DUP2_X2) + Opcodes.NOP + Opcodes.POP + Opcodes.POP2 + (Opcodes.IFEQ..Opcodes.GOTO)).toSet() + Opcodes.CHECKCAST +private val SAFE_OPCODES = buildSet { + add(Opcodes.NOP) + addAll(Opcodes.POP..Opcodes.SWAP) // POP*, DUP*, SWAP + addAll(Opcodes.IFEQ..Opcodes.GOTO) // IF*, GOTO + // CHECKCAST is technically not safe (can throw), but should be unless the type system is holey. + // Treating it as safe permits optimizing functions where a non-Any returning suspend function + // call is in a tail position (in bytecode they all return Object, so a cast is sometimes inserted). + add(Opcodes.CHECKCAST) +} private object FromSuspensionPointValue : BasicValue(AsmTypes.OBJECT_TYPE) { override fun equals(other: Any?): Boolean = other is FromSuspensionPointValue diff --git a/compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java b/compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java index dc40f00a8c2..80d151cee92 100644 --- a/compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java +++ b/compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java @@ -12770,6 +12770,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast.kt"); } + @Test + @TestMetadata("checkcast2.kt") + public void testCheckcast2() throws Exception { + runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast2.kt"); + } + @Test @TestMetadata("crossinline.kt") public void testCrossinline() throws Exception { diff --git a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast2.kt b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast2.kt new file mode 100644 index 00000000000..36fcd704569 --- /dev/null +++ b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast2.kt @@ -0,0 +1,29 @@ +// TARGET_BACKEND: JVM +// FULL_JDK +// WITH_STDLIB +// WITH_COROUTINES +// CHECK_TAIL_CALL_OPTIMIZATION +import helpers.* +import kotlin.coroutines.* +import kotlin.coroutines.intrinsics.* + +suspend fun suspendFun(x: String): String = suspendCoroutineUninterceptedOrReturn { + TailCallOptimizationChecker.saveStackTrace(it) + COROUTINE_SUSPENDED +} + +suspend fun myFunWithTailCall(x: String) { + x.let { suspendFun(it) } +} + +fun builder(c: suspend () -> Unit) { + c.startCoroutine(EmptyContinuation) +} + +fun box(): String { + builder { + myFunWithTailCall("...") + } + TailCallOptimizationChecker.checkNoStateMachineIn("myFunWithTailCall") + return "OK" +} diff --git a/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java b/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java index 339304d41be..2e0a753caff 100644 --- a/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java +++ b/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java @@ -12644,6 +12644,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest { runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast.kt"); } + @Test + @TestMetadata("checkcast2.kt") + public void testCheckcast2() throws Exception { + runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast2.kt"); + } + @Test @TestMetadata("crossinline.kt") public void testCrossinline() throws Exception { diff --git a/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java b/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java index 99b38f8a34a..2c09081e250 100644 --- a/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java +++ b/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java @@ -12770,6 +12770,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast.kt"); } + @Test + @TestMetadata("checkcast2.kt") + public void testCheckcast2() throws Exception { + runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast2.kt"); + } + @Test @TestMetadata("crossinline.kt") public void testCrossinline() throws Exception { diff --git a/compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java b/compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java index e3dce2b162f..fd4f20a2d5d 100644 --- a/compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java +++ b/compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java @@ -10203,6 +10203,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast.kt"); } + @TestMetadata("checkcast2.kt") + public void testCheckcast2() throws Exception { + runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast2.kt"); + } + @TestMetadata("crossinline.kt") public void testCrossinline() throws Exception { runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/crossinline.kt");