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.
This commit is contained in:
pyos
2021-05-28 10:38:54 +02:00
committed by max-kammerer
parent 1e953eaf01
commit 32ad747632
7 changed files with 102 additions and 101 deletions
@@ -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<ClassId, ByteArray> = SLRUMap(30, 20)
val classBytes: SLRUMap<String, ByteArray> = SLRUMap(30, 20)
val methodNodeById: SLRUMap<MethodId, SMAPAndMethodNode> = SLRUMap(60, 50)
}
@@ -28,7 +28,7 @@ abstract class InlineCodegen<out T : BaseExpressionCodegen>(
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<out T : BaseExpressionCodegen>(
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<out T : BaseExpressionCodegen>(
abstract fun extractDefaultLambdas(node: MethodNode): List<DefaultLambda>
protected inline fun <T> extractDefaultLambdas(
node: MethodNode, parameters: Map<Int, T>, block: ExtractedDefaultLambda.(T) -> DefaultLambda
): List<DefaultLambda> = expandMaskConditionsAndUpdateVariableNodes(
node, maskStartIndex, maskValues, methodHandleInDefaultMethodIndex, parameters.keys
).map {
it.block(parameters[it.offset]!!)
}
private fun generateAndInsertFinallyBlocks(
intoNode: MethodNode,
insertPoints: List<MethodInliner.PointForExternalFinallyBlocks>,
@@ -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<Type>,
final override val lambdaClassType: Type,
capturedArgs: Array<Type>,
isCrossinline: Boolean,
val offset: Int,
val needReification: Boolean
val needReification: Boolean,
sourceCompiler: SourceCompilerForInline
) : LambdaInfo(isCrossinline) {
final override var isBoundCallableReference by Delegates.notNull<Boolean>()
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<CapturedParamDesc>
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<CapturedParamDesc>
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, "<init>", constructorDescriptor, lambdaClassType)?.node
assert(constructor != null || capturedArgs.isEmpty()) {
"Can't find non-default constructor <init>$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(
@@ -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<ArgumentAndDeclIndex>, valueParameterTypes: List<Type>) = Unit
override fun extractDefaultLambdas(node: MethodNode): List<DefaultLambda> {
return expandMaskConditionsAndUpdateVariableNodes(
node, maskStartIndex, maskValues, methodHandleInDefaultMethodIndex,
extractDefaultLambdaOffsetAndDescriptor(jvmSignature, functionDescriptor),
::PsiDefaultLambda
)
}
override fun extractDefaultLambdas(node: MethodNode): List<DefaultLambda> =
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<Type>,
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<KotlinType?>
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)
}
}
@@ -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")
}
@@ -60,15 +60,15 @@ fun extractDefaultLambdaOffsetAndDescriptor(
}
}
class ExtractedDefaultLambda(val type: Type, val capturedArgs: Array<Type>, val offset: Int, val needReification: Boolean)
fun <T, R : DefaultLambda> expandMaskConditionsAndUpdateVariableNodes(
fun expandMaskConditionsAndUpdateVariableNodes(
node: MethodNode,
maskStartIndex: Int,
masks: List<Int>,
methodHandlerIndex: Int,
defaultLambdas: Map<Int, T>,
lambdaConstructor: (Type, Array<Type>, T, Int, Boolean) -> R
): List<R> {
validOffsets: Collection<Int>
): List<ExtractedDefaultLambda> {
fun isMaskIndex(varIndex: Int): Boolean {
return maskStartIndex <= varIndex && varIndex < maskStartIndex + masks.size
}
@@ -111,7 +111,8 @@ fun <T, R : DefaultLambda> expandMaskConditionsAndUpdateVariableNodes(
val toDelete = linkedSetOf<AbstractInsnNode>()
val toInsert = arrayListOf<Pair<AbstractInsnNode, AbstractInsnNode>>()
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 <T, R : DefaultLambda> 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 <T, R : DefaultLambda> expandMaskConditionsAndUpdateVariableNodes(
return defaultLambdasInfo
}
private fun <T, R : DefaultLambda> extractDefaultLambdasInfo(
private fun extractDefaultLambdasInfo(
conditions: List<Condition>,
defaultLambdas: Map<Int, T>,
toDelete: MutableCollection<AbstractInsnNode>,
toInsert: MutableList<Pair<AbstractInsnNode, AbstractInsnNode>>,
lambdaConstructor: (Type, Array<Type>, T, Int, Boolean) -> R
): List<R> {
val defaultLambdaConditions = conditions.filter { it.expandNotDelete && defaultLambdas.contains(it.varIndex) }
return defaultLambdaConditions.map {
toInsert: MutableList<Pair<AbstractInsnNode, AbstractInsnNode>>
): List<ExtractedDefaultLambda> {
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 <T, R : DefaultLambda> extractDefaultLambdasInfo(
toInsert.add(varAssignmentInstruction to defaultLambdaFakeCallStub(argTypes, it.varIndex))
lambdaConstructor(owner, argTypes, defaultLambdas[it.varIndex]!!, it.varIndex, needReification)
ExtractedDefaultLambda(owner, argTypes, it.varIndex, needReification)
}
}
@@ -65,7 +65,7 @@ class IrInlineCodegen(
val irReference = (argumentExpression as IrBlock).statements.filterIsInstance<IrFunctionReference>().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<DefaultLambda> {
return expandMaskConditionsAndUpdateVariableNodes(
node, maskStartIndex, maskValues, methodHandleInDefaultMethodIndex,
extractDefaultLambdaOffsetAndDescriptor(jvmSignature, function),
::IrDefaultLambda
)
}
override fun extractDefaultLambdas(node: MethodNode): List<DefaultLambda> =
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<Type>,
private val irValueParameter: IrValueParameter,
irValueParameter: IrValueParameter,
offset: Int,
needReification: Boolean
) : DefaultLambda(capturedArgs, irValueParameter.isCrossinline, offset, needReification) {
private lateinit var typeArguments: List<IrType>
needReification: Boolean,
sourceCompiler: IrSourceCompilerForInline
) : DefaultLambda(lambdaClassType, capturedArgs, irValueParameter.isCrossinline, offset, needReification, sourceCompiler) {
private val typeArguments: List<IrType> = (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<B>): 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<KotlinType>
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<B>): 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 }))
}
}