JVM_IR: improve suspend tail call detection.

* TailCallOptimizationLowering should go into local classes in order to transform their suspend methods; * the check for invokes of noinline lambda arguments in codegen was incorrect, as it also returned true for calls of lambdas stored in local variables; * IrInlineCodegen should mark non-inlinable arguments used as inline suspend parameters; * detection of suspend/inline call sites was incorrect (or maybe it's the `compilationContextDescriptor` that was incorrect?..)
2020-02-19 14:23:51 +01:00
parent 3cf71c1d2b
commit eff02b6e72
9 changed files with 82 additions and 180 deletions
@@ -117,7 +117,6 @@ class JvmBackendContext(
    val continuationClassBuilders = mutableMapOf<IrSimpleFunction, ClassBuilder>()
    val suspendFunctionOriginalToView = mutableMapOf<IrFunction, IrFunction>()
    val suspendFunctionOriginalToStub = mutableMapOf<IrFunction, IrFunction>()
-    val suspendTailCallsWithUnitReplacement = mutableSetOf<IrAttributeContainer>()
    val fakeContinuation: IrExpression = createFakeContinuation(this)

    val staticDefaultStubs = mutableMapOf<IrFunctionSymbol, IrFunction>()
@@ -476,28 +476,22 @@ class ExpressionCodegen(
        }
    }

-    private fun IrFunctionAccessExpression.isSuspensionPoint(): Boolean {
-        val owner = symbol.owner
-        return when {
-            // Only suspend functions are suspension points
-            !owner.isSuspend -> false
-            // But not inside continuation classes and bridges
-            irFunction.shouldNotContainSuspendMarkers() -> false
-            // Noinline function are always suspension points
-            !owner.isInline -> true
-            // The suspend intrinsics are, albeit inline, also suspension points
-            owner.fqNameForIrSerialization == FqName("kotlin.coroutines.intrinsics.IntrinsicsKt.suspendCoroutineUninterceptedOrReturn") -> true
-            // Inside $$forInline functions crossinline calls are (usually) not suspension points, otherwise, flow will be pessimized
-            (dispatchReceiver as? IrGetField)?.symbol?.owner?.origin ==
-                    LocalDeclarationsLowering.DECLARATION_ORIGIN_FIELD_FOR_CROSSINLINE_CAPTURED_VALUE ->
-                irFunction.origin != FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
-            // The same goes for inline lambdas
-            (dispatchReceiver as? IrGetValue)?.let {
-                it.origin == IrStatementOrigin.VARIABLE_AS_FUNCTION && (it.symbol.owner as? IrValueParameter)?.isNoinline != true
-            } == true -> irFunction.origin != FOR_INLINE_STATE_MACHINE_TEMPLATE
-            // Otherwise, this is normal inline call, ergo not a suspension point
-            else -> false
-        }
+    private fun IrFunctionAccessExpression.isInvokeOfInlineLambda(): Boolean = when (val receiver = dispatchReceiver) {
+        is IrGetField -> receiver.symbol.owner.origin == LocalDeclarationsLowering.DECLARATION_ORIGIN_FIELD_FOR_CROSSINLINE_CAPTURED_VALUE
+        is IrGetValue -> receiver.origin == IrStatementOrigin.VARIABLE_AS_FUNCTION &&
+                (receiver.symbol.owner as? IrValueParameter)?.isNoinline == false
+        else -> false
+    }
+
+    private fun IrFunctionAccessExpression.isSuspensionPoint(): Boolean = when {
+        !symbol.owner.isSuspend || irFunction.shouldNotContainSuspendMarkers() -> false
+        // Copy-pasted bytecode blocks are not suspension points.
+        symbol.owner.isInline ->
+            symbol.owner.fqNameForIrSerialization == FqName("kotlin.coroutines.intrinsics.IntrinsicsKt.suspendCoroutineUninterceptedOrReturn")
+        // This includes inline lambdas, but only in functions intended for the inliner; in others, they stay as `f.invoke()`.
+        isInvokeOfInlineLambda() ->
+            irFunction.origin != FOR_INLINE_STATE_MACHINE_TEMPLATE && irFunction.origin != FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
+        else -> true
    }

    override fun visitVariable(declaration: IrVariable, data: BlockInfo): PromisedValue {
@@ -661,13 +655,6 @@ class ExpressionCodegen(
        val returnType = if (owner == irFunction) signature.returnType else methodSignatureMapper.mapReturnType(owner)
        val afterReturnLabel = Label()
        expression.value.accept(this, data).coerce(returnType, owner.returnType).materialize()
-        // We replaced COERTION_TO_UNIT with IrReturn during TailCallOptimizationLowering.
-        // Generate POP GETSTATIC kotlin/Unit.INSTANCE now.
-        // Otherwise, tail-call optimization will not work. See tailSuspendUnitFun.kt test.
-        if (expression in context.suspendTailCallsWithUnitReplacement) {
-            mv.pop()
-            mv.getstatic("kotlin/Unit", "INSTANCE", "Lkotlin/Unit;")
-        }
        generateFinallyBlocksIfNeeded(returnType, afterReturnLabel, data)
        expression.markLineNumber(startOffset = true)
        if (isNonLocalReturn) {
@@ -22,11 +22,7 @@ import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.types.classOrNull
 import org.jetbrains.kotlin.ir.types.toKotlinType
-import org.jetbrains.kotlin.ir.util.dump
-import org.jetbrains.kotlin.ir.util.getArgumentsWithIr
-import org.jetbrains.kotlin.ir.util.isSuspend
-import org.jetbrains.kotlin.resolve.BindingContext
-import org.jetbrains.kotlin.ir.util.render
+import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature
 import org.jetbrains.org.objectweb.asm.Type
@@ -87,11 +83,9 @@ class IrInlineCodegen(
            super.genValueAndPut(irValueParameter, argumentExpression, parameterType, codegen, blockInfo)
        }

-        if (irValueParameter.isInlineParameter(
-                /*after transformation inlinable lambda parameter with default value would have nullable type: check default value type first*/
-                irValueParameter.defaultValue?.expression?.type ?: irValueParameter.type
-            ) && isInlineIrExpression(argumentExpression)
-        ) {
+        // after transformation inlinable lambda parameter with default value would have nullable type: check default value type first
+        val isInlineParameter = irValueParameter.isInlineParameter(irValueParameter.defaultValue?.expression?.type ?: irValueParameter.type)
+        if (isInlineParameter && isInlineIrExpression(argumentExpression)) {
            val irReference: IrFunctionReference =
                (argumentExpression as IrBlock).statements.filterIsInstance<IrFunctionReference>().single()
            val boundReceiver = argumentExpression.statements.filterIsInstance<IrVariable>().singleOrNull()
@@ -107,9 +101,14 @@ class IrInlineCodegen(
            val kind = when (irValueParameter.origin) {
                IrDeclarationOrigin.MASK_FOR_DEFAULT_FUNCTION -> ValueKind.DEFAULT_MASK
                IrDeclarationOrigin.METHOD_HANDLER_IN_DEFAULT_FUNCTION -> ValueKind.METHOD_HANDLE_IN_DEFAULT
-                else -> if (argumentExpression is IrContainerExpression && argumentExpression.origin == IrStatementOrigin.DEFAULT_VALUE)
-                    ValueKind.DEFAULT_PARAMETER
-                else ValueKind.CAPTURED
+                else -> when {
+                    argumentExpression is IrContainerExpression && argumentExpression.origin == IrStatementOrigin.DEFAULT_VALUE ->
+                        ValueKind.DEFAULT_PARAMETER
+                    // TODO ValueKind.NON_INLINEABLE_ARGUMENT_FOR_INLINE_PARAMETER_CALLED_IN_SUSPEND?
+                    isInlineParameter && irValueParameter.type.isSuspendFunctionTypeOrSubtype() ->
+                        ValueKind.NON_INLINEABLE_ARGUMENT_FOR_INLINE_SUSPEND_PARAMETER
+                    else -> ValueKind.GENERAL
+                }
            }

            val onStack = when {
@@ -130,11 +129,7 @@ class IrInlineCodegen(


            //TODO support default argument erasure
-            if (!processDefaultMaskOrMethodHandler(
-                    onStack,
-                    kind
-                )
-            ) {
+            if (!processDefaultMaskOrMethodHandler(onStack, kind)) {
                val expectedType = JvmKotlinType(parameterType, irValueParameter.type.toKotlinType())
                putArgumentOrCapturedToLocalVal(expectedType, onStack, -1, irValueParameter.index, kind)
            }
@@ -89,8 +89,8 @@ class IrSourceCompilerForInline(
                root.classCodegen.type.internalName,
                root.signature.asmMethod.name,
                root.signature.asmMethod.descriptor,
-                compilationContextFunctionDescriptor.isInlineOrInsideInline(),
-                compilationContextFunctionDescriptor.isSuspend,
+                root.irFunction.descriptor.isInlineOrInsideInline(),
+                root.irFunction.isSuspend,
                findElement()?.let { CodegenUtil.getLineNumberForElement(it, false) } ?: 0
            )
        }
@@ -11,21 +11,18 @@ import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.codegen.anyOfOverriddenFunctionsReturnsNonUnit
 import org.jetbrains.kotlin.backend.jvm.codegen.isKnownToBeTailCall
-import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
 import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.IrReturnImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrTypeOperatorCallImpl
 import org.jetbrains.kotlin.ir.types.isUnit
 import org.jetbrains.kotlin.ir.util.fqNameWhenAvailable
 import org.jetbrains.kotlin.ir.util.isSuspend
 import org.jetbrains.kotlin.ir.util.parentAsClass
-import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
-import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
-import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
-import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformer
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name

@@ -35,72 +32,60 @@ internal val tailCallOptimizationPhase = makeIrFilePhase(
    "Add or move returns to suspension points on tail-call positions"
 )

-// TODO: Make this lowering common
-private class TailCallOptimizationLowering(private val context: JvmBackendContext) : IrElementVisitorVoid, FileLoweringPass {
+// Find all tail-calls inside suspend function. We should add IrReturn before them, so the codegen will generate
+// code which is understandable by old BE's tail-call optimizer.
+private class TailCallOptimizationLowering(private val context: JvmBackendContext) : FileLoweringPass {
    override fun lower(irFile: IrFile) {
-        visitFile(irFile)
-    }
+        irFile.transformChildren(object : IrElementTransformer<TailCallOptimizationData?> {
+            override fun visitFunction(declaration: IrFunction, data: TailCallOptimizationData?) =
+                super.visitFunction(declaration, TailCallOptimizationData(declaration))

-    override fun visitElement(element: IrElement) {
-        element.acceptChildrenVoid(this)
-    }
-
-    override fun visitFunction(function: IrFunction) {
-        if (!function.isSuspend || function.origin == IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA || function.isKnownToBeTailCall()) return
-        // Disable tail-call optimization for functions, returning Unit and overriding function, returning non-unit.
-        if (function.returnType.isUnit() && function.anyOfOverriddenFunctionsReturnsNonUnit()) return
-
-        val tailCalls = mutableSetOf<IrCall>()
-
-        // Find all tail-calls inside suspend function.
-        // We should add IrReturn before them, so the codegen will generate code, which is understandable by old BE's tail-call optimizer.
-        // This function collect all possible tail-calls, if they are not prepended by IrReturns, so we can prepend them.
-        // If the call has IrReturn before it, we do not need to add another one.
-        // If the call is inside a branch of `when` statement (or `if`, which is lowered to `when`), go through the branches and
-        // find tail-calls there.
-        fun findCallsOnTailPositionWithoutImmediateReturn(statement: IrStatement, immediateReturn: Boolean = false) {
-            when (statement) {
-                is IrCall -> if (statement.isSuspend && !immediateReturn) {
-                    tailCalls += statement
-                }
-                is IrBlock -> findCallsOnTailPositionWithoutImmediateReturn(
-                    statement.statements.findTailCall(function.returnType.isUnit()) ?: return
+            override fun visitCall(expression: IrCall, data: TailCallOptimizationData?): IrExpression {
+                val transformed = super.visitCall(expression, data) as IrExpression
+                return if (data == null || expression !in data.tailCalls) transformed else IrReturnImpl(
+                    expression.startOffset, expression.endOffset, context.irBuiltIns.nothingType, data.function.symbol,
+                    if (data.returnsUnit) transformed.coerceToUnit() else transformed
                )
-                is IrWhen -> for (branch in statement.branches) {
-                    findCallsOnTailPositionWithoutImmediateReturn(branch.result)
-                }
-                is IrReturn -> findCallsOnTailPositionWithoutImmediateReturn(statement.value, immediateReturn = true)
-                is IrTypeOperatorCall -> if (statement.operator == IrTypeOperator.IMPLICIT_COERCION_TO_UNIT) {
-                    findCallsOnTailPositionWithoutImmediateReturn(statement.argument)
-                }
-                else -> {
-                    // Not a call, or something containing a tail-call, break
-                    // TODO: Support binary logical operations and elvis, though. KT-23826 and KT-23825
-                }
+            }
+        }, null)
+    }
+
+    private fun IrExpression.coerceToUnit() = IrTypeOperatorCallImpl(
+        startOffset, endOffset, context.irBuiltIns.unitType, IrTypeOperator.IMPLICIT_COERCION_TO_UNIT, context.irBuiltIns.unitType, this
+    )
+}
+
+private class TailCallOptimizationData(val function: IrFunction) {
+    val returnsUnit = function.returnType.isUnit()
+    val tailCalls = mutableSetOf<IrCall>()
+
+    // Collect all tail calls, including those nested in `when`s, which are not arguments to `return`s.
+    private fun findCallsOnTailPositionWithoutImmediateReturn(statement: IrStatement, immediateReturn: Boolean = false) {
+        when {
+            statement is IrCall && statement.isSuspend && !immediateReturn && (returnsUnit || statement.type == function.returnType) ->
+                tailCalls += statement
+            statement is IrBlock ->
+                statement.statements.findTailCall(returnsUnit)?.let(::findCallsOnTailPositionWithoutImmediateReturn)
+            statement is IrWhen ->
+                statement.branches.forEach { findCallsOnTailPositionWithoutImmediateReturn(it.result) }
+            statement is IrReturn ->
+                findCallsOnTailPositionWithoutImmediateReturn(statement.value, immediateReturn = true)
+            statement is IrTypeOperatorCall && statement.operator == IrTypeOperator.IMPLICIT_COERCION_TO_UNIT ->
+                findCallsOnTailPositionWithoutImmediateReturn(statement.argument)
+            // TODO: Support binary logical operations and elvis, though. KT-23826 and KT-23825
+        }
+    }
+
+    init {
+        if (function.isSuspend && function.origin != IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA && !function.isKnownToBeTailCall() &&
+            // See `disableTailCallOptimizationForFunctionReturningUnit` in `generateStateMachineForNamedFunction`:
+            !(returnsUnit && function.anyOfOverriddenFunctionsReturnsNonUnit())
+        ) {
+            when (val body = function.body) {
+                is IrBlockBody -> body.statements.findTailCall(returnsUnit)?.let(::findCallsOnTailPositionWithoutImmediateReturn)
+                is IrExpressionBody -> findCallsOnTailPositionWithoutImmediateReturn(body.expression)
            }
        }
-
-        when (val body = function.body) {
-            null -> return
-            is IrBlockBody -> findCallsOnTailPositionWithoutImmediateReturn(
-                body.statements.findTailCall(function.returnType.isUnit()) ?: return
-            )
-            is IrExpressionBody -> findCallsOnTailPositionWithoutImmediateReturn(body.expression)
-            else -> error("Unexpected $body")
-        }
-
-        function.transformChildrenVoid(object : IrElementTransformerVoid() {
-            override fun visitCall(call: IrCall): IrExpression {
-                if (call !in tailCalls) return call
-                if (!function.returnType.isUnit() && call.type != function.returnType) return call
-                // Replace ARETURN with { POP, GETSTATIC kotlin/Unit.INSTANCE, ARETURN } during codegen later.
-                // Otherwise, additional CHECKCAST will break tail-call optimization
-                if (function.returnType.isUnit()) {
-                    context.suspendTailCallsWithUnitReplacement += call.attributeOwnerId
-                }
-                return IrReturnImpl(call.startOffset, call.endOffset, context.irBuiltIns.nothingType, function.symbol, call)
-            }
-        })
    }
 }