SuspendFunctionLowering: optimizations of ExpressionSlicer

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

@@ -927,6 +927,8 @@ internal class SuspendFunctionsLowering(val context: Context): DeclarationContai
 
                                 suspensionPoint.transformChildrenVoid(object : IrElementTransformerVoid() {
                                     override fun visitCall(expression: IrCall): IrExpression {
+                                        expression.transformChildrenVoid(this)
+
                                         val descriptor = expression.descriptor
                                         when (descriptor) {
                                             saveStateDescriptor -> {
@@ -1111,6 +1113,7 @@ internal class SuspendFunctionsLowering(val context: Context): DeclarationContai
 
                     val newChildren = arrayOfNulls<IrExpression?>(numberOfChildren)
                     val tempStatements = mutableListOf<IrStatement>()
+                    var first = true
                     for ((index, child) in children.withIndex()) {
                         if (child == null) continue
                         val transformedChild =
@@ -1121,19 +1124,51 @@ internal class SuspendFunctionsLowering(val context: Context): DeclarationContai
                                     tempStatements += statements.take(statements.size - 1)
                                     statements.last() as IrExpression
                                 }
-                        if ((!transformedChild.isPure() && hasSuspendCallInTail[index])
+                        if (first && !hasSuspendCallInTail[index + 1]) {
-                                || (expression.isSuspendCall && transformedChild.hasExternalVarSets())) {
+                            // Don't extract suspend call to a temporary if it is the first argument and is the only suspend call.
+                            newChildren[index] = transformedChild
+                            first = false
+                            continue
+                        }
+                        first = false
+                        if (!transformedChild.isPure() && hasSuspendCallInTail[index]) {
                             // Save to temporary in order to save execution order.
                             val tmp = IrTemporaryVariableDescriptorImpl(
                                     containingDeclaration = irFunction.descriptor,
-                                    name = "tmp${tempIndex++}".synthesizedName,
+                                    name                  = "tmp${tempIndex++}".synthesizedName,
-                                    outType = transformedChild.type)
+                                    outType               = transformedChild.type)
                             tempStatements += irVar(tmp, transformedChild)
                             newChildren[index] = irGet(tmp)
                         }
                         else newChildren[index] = transformedChild
                     }
 
+                    var calledSaveState = false
+                    if (expression.isSuspendCall) {
+                        val lastChild = newChildren.last()
+                        if (lastChild != null) {
+                            // Save state at late as possible.
+                            calledSaveState = true
+                            newChildren[numberOfChildren - 1] =
+                                    if (lastChild.isPure())
+                                        irBlock(lastChild) {
+                                            +irCall(saveStateDescriptor)
+                                            +lastChild
+                                        }
+                                    else {
+                                        val tmp = IrTemporaryVariableDescriptorImpl(
+                                                containingDeclaration = irFunction.descriptor,
+                                                name                  = "tmp${tempIndex++}".synthesizedName,
+                                                outType               = lastChild.type)
+                                        irBlock(lastChild) {
+                                            +irVar(tmp, lastChild)
+                                            +irCall(saveStateDescriptor)
+                                            +irGet(tmp)
+                                        }
+                                    }
+                        }
+                    }
+
                     when (expression) {
                         is IrSetField -> {
                             expression.receiver = newChildren[0]
@@ -1158,17 +1193,14 @@ internal class SuspendFunctionsLowering(val context: Context): DeclarationContai
                             containingDeclaration = irFunction.descriptor,
                             name                  = "suspensionPointId${suspensionPointIdIndex++}".synthesizedName,
                             outType               = suspensionPointIdType)
-                    val currentSuspendResult = IrTemporaryVariableDescriptorImpl(
-                            containingDeclaration = irFunction.descriptor,
-                            name                  = "currentSuspendResult${tempIndex++}".synthesizedName,
-                            outType               = context.builtIns.nullableAnyType)
                     val suspensionPoint = IrSuspensionPointImpl(
                             startOffset                = startOffset,
                             endOffset                  = endOffset,
                             type                       = context.builtIns.nullableAnyType,
                             suspensionPointIdParameter = irVar(suspensionPointIdParameter, null),
                             result                     = irBlock(startOffset, endOffset) {
-                                +irCall(saveStateDescriptor)
+                                if (!calledSaveState)
+                                    +irCall(saveStateDescriptor)
                                 +irSetVar(suspendResult, suspendCall)
                                 +irReturnIfSuspended(suspendResult)
                                 +irGet(suspendResult)
@@ -1178,12 +1210,14 @@ internal class SuspendFunctionsLowering(val context: Context): DeclarationContai
                                 +irThrowIfNotNull(exceptionArgument)
                                 +irGet(dataArgument)
                             })
-                    tempStatements.add(irVar(currentSuspendResult, suspensionPoint))
                     val expressionResult = when {
-                        suspendCall.type.isUnit() -> IrCompositeImpl(startOffset, endOffset, suspendCall.type)
+                        suspendCall.type.isUnit() -> irImplicitCoercionToUnit(suspensionPoint)
-                        else -> irCast(irGet(currentSuspendResult), suspendCall.type, suspendCall.type)
+                        else -> irCast(suspensionPoint, suspendCall.type, suspendCall.type)
+                    }
+                    return irBlock(expression) {
+                        tempStatements.forEach { +it }
+                        +expressionResult
                     }
-                    return irWrap(expressionResult, tempStatements)
                 }
 
             }
@@ -1211,6 +1245,7 @@ internal class SuspendFunctionsLowering(val context: Context): DeclarationContai
                     }
 
                     override fun visitCall(expression: IrCall) {
+                        expression.acceptChildrenVoid(this)
                         hasSuspendCalls = hasSuspendCalls || expression.isSuspendCall
                     }
                 })
@@ -1228,20 +1263,6 @@ internal class SuspendFunctionsLowering(val context: Context): DeclarationContai
                 }
             }
 
-            private fun IrElement.hasExternalVarSets(): Boolean {
-                var hasExternalVarSets = false
-                acceptVoid(object: VariablesScopeTracker() {
-
-                    override fun visitSetVariable(expression: IrSetVariable) {
-                        super.visitSetVariable(expression)
-                        if (scopeStack.all { !it.contains(expression.descriptor) })
-                            hasExternalVarSets = true
-                    }
-                })
-
-                return hasExternalVarSets
-            }
-
             private val IrExpression.isReturnIfSuspendedCall: Boolean
                 get() = this is IrCall && this.descriptor.original == returnIfSuspendedDescriptor
         }
@@ -1316,6 +1337,9 @@ internal class SuspendFunctionsLowering(val context: Context): DeclarationContai
     fun IrBuilderWithScope.irCast(arg: IrExpression, type: KotlinType, typeOperand: KotlinType) =
             IrTypeOperatorCallImpl(startOffset, endOffset, type, IrTypeOperator.CAST, typeOperand, arg)
 
+    fun IrBuilderWithScope.irImplicitCoercionToUnit(arg: IrExpression) =
+            IrTypeOperatorCallImpl(startOffset, endOffset, context.builtIns.unitType, IrTypeOperator.IMPLICIT_COERCION_TO_UNIT, context.builtIns.unitType, arg)
+
     fun IrBuilderWithScope.irGetField(receiver: IrExpression, descriptor: PropertyDescriptor) =
             IrGetFieldImpl(startOffset, endOffset, descriptor, receiver)