azalea/kotlin-fork
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

JVM: streamline handling of tail-call suspend functions

Browse Source 
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
This commit is contained in:
pyos
2022-04-08 12:31:07 +02:00
committed by teamcity
parent 97fc97e63d
commit a07e21d913
7 changed files with 124 additions and 202 deletions
Download Patch File Download Diff File Expand all files Collapse all files
compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/CoroutineTransformerMethodVisitor.kt
+2 -9
View File
@@ -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
compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/TailCallOptimization.kt
+70 -193
View File
@@ -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<SuspensionPoint>,
disableTailCallOptimizationForFunctionReturningUnit: Boolean
) {
private val frames: Array<Frame<BasicValue>?> =
MethodTransformer.analyze(containingClassInternalName, methodNode, TcoInterpreter(suspensionPoints))
private val controlFlowGraph = ControlFlowGraph.build(methodNode)
internal fun MethodNode.allSuspensionPointsAreTailCalls(suspensionPoints: List<SuspensionPoint>, optimizeReturnUnit: Boolean): Boolean {
val frames = MethodTransformer.analyze("fake", this, TcoInterpreter(suspensionPoints))
val controlFlowGraph = ControlFlowGraph.build(this)
private val safeUnitInstances = mutableSetOf<AbstractInsnNode>()
private val popsBeforeSafeUnitInstances = mutableSetOf<AbstractInsnNode>()
private val areturnsAfterSafeUnitInstances = mutableSetOf<AbstractInsnNode>()
private val meaningfulSuccessorsCache = hashMapOf<AbstractInsnNode, List<AbstractInsnNode>>()
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<AbstractInsnNode>()
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<AbstractInsnNode> = meaningfulSuccessorsCache.getOrPut(this) {
fun AbstractInsnNode.isMeaningful() = isMeaningful && opcode != Opcodes.NOP && opcode != Opcodes.GOTO && this !is LineNumberNode
val visited = mutableSetOf<AbstractInsnNode>()
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<SuspensionPoint>): 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<Set<Int>?> {
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<Int>?, Set<Int>?>(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<SuspensionPoint>) {
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
compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java
Generated
+6
View File
@@ -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 {
compiler/testData/codegen/box/coroutines/tailCallOptimizations/checkcast2.kt
Vendored
+29
View File
@@ -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"
}
compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java
Generated
+6
View File
@@ -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 {
compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java
Generated
+6
View File
@@ -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 {
compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java
Generated
+5
View File
@@ -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");