JVM_IR: inline findSuspendAndInlineLambdas

This commit is contained in:
pyos
2020-02-17 13:01:51 +01:00
committed by Ilmir Usmanov
parent 5d603a8be4
commit e6efb81014
@@ -23,7 +23,6 @@ import org.jetbrains.kotlin.codegen.inline.coroutines.FOR_INLINE_SUFFIX
import org.jetbrains.kotlin.descriptors.Modality
import org.jetbrains.kotlin.descriptors.Visibilities
import org.jetbrains.kotlin.descriptors.Visibility
import org.jetbrains.kotlin.ir.IrElement
import org.jetbrains.kotlin.ir.IrStatement
import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
import org.jetbrains.kotlin.ir.builders.*
@@ -53,8 +52,7 @@ internal val addContinuationPhase = makeIrFilePhase(
private class AddContinuationLowering(private val context: JvmBackendContext) : FileLoweringPass {
override fun lower(irFile: IrFile) {
val suspendLambdas = findSuspendAndInlineLambdas(irFile)
transformSuspendLambdasIntoContinuations(irFile, suspendLambdas)
transformSuspendLambdasIntoContinuations(irFile)
// This should be done after converting lambdas into classes to avoid breaking the invariant that
// each lambda is referenced at most once while creating `$$forInline` methods.
addContinuationObjectAndContinuationParameterToSuspendFunctions(irFile)
@@ -95,15 +93,28 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
})
}
private fun transformSuspendLambdasIntoContinuations(irFile: IrFile, suspendLambdas: List<SuspendLambdaInfo>) {
for (lambda in suspendLambdas) {
private fun transformSuspendLambdasIntoContinuations(irFile: IrFile) {
val inlineReferences = mutableSetOf<IrCallableReference>()
val suspendLambdas = mutableMapOf<IrFunctionReference, SuspendLambdaInfo>()
irFile.acceptChildrenVoid(object : IrInlineReferenceLocator(context) {
override fun visitInlineReference(argument: IrCallableReference) {
inlineReferences.add(argument)
}
override fun visitFunctionReference(expression: IrFunctionReference) {
expression.acceptChildrenVoid(this)
if (expression.isSuspend && expression.origin == IrStatementOrigin.LAMBDA && expression !in inlineReferences) {
suspendLambdas[expression] = SuspendLambdaInfo(expression)
}
}
})
for (lambda in suspendLambdas.values) {
(lambda.function.parent as IrDeclarationContainer).declarations.remove(lambda.function)
}
irFile.transformChildrenVoid(object : IrElementTransformerVoidWithContext() {
override fun visitFunctionReference(expression: IrFunctionReference): IrExpression {
if (!expression.isSuspend)
return expression
val info = suspendLambdas.singleOrNull { it.function == expression.symbol.owner } ?: return expression
val info = suspendLambdas[expression] ?: return super.visitFunctionReference(expression)
return context.createIrBuilder(expression.symbol, expression.startOffset, expression.endOffset).run {
val expressionArguments = expression.getArguments().map { it.second }
irBlock {
@@ -117,7 +128,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
"Inconsistency between callable reference to suspend lambda and the corresponding continuation"
}
+irCall(constructor.symbol).apply {
for (typeParameter in info.constructor.parentAsClass.typeParameters) {
for (typeParameter in constructor.parentAsClass.typeParameters) {
putTypeArgument(typeParameter.index, expression.getTypeArgument(typeParameter.index))
}
expressionArguments.forEachIndexed { index, argument ->
@@ -664,25 +675,6 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
})
}
private fun findSuspendAndInlineLambdas(irElement: IrElement): List<SuspendLambdaInfo> {
val suspendLambdas = arrayListOf<SuspendLambdaInfo>()
val inlineReferences = mutableSetOf<IrCallableReference>()
irElement.acceptChildrenVoid(object : IrInlineReferenceLocator(context) {
override fun visitInlineReference(argument: IrCallableReference) {
inlineReferences.add(argument)
}
override fun visitFunctionReference(expression: IrFunctionReference) {
expression.acceptChildrenVoid(this)
if (expression.isSuspend && expression.origin == IrStatementOrigin.LAMBDA && expression !in inlineReferences) {
suspendLambdas += SuspendLambdaInfo(expression)
}
}
})
return suspendLambdas
}
private class SuspendLambdaInfo(val reference: IrFunctionReference) {
val function = reference.symbol.owner
val arity = (reference.type as IrSimpleType).arguments.size - 1