diff --git a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/NativeGenerationState.kt b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/NativeGenerationState.kt index c824f241149..2050d0029c7 100644 --- a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/NativeGenerationState.kt +++ b/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() val constructedFromExportedInlineFunctions = mutableSetOf() val inlineFunctionOrigins = mutableMapOf() + val liveVariablesAtSuspensionPoints = mutableMapOf>() + val visibleVariablesAtSuspensionPoints = mutableMapOf>() private val localClassNames = mutableMapOf() fun getLocalClassName(container: IrAttributeContainer): String? = localClassNames[container.attributeOwnerId] diff --git a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/driver/phases/NativeLoweringPhases.kt b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/driver/phases/NativeLoweringPhases.kt index e75338faca3..5e3dac5e7c2 100644 --- a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/driver/phases/NativeLoweringPhases.kt +++ b/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.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, diff --git a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/CoroutinesVarSpillingLowering.kt b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/CoroutinesVarSpillingLowering.kt index 18ba1b96ded..122a9ef158d 100644 --- a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/CoroutinesVarSpillingLowering.kt +++ b/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.Context +import org.jetbrains.kotlin.backend.konan.NativeGenerationState +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.IrCall -import org.jetbrains.kotlin.ir.expressions.IrExpression -import org.jetbrains.kotlin.ir.expressions.IrSuspensionPoint +import org.jetbrains.kotlin.ir.expressions.* 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() @@ -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> { override fun visitSuspensionPoint(expression: IrSuspensionPoint, data: List): 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()) } -} \ No newline at end of file +} + +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>(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() + scopeStack.forEach { visibleVariables += it } + generationState.visibleVariablesAtSuspensionPoints[expression] = visibleVariables + } + }) + } +}