[K/N] Fallback in case turned off liveness analysis

#KT-65153 Fixed

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

@@ -18,6 +18,7 @@ import org.jetbrains.kotlin.backend.konan.serialization.SerializedClassFields
 import org.jetbrains.kotlin.backend.konan.serialization.SerializedEagerInitializedFile
 import org.jetbrains.kotlin.backend.konan.serialization.SerializedInlineFunctionReference
 import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.expressions.IrSuspensionPoint
 
 internal class InlineFunctionOriginInfo(val irFunction: IrFunction, val irFile: IrFile, val startOffset: Int, val endOffset: Int)
 
@@ -73,6 +74,8 @@ internal class NativeGenerationState(
     val calledFromExportedInlineFunctions = mutableSetOf<IrFunction>()
     val constructedFromExportedInlineFunctions = mutableSetOf<IrClass>()
     val inlineFunctionOrigins = mutableMapOf<IrFunction, InlineFunctionOriginInfo>()
+    val liveVariablesAtSuspensionPoints = mutableMapOf<IrSuspensionPoint, List<IrVariable>>()
+    val visibleVariablesAtSuspensionPoints = mutableMapOf<IrSuspensionPoint, List<IrVariable>>()
 
     private val localClassNames = mutableMapOf<IrAttributeContainer, String>()
     fun getLocalClassName(container: IrAttributeContainer): String? = localClassNames[container.attributeOwnerId]

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.backend.konan.driver.phases
 
+import org.jetbrains.kotlin.backend.common.BodyLoweringPass
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.lower.*
 import org.jetbrains.kotlin.backend.common.lower.coroutines.AddContinuationToNonLocalSuspendFunctionsLowering
@@ -13,14 +14,13 @@ import org.jetbrains.kotlin.backend.common.lower.inline.LocalClassesExtractionFr
 import org.jetbrains.kotlin.backend.common.lower.inline.LocalClassesInInlineFunctionsLowering
 import org.jetbrains.kotlin.backend.common.lower.inline.LocalClassesInInlineLambdasLowering
 import org.jetbrains.kotlin.backend.common.lower.loops.ForLoopsLowering
+import org.jetbrains.kotlin.backend.common.lower.optimizations.LivenessAnalysis
 import org.jetbrains.kotlin.backend.common.lower.optimizations.PropertyAccessorInlineLowering
 import org.jetbrains.kotlin.backend.common.phaser.*
 import org.jetbrains.kotlin.backend.common.runOnFilePostfix
 import org.jetbrains.kotlin.backend.konan.Context
 import org.jetbrains.kotlin.backend.konan.NativeGenerationState
 import org.jetbrains.kotlin.backend.konan.driver.PhaseEngine
-import org.jetbrains.kotlin.ir.declarations.IrFile
-import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.backend.konan.*
 import org.jetbrains.kotlin.backend.konan.driver.utilities.getDefaultIrActions
 import org.jetbrains.kotlin.backend.konan.ir.FunctionsWithoutBoundCheckGenerator
@@ -33,7 +33,9 @@ import org.jetbrains.kotlin.backend.konan.lower.ReturnsInsertionLowering
 import org.jetbrains.kotlin.backend.konan.lower.UnboxInlineLowering
 import org.jetbrains.kotlin.backend.konan.optimizations.KonanBCEForLoopBodyTransformer
 import org.jetbrains.kotlin.ir.IrElement
-import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.expressions.IrBody
+import org.jetbrains.kotlin.ir.expressions.IrSuspensionPoint
 import org.jetbrains.kotlin.ir.interpreter.IrInterpreterConfiguration
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
@@ -385,8 +387,32 @@ private val coroutinesPhase = createFileLoweringPhase(
         prerequisite = setOf(localFunctionsPhase, finallyBlocksPhase, kotlinNothingValueExceptionPhase)
 )
 
+private val coroutinesLivenessAnalysisFallbackPhase = createFileLoweringPhase(
+        lowering = ::CoroutinesLivenessAnalysisFallback,
+        name = "CoroutinesLivenessAnalysisFallback",
+        description = "Compute visible variables at suspension points",
+        prerequisite = setOf(coroutinesPhase)
+)
+
+private val coroutinesLivenessAnalysisPhase = createFileLoweringPhase(
+        lowering = { context: NativeGenerationState ->
+            object : BodyLoweringPass {
+                override fun lower(irBody: IrBody, container: IrDeclaration) {
+                    val liveVariablesAtSuspensionPoints = context.liveVariablesAtSuspensionPoints
+                    LivenessAnalysis.run(irBody) { it is IrSuspensionPoint }
+                            .forEach { (irElement, liveVariables) ->
+                                liveVariablesAtSuspensionPoints[irElement as IrSuspensionPoint] = liveVariables
+                            }
+                }
+            }
+        },
+        name = "CoroutinesLivenessAnalysis",
+        description = "Run liveness analysis for coroutines",
+        prerequisite = setOf(coroutinesPhase)
+)
+
 private val coroutinesVarSpillingPhase = createFileLoweringPhase(
-        ::CoroutinesVarSpillingLowering,
+        lowering = ::CoroutinesVarSpillingLowering,
         name = "CoroutinesVarSpilling",
         description = "Save/restore coroutines variables before/after suspension",
         prerequisite = setOf(coroutinesPhase)
@@ -554,6 +580,9 @@ private fun PhaseEngine<NativeGenerationState>.getAllLowerings() = listOfNotNull
         varargPhase,
         kotlinNothingValueExceptionPhase,
         coroutinesPhase,
+        // Either of these could be turned off without losing correctness.
+        coroutinesLivenessAnalysisPhase, // This is more optimal
+        coroutinesLivenessAnalysisFallbackPhase, // While this is simple
         coroutinesVarSpillingPhase,
         typeOperatorPhase,
         expressionBodyTransformPhase,

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

@@ -5,43 +5,45 @@
 
 package org.jetbrains.kotlin.backend.konan.lower
 
-import org.jetbrains.kotlin.backend.common.BodyLoweringPass
+import org.jetbrains.kotlin.backend.common.*
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.common.lower.irBlock
-import org.jetbrains.kotlin.backend.common.lower.optimizations.LivenessAnalysis
+import org.jetbrains.kotlin.backend.konan.NativeGenerationState
-import org.jetbrains.kotlin.backend.konan.Context
+import org.jetbrains.kotlin.backend.konan.ir.KonanSymbols
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
+import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.builders.declarations.buildField
 import org.jetbrains.kotlin.ir.builders.irGet
 import org.jetbrains.kotlin.ir.builders.irGetField
 import org.jetbrains.kotlin.ir.builders.irSet
 import org.jetbrains.kotlin.ir.builders.irSetField
 import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.expressions.IrBody
+import org.jetbrains.kotlin.ir.expressions.*
-import org.jetbrains.kotlin.ir.expressions.IrCall
-import org.jetbrains.kotlin.ir.expressions.IrExpression
-import org.jetbrains.kotlin.ir.expressions.IrSuspensionPoint
 import org.jetbrains.kotlin.ir.symbols.IrVariableSymbol
 import org.jetbrains.kotlin.ir.util.addChild
 import org.jetbrains.kotlin.ir.util.overrides
 import org.jetbrains.kotlin.ir.util.parentAsClass
+import org.jetbrains.kotlin.ir.util.render
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformer
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
+import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 
 internal val DECLARATION_ORIGIN_COROUTINE_VAR_SPILLING = IrDeclarationOriginImpl("COROUTINE_VAR_SPILLING")
 
-internal class CoroutinesVarSpillingLowering(val context: Context) : BodyLoweringPass {
+internal class CoroutinesVarSpillingLowering(val generationState: NativeGenerationState) : BodyLoweringPass {
+    private val context = generationState.context
     private val irFactory = context.irFactory
     private val symbols = context.ir.symbols
+    private val invokeSuspendFunction = symbols.invokeSuspendFunction
     private val saveCoroutineState = symbols.saveCoroutineState
     private val restoreCoroutineState = symbols.restoreCoroutineState
 
     override fun lower(irBody: IrBody, container: IrDeclaration) {
         val thisReceiver = (container as? IrSimpleFunction)?.dispatchReceiverParameter
-        if (thisReceiver == null || !container.overrides(context.ir.symbols.invokeSuspendFunction.owner))
+        if (thisReceiver == null || !container.overrides(invokeSuspendFunction.owner))
             return
 
         val coroutineClass = container.parentAsClass
-        val liveLocals = LivenessAnalysis.run(irBody) { it is IrSuspensionPoint }
 
         // TODO: optimize by using the same property for different locals.
         val localToPropertyMap = mutableMapOf<IrVariableSymbol, IrField>()
@@ -64,7 +66,10 @@ internal class CoroutinesVarSpillingLowering(val context: Context) : BodyLowerin
         val irBuilder = context.createIrBuilder(container.symbol, container.startOffset, container.endOffset)
         irBody.transformChildren(object : IrElementTransformer<List<IrVariable>> {
             override fun visitSuspensionPoint(expression: IrSuspensionPoint, data: List<IrVariable>): IrExpression {
-                expression.transformChildren(this, liveLocals[expression]!!)
+                val liveVariables = generationState.liveVariablesAtSuspensionPoints[expression]
+                        ?: generationState.visibleVariablesAtSuspensionPoints[expression]
+                        ?: error("No live variables for ${container.render()} at ${expression.suspensionPointIdParameter.name}")
+                expression.transformChildren(this, liveVariables)
 
                 return expression
             }
@@ -94,4 +99,65 @@ internal class CoroutinesVarSpillingLowering(val context: Context) : BodyLowerin
             }
         }, data = emptyList())
     }
-}
+}
+
+internal class CoroutinesLivenessAnalysisFallback(val generationState: NativeGenerationState) : FileLoweringPass, IrElementVisitorVoid {
+    private val invokeSuspendFunction = generationState.context.ir.symbols.invokeSuspendFunction
+
+    override fun lower(irFile: IrFile) {
+        if (generationState.liveVariablesAtSuspensionPoints.isEmpty())
+            irFile.acceptChildrenVoid(this)
+    }
+
+    override fun visitElement(element: IrElement) {
+        element.acceptChildrenVoid(this)
+    }
+
+    override fun visitFunction(declaration: IrFunction) {
+        val body = declaration.body
+        if (body != null && declaration.dispatchReceiverParameter != null
+                && (declaration as? IrSimpleFunction)?.overrides(invokeSuspendFunction.owner) == true
+        ) {
+            computeVisibleVariablesAtSuspensionPoints(body)
+        }
+    }
+
+    private fun computeVisibleVariablesAtSuspensionPoints(body: IrBody) {
+        body.acceptChildrenVoid(object : IrElementVisitorVoid {
+            val scopeStack = mutableListOf<MutableSet<IrVariable>>(mutableSetOf())
+
+            override fun visitElement(element: IrElement) {
+                element.acceptChildrenVoid(this)
+            }
+
+            override fun visitContainerExpression(expression: IrContainerExpression) {
+                if (!expression.isTransparentScope)
+                    scopeStack.push(mutableSetOf())
+                super.visitContainerExpression(expression)
+                if (!expression.isTransparentScope)
+                    scopeStack.pop()
+            }
+
+            override fun visitCatch(aCatch: IrCatch) {
+                scopeStack.push(mutableSetOf())
+                super.visitCatch(aCatch)
+                scopeStack.pop()
+            }
+
+            override fun visitVariable(declaration: IrVariable) {
+                super.visitVariable(declaration)
+                scopeStack.peek()!!.add(declaration)
+            }
+
+            override fun visitSuspensionPoint(expression: IrSuspensionPoint) {
+                // Skip suspensionPointIdParameter, because we don't want to save it.
+                expression.result.acceptChildrenVoid(this)
+                expression.resumeResult.acceptChildrenVoid(this)
+
+                val visibleVariables = mutableListOf<IrVariable>()
+                scopeStack.forEach { visibleVariables += it }
+                generationState.visibleVariablesAtSuspensionPoints[expression] = visibleVariables
+            }
+        })
+    }
+}