[K/N][IR] Implemented suspend tail call optimization

If for a suspend function all calls to other suspend functions a tail calls,
it is not needed to build the coroutine implementation.

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/TailSuspendCallsCollector.kt

@@ -0,0 +1,110 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.backend.common
+
+import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.symbols.IrReturnableBlockSymbol
+import org.jetbrains.kotlin.ir.types.isUnit
+import org.jetbrains.kotlin.ir.util.isSuspend
+import org.jetbrains.kotlin.ir.util.render
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
+
+data class TailSuspendCalls(val callSites: Set<IrCall>, val hasNotTailSuspendCalls: Boolean)
+
+/**
+ * Collects calls to be treated as tail calls: "last" expressions which are either direct return statement with a call
+ * to other suspend function, or, for a Unit-returning function, a call to other suspend function, also returning Unit.
+ */
+fun collectTailSuspendCalls(context: CommonBackendContext, irFunction: IrSimpleFunction): TailSuspendCalls {
+    require(irFunction.isSuspend) { "A suspend function expected: ${irFunction.render()}" }
+    val body = irFunction.body ?: return TailSuspendCalls(emptySet(), false)
+
+    class VisitorState(val insideTryBlock: Boolean, val isTailExpression: Boolean)
+
+    val isUnitReturn = irFunction.returnType.isUnit()
+    var hasNotTailSuspendCall = false
+    val tailSuspendCalls = mutableSetOf<IrCall>()
+    val tailReturnableBlocks = mutableSetOf<IrReturnableBlockSymbol>()
+
+    val visitor = object : IrElementVisitor<Unit, VisitorState> {
+        override fun visitElement(element: IrElement, data: VisitorState) {
+            element.acceptChildren(this, VisitorState(data.insideTryBlock, isTailExpression = false))
+        }
+
+        override fun visitTry(aTry: IrTry, data: VisitorState) {
+            aTry.tryResult.accept(this, VisitorState(insideTryBlock = true, isTailExpression = false))
+            aTry.catches.forEach { it.result.accept(this, data) }
+            require(aTry.finallyExpression == null) { "All finally clauses should've been lowered out" }
+        }
+
+        private fun isTailReturn(expression: IrReturn) =
+            expression.returnTargetSymbol == irFunction.symbol || expression.returnTargetSymbol in tailReturnableBlocks
+
+        override fun visitReturn(expression: IrReturn, data: VisitorState) {
+            expression.value.accept(this, VisitorState(data.insideTryBlock, isTailReturn(expression)))
+        }
+
+        override fun visitExpressionBody(body: IrExpressionBody, data: VisitorState) =
+            body.acceptChildren(this, data)
+
+        override fun visitBlockBody(body: IrBlockBody, data: VisitorState) =
+            visitStatementContainer(body, data)
+
+        override fun visitContainerExpression(expression: IrContainerExpression, data: VisitorState) {
+            if (expression is IrReturnableBlock && data.isTailExpression)
+                tailReturnableBlocks.add(expression.symbol)
+            visitStatementContainer(expression, data)
+        }
+
+        private fun visitStatementContainer(expression: IrStatementContainer, data: VisitorState) {
+            expression.statements.forEachIndexed { index, irStatement ->
+                val isTailStatement = if (index == expression.statements.lastIndex) {
+                    // The last statement defines the result of the container expression, so it has the same kind.
+                    // Note: this is even true for returnable blocks: if it is a Unit-returning block, this is exactly
+                    // like a usual block; if it is a non-Unit block, then it must end with an explicit return statement.
+                    data.isTailExpression
+                } else {
+                    // In a Unit-returning function, any statement directly followed by a `return` is a tail statement.
+                    isUnitReturn && expression.statements[index + 1].let {
+                        it is IrReturn && isTailReturn(it) && it.value.isUnitRead()
+                    }
+                }
+                irStatement.accept(this, VisitorState(data.insideTryBlock, isTailStatement))
+            }
+        }
+
+        override fun visitWhen(expression: IrWhen, data: VisitorState) {
+            expression.branches.forEach {
+                it.condition.accept(this, VisitorState(data.insideTryBlock, isTailExpression = false))
+                it.result.accept(this, data)
+            }
+        }
+
+        override fun visitCall(expression: IrCall, data: VisitorState) {
+            if (expression.isSuspend) {
+                if (!data.insideTryBlock && data.isTailExpression)
+                    tailSuspendCalls.add(expression)
+                else
+                    hasNotTailSuspendCall = true
+            }
+
+            // For a tail call, [returnIfSuspended] can be optimized away.
+            val isTailExpression = data.isTailExpression && expression.isReturnIfSuspendedCall()
+            expression.acceptChildren(this, VisitorState(data.insideTryBlock, isTailExpression))
+        }
+
+        private fun IrExpression.isUnitRead(): Boolean =
+            this is IrGetObjectValue && symbol == context.irBuiltIns.unitClass
+
+        private fun IrCall.isReturnIfSuspendedCall() =
+            symbol == context.ir.symbols.returnIfSuspended
+    }
+
+    body.accept(visitor, VisitorState(insideTryBlock = false, isTailExpression = true))
+    return TailSuspendCalls(tailSuspendCalls, hasNotTailSuspendCall)
+}

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/AbstractSuspendFunctionsLowering.kt

@@ -63,19 +63,16 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
 
             override fun visitClass(declaration: IrClass) {
                 declaration.acceptChildrenVoid(this)
-                declaration.transformDeclarationsFlat(::tryTransformSuspendFunction)
+                if (declaration.origin != DECLARATION_ORIGIN_COROUTINE_IMPL)
+                    declaration.transformDeclarationsFlat(::tryTransformSuspendFunction)
             }
         })
     }
 
-
-    // Suppress since it is used in native
-    @Suppress("MemberVisibilityCanBePrivate")
     protected fun IrCall.isReturnIfSuspendedCall() =
-        symbol.signature == context.ir.symbols.returnIfSuspended.signature
+        symbol == context.ir.symbols.returnIfSuspended
 
     private fun tryTransformSuspendFunction(element: IrElement) =
-
         if (element is IrSimpleFunction && element.isSuspend && element.modality != Modality.ABSTRACT)
             transformSuspendFunction(element, suspendLambdas[element])
         else null
@@ -123,103 +120,52 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
         })
     }
 
-    private sealed class SuspendFunctionKind {
+    private fun transformSuspendFunction(irFunction: IrSimpleFunction, functionReference: IrFunctionReference?): List<IrDeclaration>? {
-        object NO_SUSPEND_CALLS : SuspendFunctionKind()
+        val (tailSuspendCalls, hasNotTailSuspendCalls) = collectTailSuspendCalls(context, irFunction)
-        class DELEGATING(val delegatingCall: IrCall) : SuspendFunctionKind()
+        return when {
-        object NEEDS_STATE_MACHINE : SuspendFunctionKind()
+            irFunction in suspendLambdas -> {
+                // Suspend lambdas always need coroutine implementation.
+                // They are called through factory method <create>, thus we can eliminate original body.
+                listOf(buildCoroutine(irFunction, functionReference))
+            }
+
+            // TODO: May be optimize tail suspend calls even in case of a state machine?
+            hasNotTailSuspendCalls -> listOf<IrDeclaration>(buildCoroutine(irFunction, functionReference), irFunction)
+
+            else -> {
+                // Otherwise, no suspend calls at all or all of them are tail calls - no need in a state machine.
+                // Have to simplify them though (convert them to proper return statements).
+                simplifyTailSuspendCalls(irFunction, tailSuspendCalls)
+                null
+            }
+        }
     }
 
-    private fun transformSuspendFunction(irFunction: IrSimpleFunction, functionReference: IrFunctionReference?) =
+    private fun simplifyTailSuspendCalls(irFunction: IrSimpleFunction, tailSuspendCalls: Set<IrCall>) {
-        when (val suspendFunctionKind = getSuspendFunctionKind(irFunction)) {
+        if (tailSuspendCalls.isEmpty()) return
-            is SuspendFunctionKind.NO_SUSPEND_CALLS -> {
-                null                                                            // No suspend function calls - just an ordinary function.
-            }
 
-            is SuspendFunctionKind.DELEGATING -> {                              // Calls another suspend function at the end.
+        val irBuilder = context.createIrBuilder(irFunction.symbol)
-                removeReturnIfSuspendedCallAndSimplifyDelegatingCall(irFunction, suspendFunctionKind.delegatingCall)
+        irFunction.body!!.transformChildrenVoid(object : IrElementTransformerVoid() {
-                null                                                            // No need in state machine.
+            override fun visitCall(expression: IrCall): IrExpression {
-            }
+                val shortCut = if (expression.isReturnIfSuspendedCall())
+                    expression.getValueArgument(0)!!
+                else expression
 
-            is SuspendFunctionKind.NEEDS_STATE_MACHINE -> {
+                shortCut.transformChildrenVoid(this)
-                val coroutine = buildCoroutine(irFunction, functionReference)   // Coroutine implementation.
-                if (irFunction in suspendLambdas)             // Suspend lambdas are called through factory method <create>,
-                    listOf(coroutine)                                           // thus we can eliminate original body.
-                else
-                    listOf<IrDeclaration>(coroutine, irFunction)
-            }
-        }
 
-    private fun getSuspendFunctionKind(irFunction: IrSimpleFunction): SuspendFunctionKind {
+                return if (!expression.isSuspend)
-        if (irFunction in suspendLambdas)
+                    shortCut
-            return SuspendFunctionKind.NEEDS_STATE_MACHINE            // Suspend lambdas always need coroutine implementation.
+                else irBuilder.at(expression).irReturn(
-
+                    irBuilder.generateDelegatedCall(irFunction.returnType, shortCut)
-        val body = irFunction.body ?: return SuspendFunctionKind.NO_SUSPEND_CALLS
+                )
-
-        var numberOfSuspendCalls = 0
-        body.acceptVoid(object : IrElementVisitorVoid {
-            override fun visitElement(element: IrElement) {
-                element.acceptChildrenVoid(this)
-            }
-
-            override fun visitCall(expression: IrCall) {
-                expression.acceptChildrenVoid(this)
-                if (expression.isSuspend)
-                    ++numberOfSuspendCalls
             }
         })
-        // It is important to optimize the case where there is only one suspend call and it is the last statement
-        // because we don't need to build a fat coroutine class in that case.
-        // This happens a lot in practise because of suspend functions with default arguments.
-        // TODO: use TailRecursionCallsCollector.
-        val lastCall = when (val lastStatement = (body as IrBlockBody).statements.lastOrNull()) {
-            is IrCall -> lastStatement
-            is IrReturn -> {
-                var value: IrElement = lastStatement
-                /*
-                 * Check if matches this pattern:
-                 * block/return {
-                 *     block/return {
-                 *         .. suspendCall()
-                 *     }
-                 * }
-                 */
-                loop@ while (true) {
-                    value = when {
-                        value is IrBlock && value.statements.size == 1 -> value.statements.first()
-                        value is IrReturn -> value.value
-                        else -> break@loop
-                    }
-                }
-                value as? IrCall
-            }
-            else -> null
-        }
-        val suspendCallAtEnd = lastCall != null && lastCall.isSuspend    // Suspend call.
-        return when {
-            numberOfSuspendCalls == 0 -> SuspendFunctionKind.NO_SUSPEND_CALLS
-            numberOfSuspendCalls == 1
-                    && suspendCallAtEnd -> SuspendFunctionKind.DELEGATING(lastCall!!)
-            else -> SuspendFunctionKind.NEEDS_STATE_MACHINE
-        }
     }
 
     private val symbols = context.ir.symbols
     private val getContinuationSymbol = symbols.getContinuation
     private val continuationClassSymbol = getContinuationSymbol.owner.returnType.classifierOrFail as IrClassSymbol
 
-    private fun removeReturnIfSuspendedCallAndSimplifyDelegatingCall(irFunction: IrFunction, delegatingCall: IrCall) {
-        val returnValue =
-            if (delegatingCall.isReturnIfSuspendedCall())
-                delegatingCall.getValueArgument(0)!!
-            else delegatingCall
-        context.createIrBuilder(irFunction.symbol).run {
-            val statements = (irFunction.body as IrBlockBody).statements
-            val lastStatement = statements.last()
-            assert(lastStatement == delegatingCall || lastStatement is IrReturn) { "Unexpected statement $lastStatement" }
-            statements[statements.lastIndex] = irReturn(generateDelegatedCall(irFunction.returnType, returnValue))
-        }
-    }
-
     private fun buildCoroutine(irFunction: IrSimpleFunction, functionReference: IrFunctionReference?): IrClass {
         val coroutine = CoroutineBuilder(irFunction, functionReference).build()
         builtCoroutines[irFunction] = coroutine