[K/N][IR] Suspend tail call optimization for coroutines

Even if a coroutine class has to be built for a suspend function, we
still can perform tail call optimization for it.

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

@@ -35,7 +35,8 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
     protected abstract fun buildStateMachine(
         stateMachineFunction: IrFunction,
         transformingFunction: IrFunction,
-        argumentToPropertiesMap: Map<IrValueParameter, IrField>
+        argumentToPropertiesMap: Map<IrValueParameter, IrField>,
+        tailSuspendCalls: Set<IrCall>
     )
 
     protected abstract fun IrBlockBodyBuilder.generateCoroutineStart(invokeSuspendFunction: IrFunction, receiver: IrExpression)
@@ -126,11 +127,12 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
             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))
+                // Tail call optimization for suspend lambdas isn't yet supported by stdlib.
+                listOf(buildCoroutine(irFunction, functionReference, emptySet()))
             }
 
-            // TODO: May be optimize tail suspend calls even in case of a state machine?
+            hasNotTailSuspendCalls ->
-            hasNotTailSuspendCalls -> listOf<IrDeclaration>(buildCoroutine(irFunction, functionReference), irFunction)
+                listOf<IrDeclaration>(buildCoroutine(irFunction, functionReference, tailSuspendCalls), irFunction)
 
             else -> {
                 // Otherwise, no suspend calls at all or all of them are tail calls - no need in a state machine.
@@ -153,7 +155,7 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
 
                 shortCut.transformChildrenVoid(this)
 
-                return if (!expression.isSuspend)
+                return if (!expression.isSuspend || expression !in tailSuspendCalls)
                     shortCut
                 else irBuilder.at(expression).irReturn(
                     irBuilder.generateDelegatedCall(irFunction.returnType, shortCut)
@@ -166,8 +168,8 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
     private val getContinuationSymbol = symbols.getContinuation
     private val continuationClassSymbol = getContinuationSymbol.owner.returnType.classifierOrFail as IrClassSymbol
 
-    private fun buildCoroutine(irFunction: IrSimpleFunction, functionReference: IrFunctionReference?): IrClass {
+    private fun buildCoroutine(irFunction: IrSimpleFunction, functionReference: IrFunctionReference?, tailSuspendCalls: Set<IrCall>): IrClass {
-        val coroutine = CoroutineBuilder(irFunction, functionReference).build()
+        val coroutine = CoroutineBuilder(irFunction, functionReference, tailSuspendCalls).build()
         builtCoroutines[irFunction] = coroutine
 
         if (functionReference == null) {
@@ -204,7 +206,11 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
         val stateMachineFunction: IrFunction
     )
 
-    private inner class CoroutineBuilder(val irFunction: IrFunction, val functionReference: IrFunctionReference?) {
+    private inner class CoroutineBuilder(
+        val irFunction: IrFunction,
+        val functionReference: IrFunctionReference?,
+        val tailSuspendCalls: Set<IrCall>
+    ) {
         private val functionParameters = irFunction.explicitParameters
         private val boundFunctionParameters = functionReference?.getArgumentsWithIr()?.map { it.first }
         private val unboundFunctionParameters = boundFunctionParameters?.let { functionParameters - it }
@@ -251,6 +257,7 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
                 suspendFunctionClassTypeArguments = unboundParameterTypes + irFunction.returnType
                 superTypes += suspendFunctionClass.typeWith(suspendFunctionClassTypeArguments)
             }
+            coroutineClass.superTypes = superTypes
 
             val coroutineConstructor = buildConstructor()
 
@@ -285,7 +292,6 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
                 )
             }
 
-            coroutineClass.superTypes += superTypes
             coroutineClass.addFakeOverrides(
                 context.typeSystem,
                 ignoredParentSymbols = ignoredParentSymbols
@@ -515,7 +521,7 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
                 overriddenSymbols += stateMachineFunction.symbol
             }
 
-            buildStateMachine(function, irFunction, argumentToPropertiesMap)
+            buildStateMachine(function, irFunction, argumentToPropertiesMap, tailSuspendCalls)
             return function
         }
     }

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/ir/Ir.kt

@@ -351,6 +351,12 @@ internal class KonanSymbols(
     val invokeSuspendFunction =
             irBuiltIns.findBuiltInClassMemberFunctions(baseContinuationImpl, Name.identifier("invokeSuspend")).single()
 
+    val completionGetter = symbolTable.referenceSimpleFunction(
+            baseContinuationImpl.descriptor.unsubstitutedMemberScope
+                    .getContributedVariables(Name.identifier("completion"), NoLookupLocation.FROM_BACKEND).single()
+                    .getter!!
+    )
+
     override val coroutineSuspendedGetter = symbolTable.referenceSimpleFunction(
             coroutinesIntrinsicsPackage
                     .getContributedVariables(COROUTINE_SUSPENDED_NAME, NoLookupLocation.FROM_BACKEND)

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/NativeSuspendFunctionLowering.kt

@@ -2,6 +2,7 @@ package org.jetbrains.kotlin.backend.konan.lower
 
 import org.jetbrains.kotlin.backend.common.descriptors.synthesizedName
 import org.jetbrains.kotlin.backend.common.lower.*
+import org.jetbrains.kotlin.backend.common.lower.coroutines.getOrCreateFunctionWithContinuationStub
 import org.jetbrains.kotlin.backend.konan.Context
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
@@ -58,15 +59,20 @@ internal class NativeSuspendFunctionsLowering(ctx: Context) : AbstractSuspendFun
         )
     }
 
+    private val getContinuation = context.ir.symbols.getContinuation
+    private val completionGetter = context.ir.symbols.completionGetter
 
     override fun buildStateMachine(stateMachineFunction: IrFunction,
                                    transformingFunction: IrFunction,
-                                   argumentToPropertiesMap: Map<IrValueParameter, IrField>) {
+                                   argumentToPropertiesMap: Map<IrValueParameter, IrField>,
+                                   tailSuspendCalls: Set<IrCall>) {
         val originalBody = transformingFunction.body!!
         val resultArgument = stateMachineFunction.valueParameters.single()
 
         val coroutineClass = stateMachineFunction.parentAsClass
 
+        val thisReceiver = stateMachineFunction.dispatchReceiverParameter!!
+
         val labelField = coroutineClass.addField(Name.identifier("label"), symbols.nativePtrType, true)
 
         val startOffset = transformingFunction.startOffset
@@ -74,6 +80,71 @@ internal class NativeSuspendFunctionsLowering(ctx: Context) : AbstractSuspendFun
 
         val irBuilder = context.createIrBuilder(stateMachineFunction.symbol, startOffset, endOffset)
         stateMachineFunction.body = irBuilder.irBlockBody(startOffset, endOffset) {
+            if (tailSuspendCalls.isNotEmpty()) {
+                /*
+                 * Usual suspend call of, say, function foo will be transformed to something like this:
+                 * val result = foo(.., continuation = this)
+                 * if (result == COROUTINE_SUSPENDED) return COROUTINE_SUSPENDED
+                 * if (result.failed)
+                 *     this.completion.resumeWithException(result.exception)
+                 * else this.completion.resumeWith(result.value)
+                 *
+                 * But, if a call to foo is a tail call, we can replace the above with just:
+                 * return foo(.., continuation = this.completion)
+                 *
+                 * The visitor below does exactly this.
+                 */
+                originalBody.transformChildren(object : IrElementTransformer<Boolean> {
+                    fun IrBuilderWithScope.irGetCompletion() = irCall(completionGetter).apply { dispatchReceiver = irGet(thisReceiver) }
+
+                    override fun visitCall(expression: IrCall, /* substituteContinuation */ data: Boolean): IrExpression {
+                        val isTailSuspendCall = expression in tailSuspendCalls
+
+                        return when {
+                            expression.symbol == getContinuation -> {
+                                if (data)
+                                    irBuilder.at(expression).irGetCompletion()
+                                else
+                                    expression
+                            }
+
+                            expression.isReturnIfSuspendedCall() -> {
+                                if (!isTailSuspendCall)
+                                    expression
+                                else {
+                                    expression.transformChildren(this, /* substituteContinuation = */ true)
+
+                                    irBuilder.at(expression).irReturn(expression.getValueArgument(0)!!)
+                                }
+                            }
+
+                            else -> {
+                                // Only the top-most call should have its continuation substituted.
+                                val substituteContinuation = data && !expression.isSuspendCall
+                                expression.transformChildren(this, substituteContinuation)
+
+                                if (!isTailSuspendCall)
+                                    expression
+                                else {
+                                    val oldFun = expression.symbol.owner
+                                    val newFun: IrSimpleFunction = oldFun.getOrCreateFunctionWithContinuationStub(this@NativeSuspendFunctionsLowering.context)
+
+                                    irBuilder.at(expression).irReturn(
+                                            irCall(
+                                                    expression,
+                                                    newFun.symbol,
+                                                    newReturnType = newFun.returnType,
+                                                    newSuperQualifierSymbol = expression.superQualifierSymbol
+                                            ).also {
+                                                it.putValueArgument(it.valueArgumentsCount - 1, irGetCompletion())
+                                            }
+                                    )
+                                }
+                            }
+                        }
+                    }
+                }, /* substituteContinuation = */ false)
+            }
 
             val suspendResult = irVar("suspendResult".synthesizedName, context.irBuiltIns.anyNType, true)
 
@@ -100,8 +171,6 @@ internal class NativeSuspendFunctionsLowering(ctx: Context) : AbstractSuspendFun
 
             originalBody.transformChildrenVoid(object : IrElementTransformerVoid() {
 
-                private val thisReceiver = stateMachineFunction.dispatchReceiverParameter!!
-
                 // Replace returns to refer to the new function.
                 override fun visitReturn(expression: IrReturn): IrExpression {
                     expression.transformChildrenVoid(this)