From 32ad747632a589609b5501e449c319cd3a294d5f Mon Sep 17 00:00:00 2001 From: pyos Date: Fri, 28 May 2021 10:38:54 +0200 Subject: [PATCH] JVM: load default lambda method nodes immediately The ones that are not needed are filtered out before DefaultLambda is even constructed anyway, and this way we need fewer lateinit vars. --- .../kotlin/codegen/inline/InlineCache.kt | 2 +- .../kotlin/codegen/inline/InlineCodegen.kt | 14 +++- .../kotlin/codegen/inline/LambdaInfo.kt | 65 +++++++++---------- .../kotlin/codegen/inline/PsiInlineCodegen.kt | 28 ++++---- .../codegen/inline/SourceCompilerForInline.kt | 2 +- .../codegen/inline/defaultMethodUtil.kt | 27 ++++---- .../backend/jvm/codegen/IrInlineCodegen.kt | 65 +++++++++---------- 7 files changed, 102 insertions(+), 101 deletions(-) diff --git a/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/InlineCache.kt b/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/InlineCache.kt index 8e7d1d5a063..bbc310be0f2 100644 --- a/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/InlineCache.kt +++ b/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/InlineCache.kt @@ -23,7 +23,7 @@ import org.jetbrains.org.objectweb.asm.commons.Method data class MethodId(val ownerInternalName: String, val method: Method) class InlineCache { - val classBytes: SLRUMap = SLRUMap(30, 20) + val classBytes: SLRUMap = SLRUMap(30, 20) val methodNodeById: SLRUMap = SLRUMap(60, 50) } 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 c51098d36ee..43c80b70f71 100644 --- a/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/InlineCodegen.kt +++ b/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/InlineCodegen.kt @@ -28,7 +28,7 @@ abstract class InlineCodegen( protected val jvmSignature: JvmMethodSignature, private val typeParameterMappings: TypeParameterMappings<*>, protected val sourceCompiler: SourceCompilerForInline, - protected val reifiedTypeInliner: ReifiedTypeInliner<*> + private val reifiedTypeInliner: ReifiedTypeInliner<*> ) { // TODO: implement AS_FUNCTION inline strategy private val asFunctionInline = false @@ -147,7 +147,9 @@ abstract class InlineCodegen( 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) + if (lambda.needReification) { + lambda.reifiedTypeParametersUsages.mergeAll(reifiedTypeInliner.reifyInstructions(lambda.node.node)) + } rememberCapturedForDefaultLambda(lambda) } } @@ -206,6 +208,14 @@ abstract class InlineCodegen( abstract fun extractDefaultLambdas(node: MethodNode): List + protected inline fun extractDefaultLambdas( + node: MethodNode, parameters: Map, block: ExtractedDefaultLambda.(T) -> DefaultLambda + ): List = expandMaskConditionsAndUpdateVariableNodes( + node, maskStartIndex, maskValues, methodHandleInDefaultMethodIndex, parameters.keys + ).map { + it.block(parameters[it.offset]!!) + } + private fun generateAndInsertFinallyBlocks( intoNode: MethodNode, insertPoints: List, diff --git a/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/LambdaInfo.kt b/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/LambdaInfo.kt index 5bed790ce32..b9eec5546e5 100644 --- a/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/LambdaInfo.kt +++ b/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/LambdaInfo.kt @@ -14,7 +14,6 @@ import org.jetbrains.org.objectweb.asm.Type import org.jetbrains.org.objectweb.asm.Label import org.jetbrains.org.objectweb.asm.commons.Method import org.jetbrains.org.objectweb.asm.tree.FieldInsnNode -import kotlin.properties.Delegates interface FunctionalArgument @@ -41,8 +40,6 @@ abstract class LambdaInfo(@JvmField val isCrossInline: Boolean) : FunctionalArgu val reifiedTypeParametersUsages = ReifiedTypeParametersUsages() - abstract fun generateLambdaBody(sourceCompiler: SourceCompilerForInline, reifiedTypeInliner: ReifiedTypeInliner<*>) - open val hasDispatchReceiver = true fun addAllParameters(remapper: FieldRemapper): Parameters { @@ -71,49 +68,52 @@ object NonInlineableArgumentForInlineableParameterCalledInSuspend : FunctionalAr object NonInlineableArgumentForInlineableSuspendParameter : FunctionalArgument abstract class ExpressionLambda(isCrossInline: Boolean) : LambdaInfo(isCrossInline) { - override fun generateLambdaBody(sourceCompiler: SourceCompilerForInline, reifiedTypeInliner: ReifiedTypeInliner<*>) { + fun generateLambdaBody(sourceCompiler: SourceCompilerForInline) { node = sourceCompiler.generateLambdaBody(this, reifiedTypeParametersUsages) node.node.preprocessSuspendMarkers(forInline = true, keepFakeContinuation = false) } } abstract class DefaultLambda( - private val capturedArgs: Array, + final override val lambdaClassType: Type, + capturedArgs: Array, isCrossinline: Boolean, val offset: Int, - val needReification: Boolean + val needReification: Boolean, + sourceCompiler: SourceCompilerForInline ) : LambdaInfo(isCrossinline) { - final override var isBoundCallableReference by Delegates.notNull() - private set + final override val isSuspend + get() = false // TODO: it should probably be true sometimes, but it never was + final override val isBoundCallableReference: Boolean + final override val capturedVars: List - final override lateinit var invokeMethod: Method - private set + final override val invokeMethod: Method + get() = Method(node.node.name, node.node.desc) - final override lateinit var capturedVars: List - private set + val originalBoundReceiverType: Type? - var originalBoundReceiverType: Type? = null - private set + protected val isPropertyReference: Boolean + protected val isFunctionReference: Boolean - override val isSuspend = false // TODO: it should probably be true sometimes, but it never was - - override fun generateLambdaBody(sourceCompiler: SourceCompilerForInline, reifiedTypeInliner: ReifiedTypeInliner<*>) { - val classBytes = loadClassBytesByInternalName(sourceCompiler.state, lambdaClassType.internalName) + init { + val classBytes = loadClass(sourceCompiler) val superName = ClassReader(classBytes).superName - val isPropertyReference = superName in PROPERTY_REFERENCE_SUPER_CLASSES - val isFunctionReference = superName == FUNCTION_REFERENCE.internalName || superName == FUNCTION_REFERENCE_IMPL.internalName + isPropertyReference = superName in PROPERTY_REFERENCE_SUPER_CLASSES + isFunctionReference = superName == FUNCTION_REFERENCE.internalName || superName == FUNCTION_REFERENCE_IMPL.internalName val constructorDescriptor = Type.getMethodDescriptor(Type.VOID_TYPE, *capturedArgs) val constructor = getMethodNode(classBytes, "", constructorDescriptor, lambdaClassType)?.node assert(constructor != null || capturedArgs.isEmpty()) { "Can't find non-default constructor $constructorDescriptor for default lambda $lambdaClassType" } + // This only works for primitives but not inline classes, since information about the Kotlin type of the bound + // receiver is not present anywhere. This is why with JVM_IR the constructor argument of bound references + // is already `Object`, and this field is never used. + originalBoundReceiverType = + capturedArgs.singleOrNull()?.takeIf { (isFunctionReference || isPropertyReference) && AsmUtil.isPrimitive(it) } capturedVars = if (isFunctionReference || isPropertyReference) capturedArgs.singleOrNull()?.let { - if (AsmUtil.isPrimitive(it)) { - originalBoundReceiverType = it - } listOf(capturedParamDesc(AsmUtil.RECEIVER_PARAMETER_NAME, AsmUtil.boxType(it), isSuspend = false)) } ?: emptyList() else @@ -121,23 +121,22 @@ abstract class DefaultLambda( capturedParamDesc(fieldNode.name, Type.getType(fieldNode.desc), isSuspend = false) }?.toList() ?: emptyList() isBoundCallableReference = (isFunctionReference || isPropertyReference) && capturedVars.isNotEmpty() + } + private fun loadClass(sourceCompiler: SourceCompilerForInline): ByteArray = + sourceCompiler.state.inlineCache.classBytes.getOrPut(lambdaClassType.internalName) { + loadClassBytesByInternalName(sourceCompiler.state, lambdaClassType.internalName) + } + + protected fun loadInvoke(sourceCompiler: SourceCompilerForInline, invokeMethod: Method) { + val classBytes = loadClass(sourceCompiler) val invokeNameFallback = (if (isPropertyReference) OperatorNameConventions.GET else OperatorNameConventions.INVOKE).asString() - val invokeMethod = mapAsmMethod(sourceCompiler, isPropertyReference) - // TODO: `signatureAmbiguity = true` ignores the argument types from `invokeDescriptor` and only looks at the count. - // This effectively masks incorrect results from `mapAsmDescriptor`, which hopefully won't manifest in another way. + // TODO: `signatureAmbiguity = true` ignores the argument types from `invokeMethod` and only looks at the count. node = getMethodNode(classBytes, invokeMethod.name, invokeMethod.descriptor, lambdaClassType, signatureAmbiguity = true) ?: getMethodNode(classBytes, invokeNameFallback, invokeMethod.descriptor, lambdaClassType, signatureAmbiguity = true) ?: error("Can't find method '$invokeMethod' in '${lambdaClassType.internalName}'") - this.invokeMethod = Method(node.node.name, node.node.desc) - if (needReification) { - //nested classes could also require reification - reifiedTypeParametersUsages.mergeAll(reifiedTypeInliner.reifyInstructions(node.node)) - } } - protected abstract fun mapAsmMethod(sourceCompiler: SourceCompilerForInline, isPropertyReference: Boolean): Method - private companion object { val PROPERTY_REFERENCE_SUPER_CLASSES = listOf( 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 33e2d0e581a..23027dfe97a 100644 --- a/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/PsiInlineCodegen.kt +++ b/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/PsiInlineCodegen.kt @@ -76,7 +76,7 @@ class PsiInlineCodegen( for (info in expressionMap.values) { if (info is PsiExpressionLambda) { // Can't be done immediately in `rememberClosure` for some reason: - info.generateLambdaBody(sourceCompiler, reifiedTypeInliner) + info.generateLambdaBody(sourceCompiler) // 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) { @@ -214,13 +214,10 @@ class PsiInlineCodegen( override fun reorderArgumentsIfNeeded(actualArgsWithDeclIndex: List, valueParameterTypes: List) = Unit - override fun extractDefaultLambdas(node: MethodNode): List { - return expandMaskConditionsAndUpdateVariableNodes( - node, maskStartIndex, maskValues, methodHandleInDefaultMethodIndex, - extractDefaultLambdaOffsetAndDescriptor(jvmSignature, functionDescriptor), - ::PsiDefaultLambda - ) - } + override fun extractDefaultLambdas(node: MethodNode): List = + extractDefaultLambdas(node, extractDefaultLambdaOffsetAndDescriptor(jvmSignature, functionDescriptor)) { parameter -> + PsiDefaultLambda(type, capturedArgs, parameter, offset, needReification, sourceCompiler) + } } private val FunctionDescriptor.explicitParameters @@ -323,13 +320,14 @@ class PsiExpressionLambda( } class PsiDefaultLambda( - override val lambdaClassType: Type, + lambdaClassType: Type, capturedArgs: Array, - private val parameterDescriptor: ValueParameterDescriptor, + parameterDescriptor: ValueParameterDescriptor, offset: Int, - needReification: Boolean -) : DefaultLambda(capturedArgs, parameterDescriptor.isCrossinline, offset, needReification) { - private lateinit var invokeMethodDescriptor: FunctionDescriptor + needReification: Boolean, + sourceCompiler: SourceCompilerForInline +) : DefaultLambda(lambdaClassType, capturedArgs, parameterDescriptor.isCrossinline, offset, needReification, sourceCompiler) { + private val invokeMethodDescriptor: FunctionDescriptor override val invokeMethodParameters: List get() = invokeMethodDescriptor.explicitParameters.map { it.returnType } @@ -337,7 +335,7 @@ class PsiDefaultLambda( override val invokeMethodReturnType: KotlinType? get() = invokeMethodDescriptor.returnType - override fun mapAsmMethod(sourceCompiler: SourceCompilerForInline, isPropertyReference: Boolean): Method { + init { val substitutedDescriptor = parameterDescriptor.type.memberScope .getContributedFunctions(OperatorNameConventions.INVOKE, NoLookupLocation.FROM_BACKEND) .single() @@ -355,6 +353,6 @@ class PsiDefaultLambda( // Non-suspend function references and lambdas: `(A) -> B` => `invoke(A): B` else -> substitutedDescriptor } - return sourceCompiler.state.typeMapper.mapSignatureSkipGeneric(invokeMethodDescriptor).asmMethod + loadInvoke(sourceCompiler, sourceCompiler.state.typeMapper.mapSignatureSkipGeneric(invokeMethodDescriptor).asmMethod) } } diff --git a/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/SourceCompilerForInline.kt b/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/SourceCompilerForInline.kt index c32551dd607..d8355c73a8e 100644 --- a/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/SourceCompilerForInline.kt +++ b/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/SourceCompilerForInline.kt @@ -82,7 +82,7 @@ fun loadCompiledInlineFunction( state: GenerationState ): SMAPAndMethodNode { val containerType = AsmUtil.asmTypeByClassId(containerId) - val bytes = state.inlineCache.classBytes.getOrPut(containerId) { + val bytes = state.inlineCache.classBytes.getOrPut(containerType.internalName) { findVirtualFile(state, containerId)?.contentsToByteArray() ?: throw IllegalStateException("Couldn't find declaration file for $containerId") } diff --git a/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/defaultMethodUtil.kt b/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/defaultMethodUtil.kt index f83e2d48305..9aa051cf06d 100644 --- a/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/defaultMethodUtil.kt +++ b/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/defaultMethodUtil.kt @@ -60,15 +60,15 @@ fun extractDefaultLambdaOffsetAndDescriptor( } } +class ExtractedDefaultLambda(val type: Type, val capturedArgs: Array, val offset: Int, val needReification: Boolean) -fun expandMaskConditionsAndUpdateVariableNodes( +fun expandMaskConditionsAndUpdateVariableNodes( node: MethodNode, maskStartIndex: Int, masks: List, methodHandlerIndex: Int, - defaultLambdas: Map, - lambdaConstructor: (Type, Array, T, Int, Boolean) -> R -): List { + validOffsets: Collection +): List { fun isMaskIndex(varIndex: Int): Boolean { return maskStartIndex <= varIndex && varIndex < maskStartIndex + masks.size } @@ -111,7 +111,8 @@ fun expandMaskConditionsAndUpdateVariableNodes( val toDelete = linkedSetOf() val toInsert = arrayListOf>() - val defaultLambdasInfo = extractDefaultLambdasInfo(conditions, defaultLambdas, toDelete, toInsert, lambdaConstructor) + val extractable = conditions.filter { it.expandNotDelete && it.varIndex in validOffsets } + val defaultLambdasInfo = extractDefaultLambdasInfo(extractable, toDelete, toInsert) val indexToVarNode = node.localVariables?.filter { it.index < maskStartIndex }?.associateBy { it.index } ?: emptyMap() conditions.forEach { @@ -131,7 +132,7 @@ fun expandMaskConditionsAndUpdateVariableNodes( } node.localVariables.removeIf { - (it.start in toDelete && it.end in toDelete) || defaultLambdas.contains(it.index) + (it.start in toDelete && it.end in toDelete) || validOffsets.contains(it.index) } node.remove(toDelete) @@ -139,16 +140,12 @@ fun expandMaskConditionsAndUpdateVariableNodes( return defaultLambdasInfo } - -private fun extractDefaultLambdasInfo( +private fun extractDefaultLambdasInfo( conditions: List, - defaultLambdas: Map, toDelete: MutableCollection, - toInsert: MutableList>, - lambdaConstructor: (Type, Array, T, Int, Boolean) -> R -): List { - val defaultLambdaConditions = conditions.filter { it.expandNotDelete && defaultLambdas.contains(it.varIndex) } - return defaultLambdaConditions.map { + toInsert: MutableList> +): List { + return conditions.map { val varAssignmentInstruction = it.varInsNode!! var instanceInstuction = varAssignmentInstruction.previous if (instanceInstuction is TypeInsnNode && instanceInstuction.opcode == Opcodes.CHECKCAST) { @@ -190,7 +187,7 @@ private fun extractDefaultLambdasInfo( toInsert.add(varAssignmentInstruction to defaultLambdaFakeCallStub(argTypes, it.varIndex)) - lambdaConstructor(owner, argTypes, defaultLambdas[it.varIndex]!!, it.varIndex, needReification) + ExtractedDefaultLambda(owner, argTypes, it.varIndex, needReification) } } 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 f491d8adbf2..edbf99cc8c1 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 @@ -65,7 +65,7 @@ class IrInlineCodegen( val irReference = (argumentExpression as IrBlock).statements.filterIsInstance().single() val lambdaInfo = IrExpressionLambdaImpl(codegen, irReference, irValueParameter) rememberClosure(parameterType, irValueParameter.index, lambdaInfo) - lambdaInfo.generateLambdaBody(sourceCompiler, reifiedTypeInliner) + lambdaInfo.generateLambdaBody(sourceCompiler) lambdaInfo.reference.getArgumentsWithIr().forEachIndexed { index, (_, ir) -> putCapturedValueOnStack(ir, lambdaInfo.capturedVars[index], index) } @@ -131,13 +131,10 @@ class IrInlineCodegen( generateStub(text, codegen) } - override fun extractDefaultLambdas(node: MethodNode): List { - return expandMaskConditionsAndUpdateVariableNodes( - node, maskStartIndex, maskValues, methodHandleInDefaultMethodIndex, - extractDefaultLambdaOffsetAndDescriptor(jvmSignature, function), - ::IrDefaultLambda - ) - } + override fun extractDefaultLambdas(node: MethodNode): List = + extractDefaultLambdas(node, extractDefaultLambdaOffsetAndDescriptor(jvmSignature, function)) { parameter -> + IrDefaultLambda(type, capturedArgs, parameter, offset, needReification, sourceCompiler as IrSourceCompilerForInline) + } } class IrExpressionLambdaImpl( @@ -199,13 +196,30 @@ class IrExpressionLambdaImpl( } class IrDefaultLambda( - override val lambdaClassType: Type, + lambdaClassType: Type, capturedArgs: Array, - private val irValueParameter: IrValueParameter, + irValueParameter: IrValueParameter, offset: Int, - needReification: Boolean -) : DefaultLambda(capturedArgs, irValueParameter.isCrossinline, offset, needReification) { - private lateinit var typeArguments: List + needReification: Boolean, + sourceCompiler: IrSourceCompilerForInline +) : DefaultLambda(lambdaClassType, capturedArgs, irValueParameter.isCrossinline, offset, needReification, sourceCompiler) { + private val typeArguments: List = (irValueParameter.type as IrSimpleType).arguments.let { + val context = sourceCompiler.codegen.context + if (isPropertyReference) { + // Property references: `(A) -> B` => `get(Any?): Any?` + List(it.size) { context.irBuiltIns.anyNType } + } else { + // Non-suspend function references and lambdas: `(A) -> B` => `invoke(A): B` + // Suspend function references: `suspend (A) -> B` => `invoke(A, Continuation): Any?` + // TODO: default suspend lambdas are currently uninlinable + it.mapTo(mutableListOf()) { argument -> (argument as IrTypeProjection).type }.apply { + if (irValueParameter.type.isSuspendFunction()) { + set(size - 1, context.ir.symbols.continuationClass.typeWith(get(size - 1))) + add(context.irBuiltIns.anyNType) + } + } + } + } override val invokeMethodParameters: List get() = typeArguments.dropLast(1).map { it.toIrBasedKotlinType() } @@ -213,33 +227,16 @@ class IrDefaultLambda( override val invokeMethodReturnType: KotlinType get() = typeArguments.last().toIrBasedKotlinType() - override fun mapAsmMethod(sourceCompiler: SourceCompilerForInline, isPropertyReference: Boolean): Method { - val context = (sourceCompiler as IrSourceCompilerForInline).codegen.context - typeArguments = (irValueParameter.type as IrSimpleType).arguments.let { - if (isPropertyReference) { - // Property references: `(A) -> B` => `get(Any?): Any?` - List(it.size) { context.irBuiltIns.anyNType } - } else { - // Non-suspend function references and lambdas: `(A) -> B` => `invoke(A): B` - // Suspend function references: `suspend (A) -> B` => `invoke(A, Continuation): Any?` - // TODO: default suspend lambdas are currently uninlinable - it.mapTo(mutableListOf()) { argument -> (argument as IrTypeProjection).type }.apply { - if (irValueParameter.type.isSuspendFunction()) { - set(size - 1, context.ir.symbols.continuationClass.typeWith(get(size - 1))) - add(context.irBuiltIns.anyNType) - } - } - } - } + init { val base = if (isPropertyReference) OperatorNameConventions.GET.asString() else OperatorNameConventions.INVOKE.asString() val name = InlineClassAbi.hashSuffix( - context.state.useOldManglingSchemeForFunctionsWithInlineClassesInSignatures, + sourceCompiler.codegen.context.state.useOldManglingSchemeForFunctionsWithInlineClassesInSignatures, typeArguments.dropLast(1), typeArguments.last().takeIf { it.isInlineClassType() } )?.let { "$base-$it" } ?: base - // TODO: while technically only the number of arguments here matters right now (see DefaultLambdaInfo.generateLambdaBody), + // TODO: while technically only the number of arguments here matters right now (see `loadInvoke`), // it would be better to map to a non-erased signature if not a property reference. - return Method(name, AsmTypes.OBJECT_TYPE, Array(typeArguments.size - 1) { AsmTypes.OBJECT_TYPE }) + loadInvoke(sourceCompiler, Method(name, AsmTypes.OBJECT_TYPE, Array(typeArguments.size - 1) { AsmTypes.OBJECT_TYPE })) } }