[K/N] Moved coroutines var spilling lowering down the pipeline

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/driver/phases/NativeLoweringPhases.kt

@@ -457,7 +457,7 @@ private val coroutinesLivenessAnalysisPhase = createFileLoweringPhase(
         prerequisite = setOf(coroutinesPhase)
 )
 
-private val coroutinesVarSpillingPhase = createFileLoweringPhase(
+internal val CoroutinesVarSpillingPhase = createFileLoweringPhase(
         lowering = ::CoroutinesVarSpillingLowering,
         name = "CoroutinesVarSpilling",
         prerequisite = setOf(coroutinesPhase)
@@ -631,7 +631,6 @@ internal fun PhaseEngine<NativeGenerationState>.getLoweringsAfterInlining(): Low
         // Either of these could be turned off without losing correctness.
         coroutinesLivenessAnalysisPhase, // This is more optimal
         coroutinesLivenessAnalysisFallbackPhase, // While this is simple
-        coroutinesVarSpillingPhase,
         typeOperatorPhase,
         expressionBodyTransformPhase,
         objectClassesPhase,

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/driver/phases/TopLevelPhases.kt

@@ -423,9 +423,10 @@ private fun PhaseEngine<NativeGenerationState>.runCodegen(module: IrModuleFragme
     val dceResult = runPhase(DCEPhase, DCEInput(module, moduleDFG, devirtualizationAnalysisResults), disable = !optimize)
     runPhase(RemoveRedundantCallsToStaticInitializersPhase, RedundantCallsInput(moduleDFG, devirtualizationAnalysisResults, module), disable = !optimize)
     runPhase(DevirtualizationPhase, DevirtualizationInput(module, devirtualizationAnalysisResults), disable = !optimize)
-    // Have to run after link dependencies phase, because fields from dependencies can be changed during lowerings.
-    // Inline accessors only in optimized builds due to separate compilation and possibility to get broken debug information.
     module.files.forEach {
+        runPhase(CoroutinesVarSpillingPhase, it)
+        // Have to run after link dependencies phase, because fields from dependencies can be changed during lowerings.
+        // Inline accessors only in optimized builds due to separate compilation and possibility to get broken debug information.
         runPhase(PropertyAccessorInlinePhase, it, disable = !optimize)
         runPhase(InlineClassPropertyAccessorsPhase, it, disable = !optimize)
         runPhase(RedundantCoercionsCleaningPhase, it)

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

@@ -106,25 +106,18 @@ internal class CoroutinesVarSpillingLowering(val generationState: NativeGenerati
 /**
  * Computes visible variables at suspension points.
  */
-internal class CoroutinesLivenessAnalysisFallback(val generationState: NativeGenerationState) : FileLoweringPass, IrElementVisitorVoid {
+internal class CoroutinesLivenessAnalysisFallback(val generationState: NativeGenerationState) : BodyLoweringPass {
     private val invokeSuspendFunction = generationState.context.ir.symbols.invokeSuspendFunction
 
-    override fun lower(irFile: IrFile) {
-        if (generationState.liveVariablesAtSuspensionPoints.isEmpty())
-            irFile.acceptChildrenVoid(this)
-    }
+    override fun lower(irBody: IrBody, container: IrDeclaration) {
+        if (generationState.liveVariablesAtSuspensionPoints.isNotEmpty())
+            return
 
-    override fun visitElement(element: IrElement) {
-        element.acceptChildrenVoid(this)
-    }
+        val thisReceiver = (container as? IrSimpleFunction)?.dispatchReceiverParameter
+        if (thisReceiver == null || !container.overrides(invokeSuspendFunction.owner))
+            return
 
-    override fun visitFunction(declaration: IrFunction) {
-        val body = declaration.body
-        if (body != null && declaration.dispatchReceiverParameter != null
-                && (declaration as? IrSimpleFunction)?.overrides(invokeSuspendFunction.owner) == true
-        ) {
-            computeVisibleVariablesAtSuspensionPoints(body)
-        }
+        computeVisibleVariablesAtSuspensionPoints(irBody)
     }
 
     private fun computeVisibleVariablesAtSuspensionPoints(body: IrBody) {

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

@@ -181,7 +181,7 @@ internal class FunctionDFGBuilder(private val generationState: NativeGenerationS
 
         val function = FunctionDFGBuilder(expressionValuesExtractor, visitor.variableValues,
                 declaration, visitor.expressions, visitor.parentLoops, visitor.returnValues,
-                visitor.thrownValues, visitor.catchParameters).build()
+                visitor.thrownValues, visitor.catchParameters, visitor.liveVariables).build()
 
         context.logMultiple {
             +function.debugString()
@@ -198,9 +198,11 @@ internal class FunctionDFGBuilder(private val generationState: NativeGenerationS
         val returnValues = mutableListOf<IrExpression>()
         val thrownValues = mutableListOf<IrExpression>()
         val catchParameters = mutableSetOf<IrVariable>()
+        val liveVariables = mutableMapOf<IrCall, List<IrVariable>>()
 
         private val suspendableExpressionStack = mutableListOf<IrSuspendableExpression>()
         private val loopStack = mutableListOf<IrLoop>()
+        private val liveVariablesStack = mutableListOf<List<IrVariable>>()
         private val currentLoop get() = loopStack.peek()
 
         override fun visitElement(element: IrElement) {
@@ -249,6 +251,9 @@ internal class FunctionDFGBuilder(private val generationState: NativeGenerationS
 
                     expressions += jobInvocation to currentLoop
                 }
+                if (expression.symbol == saveCoroutineState)
+                    liveVariables[expression] = liveVariablesStack.peek()!!
+
                 val intrinsicType = tryGetIntrinsicType(expression)
                 if (intrinsicType == IntrinsicType.COMPARE_AND_SET || intrinsicType == IntrinsicType.COMPARE_AND_EXCHANGE) {
                     expressions += IrSetFieldImpl(
@@ -290,8 +295,10 @@ internal class FunctionDFGBuilder(private val generationState: NativeGenerationS
                 suspendableExpressionStack.push(expression)
                 suspendableExpressionValues.put(expression, mutableListOf())
             }
-            if (expression is IrSuspensionPoint)
+            if (expression is IrSuspensionPoint) {
                 suspendableExpressionValues[suspendableExpressionStack.peek()!!]!!.add(expression)
+                liveVariablesStack.push(generationState.liveVariablesAtSuspensionPoints[expression]!!)
+            }
             if (expression is IrLoop) {
                 parentLoops[expression] = currentLoop
                 loopStack.push(expression)
@@ -303,6 +310,8 @@ internal class FunctionDFGBuilder(private val generationState: NativeGenerationS
                 loopStack.pop()
             if (expression is IrSuspendableExpression)
                 suspendableExpressionStack.pop()
+            if (expression is IrSuspensionPoint)
+                liveVariablesStack.pop()
         }
 
         override fun visitSetField(expression: IrSetField) {
@@ -373,6 +382,7 @@ internal class FunctionDFGBuilder(private val generationState: NativeGenerationS
     }
 
     private val symbols = context.ir.symbols
+    private val unitType = context.irBuiltIns.unitType
     private val arrayGetSymbols = symbols.arrayGet.values
     private val arraySetSymbols = symbols.arraySet.values
     private val createUninitializedInstanceSymbol = symbols.createUninitializedInstance
@@ -383,6 +393,8 @@ internal class FunctionDFGBuilder(private val generationState: NativeGenerationS
     private val executeImplProducerInvoke = executeImplProducerClass.simpleFunctions()
             .single { it.name == OperatorNameConventions.INVOKE }
     private val reinterpret = symbols.reinterpret
+    private val saveCoroutineState = symbols.saveCoroutineState
+    private val restoreCoroutineState = symbols.restoreCoroutineState
     private val objCObjectRawValueGetter = symbols.interopObjCObjectRawValueGetter
 
     private class Scoped<out T : Any>(val value: T, val scope: DataFlowIR.Node.Scope)
@@ -396,6 +408,7 @@ internal class FunctionDFGBuilder(private val generationState: NativeGenerationS
             val returnValues: List<IrExpression>,
             val thrownValues: List<IrExpression>,
             val catchParameters: Set<IrVariable>,
+            val liveVariables: Map<IrCall, List<IrVariable>>,
     ) {
 
         private val rootScope = DataFlowIR.Node.Scope(0, emptyList())
@@ -622,6 +635,16 @@ internal class FunctionDFGBuilder(private val generationState: NativeGenerationS
 
                                 initInstanceSymbol -> error("Should've been lowered: ${value.render()}")
 
+                                saveCoroutineState -> DataFlowIR.Node.SaveCoroutineState(liveVariables[value]!!.map { variables[it]!!.value })
+
+                                restoreCoroutineState -> // This is a no-op for all analyses so far.
+                                    DataFlowIR.Node.StaticCall(
+                                            symbolTable.mapFunction(symbols.theUnitInstance.owner),
+                                            emptyList(),
+                                            symbolTable.mapType(unitType),
+                                            null
+                                    )
+
                                 else -> {
                                     /*
                                      * Resolve owner of the call with special handling of Any methods:

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

@@ -254,6 +254,8 @@ internal object DataFlowIR {
 
         class ArrayWrite(val callee: FunctionSymbol, val array: Edge, val index: Edge, val value: Edge, val type: Type) : Node()
 
+        class SaveCoroutineState(val liveVariables: List<Variable>) : Node()
+
         class Variable(values: List<Edge>, val type: Type, val kind: VariableKind) : Node() {
             val values = mutableListOf<Edge>().also { it += values }
         }
@@ -380,6 +382,13 @@ internal object DataFlowIR {
                     appendCastTo(node.value.castToType)
                 }
 
+                is Node.SaveCoroutineState -> buildString {
+                    appendLine("        SAVE COROUTINE STATE")
+                    appendList(node.liveVariables) {
+                        append("            VAL #${ids[it]!!}")
+                    }
+                }
+
                 is Node.Variable -> buildString {
                     append("       ${node.kind}")
                     appendList(node.values) {

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

@@ -1286,6 +1286,10 @@ internal object DevirtualizationAnalysis {
                             constraintGraph.voidNode
                         }
 
+                        is DataFlowIR.Node.SaveCoroutineState -> {
+                            constraintGraph.voidNode
+                        }
+
                         is DataFlowIR.Node.Variable ->
                             node.values.map { edgeToConstraintNode(it) }.let { values ->
                                 ordinaryNode { "TempVar\$${function.symbol}" }.also { node ->

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

@@ -857,6 +857,23 @@ internal object EscapeAnalysis {
 
                 traverseAndConvert(body.rootScope, Depths.ROOT_SCOPE - 1)
 
+                val dummyParameter = DataFlowIR.Node.Parameter(-1)
+                val parameters = Array(function.symbol.parameters.size) { dummyParameter } // Put dummy in order to not bother with nullability.
+                // Parameters are declared in the root scope
+                for (node in body.rootScope.nodes) {
+                    (node as? DataFlowIR.Node.Parameter)?.let {
+                        require(parameters[it.index] == dummyParameter) {
+                            "Two different parameters with the same index ${it.index} for $functionSymbol"
+                        }
+                        parameters[it.index] = it
+                    }
+                }
+                parameters.forEachIndexed { index, parameter ->
+                    require(parameter != dummyParameter) {
+                        "No parameter with index $index for $functionSymbol"
+                    }
+                }
+
                 val nothing = moduleDFG.symbolTable.mapClassReferenceType(context.ir.symbols.nothing.owner)
                 body.forEachNonScopeNode { node ->
                     when (node) {
@@ -900,6 +917,12 @@ internal object EscapeAnalysis {
                             readResult.addAssignmentEdge(array.getFieldNode(intestinesField, this))
                         }
 
+                        is DataFlowIR.Node.SaveCoroutineState -> {
+                            val thisIntestines = nodes[parameters[0]]!!.getFieldNode(intestinesField, this)
+                            for (variable in node.liveVariables)
+                                thisIntestines.addAssignmentEdge(nodes[variable]!!)
+                        }
+
                         is DataFlowIR.Node.Variable -> {
                             val variable = nodes[node]!!
                             for (value in node.values) {
@@ -925,9 +948,6 @@ internal object EscapeAnalysis {
 
                 val escapes = functionSymbol.escapes
                 if (escapes != null) {
-                    // Parameters are declared in the root scope
-                    val parameters = function.body.rootScope.nodes
-                            .filterIsInstance<DataFlowIR.Node.Parameter>()
                     for (parameter in parameters)
                         if (escapes.escapesAt(parameter.index))
                             escapeOrigins += nodes[parameter]!!