From 8a3af851d9a5ea782d142168fbdd4ae7668d76fd Mon Sep 17 00:00:00 2001 From: pyos Date: Tue, 25 May 2021 16:23:05 +0200 Subject: [PATCH] JVM: pregenerate inline lambda bodies as early as possible Which in JVM_IR is immediately in `genValueAndPut`, but for the old backend needs to be delayed until `genCallInner` for some reason. --- .../kotlin/codegen/inline/InlineCodegen.kt | 39 +++------ .../kotlin/codegen/inline/PsiInlineCodegen.kt | 81 +++++++++---------- .../backend/jvm/codegen/IrInlineCodegen.kt | 20 +---- 3 files changed, 55 insertions(+), 85 deletions(-) diff --git a/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/InlineCodegen.kt b/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/InlineCodegen.kt index d7cba6c4f3d..4536948555f 100644 --- a/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/InlineCodegen.kt +++ b/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/InlineCodegen.kt @@ -48,7 +48,7 @@ abstract class InlineCodegen( protected val jvmSignature: JvmMethodSignature, private val typeParameterMappings: TypeParameterMappings<*>, protected val sourceCompiler: SourceCompilerForInline, - private val reifiedTypeInliner: ReifiedTypeInliner<*> + protected val reifiedTypeInliner: ReifiedTypeInliner<*> ) { init { assert(InlineUtil.isInline(functionDescriptor)) { @@ -193,21 +193,27 @@ abstract class InlineCodegen( ?: error("No stack value for continuation parameter of suspend function") } + protected fun rememberClosure(parameterType: Type, index: Int, lambdaInfo: LambdaInfo) { + val closureInfo = invocationParamBuilder.addNextValueParameter(parameterType, true, null, index) + closureInfo.functionalArgument = lambdaInfo + expressionMap[closureInfo.index] = lambdaInfo + } + private fun inlineCall(nodeAndSmap: SMAPAndMethodNode, inlineDefaultLambda: Boolean): InlineResult { assert(delayedHiddenWriting == null) { "'putHiddenParamsIntoLocals' should be called after 'processAndPutHiddenParameters(true)'" } + val node = nodeAndSmap.node if (inlineDefaultLambda) { for (lambda in extractDefaultLambdas(node)) { invocationParamBuilder.buildParameters().getParameterByDeclarationSlot(lambda.offset).functionalArgument = lambda val prev = expressionMap.put(lambda.offset, lambda) assert(prev == null) { "Lambda with offset ${lambda.offset} already exists: $prev" } + lambda.generateLambdaBody(sourceCompiler, reifiedTypeInliner) + rememberCapturedForDefaultLambda(lambda) } } - val reificationResult = reifiedTypeInliner.reifyInstructions(node) - generateClosuresBodies() - //through generation captured parameters will be added to invocationParamBuilder - putClosureParametersOnStack() + val reificationResult = reifiedTypeInliner.reifyInstructions(node) val parameters = invocationParamBuilder.buildParameters() @@ -340,14 +346,6 @@ abstract class InlineCodegen( } } - private fun generateClosuresBodies() { - for (info in expressionMap.values) { - if (info is LambdaInfo) { - info.generateLambdaBody(sourceCompiler, reifiedTypeInliner) - } - } - } - protected fun putArgumentOrCapturedToLocalVal( jvmKotlinType: JvmKotlinType, stackValue: StackValue, @@ -431,19 +429,7 @@ abstract class InlineCodegen( } } - private fun putClosureParametersOnStack() { - for (next in expressionMap.values) { - if (next is LambdaInfo) { - // Closure parameters for bound references have already been generated in between other arguments. - if (next is ExpressionLambda && next.isBoundCallableReference) continue - putClosureParametersOnStack(next) - } - } - } - - protected abstract fun putClosureParametersOnStack(next: LambdaInfo) - - protected fun rememberCapturedForDefaultLambda(defaultLambda: DefaultLambda) { + private fun rememberCapturedForDefaultLambda(defaultLambda: DefaultLambda) { for ((paramIndex, captured) in defaultLambda.capturedVars.withIndex()) { putArgumentOrCapturedToLocalVal( JvmKotlinType(captured.type), @@ -459,7 +445,6 @@ abstract class InlineCodegen( } } - protected fun processDefaultMaskOrMethodHandler(value: StackValue, kind: ValueKind): Boolean { if (kind !== ValueKind.DEFAULT_MASK && kind !== ValueKind.METHOD_HANDLE_IN_DEFAULT) { return false diff --git a/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/PsiInlineCodegen.kt b/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/PsiInlineCodegen.kt index bfec4a9806c..b504cfe7156 100644 --- a/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/PsiInlineCodegen.kt +++ b/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/PsiInlineCodegen.kt @@ -28,7 +28,6 @@ import org.jetbrains.kotlin.resolve.inline.InlineUtil import org.jetbrains.kotlin.resolve.inline.InlineUtil.isInlinableParameterExpression import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodParameterKind import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature -import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue import org.jetbrains.kotlin.serialization.deserialization.descriptors.DescriptorWithContainerSource import org.jetbrains.kotlin.types.KotlinType import org.jetbrains.kotlin.util.OperatorNameConventions @@ -72,6 +71,17 @@ class PsiInlineCodegen( } try { val registerLineNumber = registerLineNumberAfterwards(resolvedCall) + for (info in expressionMap.values) { + if (info is PsiExpressionLambda) { + // Can't be done immediately in `rememberClosure` for some reason: + info.generateLambdaBody(sourceCompiler, reifiedTypeInliner) + // Requires `generateLambdaBody` first if the closure is non-empty (for bound callable references, + // or indeed any callable references, it *is* empty, so this was done in `rememberClosure`): + if (!info.isBoundCallableReference) { + putClosureParametersOnStack(info, null) + } + } + } performInline(resolvedCall?.typeArguments?.keys?.toList(), callDefault, callDefault, codegen.typeSystem, registerLineNumber) } finally { state.globalInlineContext.exitFromInlining() @@ -102,31 +112,6 @@ class PsiInlineCodegen( delayedHiddenWriting = recordParameterValueInLocalVal(justProcess, false, *hiddenParameters.toTypedArray()) } - var activeLambda: LambdaInfo? = null - private set - - override fun putClosureParametersOnStack(next: LambdaInfo) { - activeLambda = next - when (next) { - is PsiExpressionLambda -> { - val receiverValue = next.boundReceiver?.let { boundReceiver -> - val receiver = codegen.generateReceiverValue(boundReceiver, false) - val receiverKotlinType = receiver.kotlinType - val boxedReceiver = - if (receiverKotlinType != null) - DescriptorAsmUtil.boxType(receiver.type, receiverKotlinType, state.typeMapper) - else - AsmUtil.boxType(receiver.type) - StackValue.coercion(receiver, boxedReceiver, receiverKotlinType) - } - codegen.pushClosureOnStack(next.classDescriptor, true, this, receiverValue) - } - is DefaultLambda -> rememberCapturedForDefaultLambda(next) - else -> throw RuntimeException("Unknown lambda: $next") - } - activeLambda = null - } - /*lambda or callable reference*/ private fun isInliningParameter(expression: KtExpression, valueParameterDescriptor: ValueParameterDescriptor): Boolean { //TODO deparenthesize typed @@ -147,11 +132,7 @@ class PsiInlineCodegen( } if (isInliningParameter(argumentExpression, valueParameterDescriptor)) { - val lambdaInfo = rememberClosure(argumentExpression, parameterType.type, valueParameterDescriptor) - if (lambdaInfo.boundReceiver != null) { - // Has to be done immediately to preserve evaluation order. - putClosureParametersOnStack(lambdaInfo) - } + rememberClosure(argumentExpression, parameterType.type, valueParameterDescriptor) } else { val value = codegen.gen(argumentExpression) val kind = when { @@ -170,7 +151,7 @@ class PsiInlineCodegen( private fun isCallSiteIsSuspend(descriptor: ValueParameterDescriptor): Boolean = state.bindingContext[CodegenBinding.CALL_SITE_IS_SUSPEND_FOR_CROSSINLINE_LAMBDA, descriptor] == true - private fun rememberClosure(expression: KtExpression, type: Type, parameter: ValueParameterDescriptor): PsiExpressionLambda { + private fun rememberClosure(expression: KtExpression, type: Type, parameter: ValueParameterDescriptor) { val ktLambda = KtPsiUtil.deparenthesize(expression) assert(isInlinableParameterExpression(ktLambda)) { "Couldn't find inline expression in ${expression.text}" } @@ -179,15 +160,33 @@ class PsiInlineCodegen( JvmCodegenUtil.getBoundCallableReferenceReceiver(resolvedCall) } else null - return PsiExpressionLambda( - ktLambda!!, state.typeMapper, state.languageVersionSettings, parameter.isCrossinline, boundReceiver - ).also { lambda -> - val closureInfo = invocationParamBuilder.addNextValueParameter(type, true, null, parameter.index) - closureInfo.functionalArgument = lambda - expressionMap[closureInfo.index] = lambda + val lambda = PsiExpressionLambda( + ktLambda!!, state.typeMapper, state.languageVersionSettings, parameter.isCrossinline, boundReceiver != null + ) + rememberClosure(type, parameter.index, lambda) + if (boundReceiver != null) { + // Has to be done immediately to preserve evaluation order. + val receiver = codegen.generateReceiverValue(boundReceiver, false) + val receiverKotlinType = receiver.kotlinType + val boxedReceiver = + if (receiverKotlinType != null) + DescriptorAsmUtil.boxType(receiver.type, receiverKotlinType, state.typeMapper) + else + AsmUtil.boxType(receiver.type) + val receiverValue = StackValue.coercion(receiver, boxedReceiver, receiverKotlinType) + putClosureParametersOnStack(lambda, receiverValue) } } + var activeLambda: LambdaInfo? = null + private set + + private fun putClosureParametersOnStack(next: PsiExpressionLambda, receiverValue: StackValue?) { + activeLambda = next + codegen.pushClosureOnStack(next.classDescriptor, true, this, receiverValue) + activeLambda = null + } + override fun putValueIfNeeded(parameterType: JvmKotlinType, value: StackValue, kind: ValueKind, parameterIndex: Int) { if (processDefaultMaskOrMethodHandler(value, kind)) return @@ -231,11 +230,8 @@ class PsiExpressionLambda( private val typeMapper: KotlinTypeMapper, private val languageVersionSettings: LanguageVersionSettings, isCrossInline: Boolean, - val boundReceiver: ReceiverValue? -) : ExpressionLambda(isCrossInline) { override val isBoundCallableReference: Boolean - get() = boundReceiver != null - +) : ExpressionLambda(isCrossInline) { override val lambdaClassType: Type override val invokeMethod: Method @@ -297,6 +293,7 @@ class PsiExpressionLambda( isSuspend = invokeMethodDescriptor.isSuspend } + // This can only be computed after generating the body, hence `lazy`. override val capturedVars: List by lazy { arrayListOf().apply { val captureThis = closure.capturedOuterClassDescriptor diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrInlineCodegen.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrInlineCodegen.kt index 3b5176bc11c..5d839db9234 100644 --- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrInlineCodegen.kt +++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrInlineCodegen.kt @@ -57,16 +57,6 @@ class IrInlineCodegen( } } - override fun putClosureParametersOnStack(next: LambdaInfo) { - when (next) { - is IrExpressionLambdaImpl -> next.reference.getArgumentsWithIr().forEachIndexed { index, (_, ir) -> - putCapturedValueOnStack(ir, next.capturedVars[index], index) - } - is IrDefaultLambda -> rememberCapturedForDefaultLambda(next) - else -> throw RuntimeException("Unknown lambda: $next") - } - } - override fun genValueAndPut( irValueParameter: IrValueParameter, argumentExpression: IrExpression, @@ -78,12 +68,10 @@ class IrInlineCodegen( if (isInlineParameter && argumentExpression.isInlineIrExpression()) { val irReference = (argumentExpression as IrBlock).statements.filterIsInstance().single() val lambdaInfo = IrExpressionLambdaImpl(codegen, irReference, irValueParameter) - val closureInfo = invocationParamBuilder.addNextValueParameter(parameterType, true, null, irValueParameter.index) - closureInfo.functionalArgument = lambdaInfo - expressionMap[closureInfo.index] = lambdaInfo - val boundReceiver = irReference.extensionReceiver - if (boundReceiver != null) { - putCapturedValueOnStack(boundReceiver, lambdaInfo.capturedVars.single(), 0) + rememberClosure(parameterType, irValueParameter.index, lambdaInfo) + lambdaInfo.generateLambdaBody(sourceCompiler, reifiedTypeInliner) + lambdaInfo.reference.getArgumentsWithIr().forEachIndexed { index, (_, ir) -> + putCapturedValueOnStack(ir, lambdaInfo.capturedVars[index], index) } } else { val kind = when (irValueParameter.origin) {