JVM IR: optimize special method / signature computation in BridgeLowering

First of all, put method signature caches of BridgeLowering into JvmBackendContext. Since method hierarchies can span several files, and a new instance of BridgeLowering is created to lower each file, keeping them cached for the whole module makes BridgeLowering faster. Also, do not attempt to compute special bridges if the method is irrelevant, which can be deduced by its name. With this optimization, the special method cache is no longer needed. This brings BridgeLowering time from 3.8% down to 1.5% of all compiler time on a sample project.
2020-08-12 18:20:57 +02:00
parent 0727e9055b
commit adcfbdec24
3 changed files with 81 additions and 53 deletions
@@ -60,7 +60,7 @@ class SpecialBridgeMethods(val context: CommonBackendContext) {
    private fun getSecondArg(bridge: IrSimpleFunction) =
        IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, bridge.valueParameters[1].symbol)

-    private val SPECIAL_METHODS_WITH_DEFAULTS_MAP = mapOf(
+    private val specialMethodsWithDefaults = mapOf(
        makeDescription(KotlinBuiltIns.FQ_NAMES.collection, "contains", 1) to
                SpecialMethodWithDefaultInfo(::constFalse, 1),
        makeDescription(KotlinBuiltIns.FQ_NAMES.mutableCollection, "remove", 1) to
@@ -83,7 +83,7 @@ class SpecialBridgeMethods(val context: CommonBackendContext) {
                SpecialMethodWithDefaultInfo(::constNull, 1, needsGenericSignature = true)
    )

-    private val SPECIAL_PROPERTIES_SET = mapOf(
+    private val specialProperties = mapOf(
        makeDescription(KotlinBuiltIns.FQ_NAMES.collection, "size") to BuiltInWithDifferentJvmName(),
        makeDescription(KotlinBuiltIns.FQ_NAMES.map, "size") to BuiltInWithDifferentJvmName(),
        makeDescription(KotlinBuiltIns.FQ_NAMES.charSequence.toSafe(), "length") to BuiltInWithDifferentJvmName(),
@@ -92,7 +92,7 @@ class SpecialBridgeMethods(val context: CommonBackendContext) {
        makeDescription(KotlinBuiltIns.FQ_NAMES.map, "entries") to BuiltInWithDifferentJvmName(needsGenericSignature = true)
    )

-    private val SPECIAL_METHODS_SETS = mapOf(
+    private val specialMethods = mapOf(
        makeDescription(KotlinBuiltIns.FQ_NAMES.number.toSafe(), "toByte") to BuiltInWithDifferentJvmName(),
        makeDescription(KotlinBuiltIns.FQ_NAMES.number.toSafe(), "toShort") to BuiltInWithDifferentJvmName(),
        makeDescription(KotlinBuiltIns.FQ_NAMES.number.toSafe(), "toInt") to BuiltInWithDifferentJvmName(),
@@ -103,10 +103,13 @@ class SpecialBridgeMethods(val context: CommonBackendContext) {
        makeDescription(KotlinBuiltIns.FQ_NAMES.mutableList, "removeAt", 1) to BuiltInWithDifferentJvmName(needsGenericSignature = true)
    )

+    val specialMethodNames = (specialMethodsWithDefaults + specialMethods).map { (description) -> description.name }.toHashSet()
+    val specialPropertyNames = specialProperties.map { (description) -> description.name }.toHashSet()
+
    fun findSpecialWithOverride(irFunction: IrSimpleFunction): Pair<IrSimpleFunction, SpecialMethodWithDefaultInfo>? {
        irFunction.allOverridden().forEach { overridden ->
            val description = overridden.toDescription()
-            SPECIAL_METHODS_WITH_DEFAULTS_MAP[description]?.let {
+            specialMethodsWithDefaults[description]?.let {
                return Pair(overridden, it)
            }
        }
@@ -115,7 +118,7 @@ class SpecialBridgeMethods(val context: CommonBackendContext) {

    fun getSpecialMethodInfo(irFunction: IrSimpleFunction): SpecialMethodWithDefaultInfo? {
        val description = irFunction.toDescription()
-        return SPECIAL_METHODS_WITH_DEFAULTS_MAP[description]
+        return specialMethodsWithDefaults[description]
    }

    fun getBuiltInWithDifferentJvmName(irFunction: IrSimpleFunction): BuiltInWithDifferentJvmName? {
@@ -123,10 +126,10 @@ class SpecialBridgeMethods(val context: CommonBackendContext) {
            val classFqName = irFunction.parentAsClass.fqNameWhenAvailable
                ?: return null

-            return SPECIAL_PROPERTIES_SET[makeDescription(classFqName, it.owner.name.asString())]
+            return specialProperties[makeDescription(classFqName, it.owner.name.asString())]
        }

-        return SPECIAL_METHODS_SETS[irFunction.toDescription()]
+        return specialMethods[irFunction.toDescription()]
    }
 }

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.backend.jvm.codegen.MethodSignatureMapper
 import org.jetbrains.kotlin.backend.jvm.codegen.createFakeContinuation
 import org.jetbrains.kotlin.backend.jvm.descriptors.JvmSharedVariablesManager
 import org.jetbrains.kotlin.backend.jvm.intrinsics.IrIntrinsicMethods
+import org.jetbrains.kotlin.backend.jvm.lower.BridgeLowering
 import org.jetbrains.kotlin.backend.jvm.lower.CollectionStubComputer
 import org.jetbrains.kotlin.backend.jvm.lower.JvmInnerClassesSupport
 import org.jetbrains.kotlin.backend.jvm.lower.inlineclasses.InlineClassAbi
@@ -106,6 +107,7 @@ class JvmBackendContext(
    internal val hiddenConstructors = mutableMapOf<IrConstructor, IrConstructor>()

    internal val collectionStubComputer = CollectionStubComputer(this)
+    internal val bridgeLoweringCache = BridgeLowering.BridgeLoweringCache(this)

    override var inVerbosePhase: Boolean = false

@@ -36,7 +36,6 @@ import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
-import org.jetbrains.kotlin.utils.getOrPutNullable
 import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.commons.Method

@@ -121,7 +120,7 @@ internal val bridgePhase = makeIrFilePhase(
    description = "Generate bridges"
 )

-private class BridgeLowering(val context: JvmBackendContext) : FileLoweringPass, IrElementTransformerVoid() {
+internal class BridgeLowering(val context: JvmBackendContext) : FileLoweringPass, IrElementTransformerVoid() {
    // Represents a synthetic bridge to `overridden` with a precomputed signature
    private class Bridge(
        val overridden: IrSimpleFunction,
@@ -142,7 +141,7 @@ private class BridgeLowering(val context: JvmBackendContext) : FileLoweringPass,
    //
    // Finally, we sometimes need to use INVOKESPECIAL to invoke an existing special bridge implementation in a
    // superclass, which is what `superQualifierSymbol` is for.
-    private data class SpecialBridge(
+    data class SpecialBridge(
        val overridden: IrSimpleFunction,
        val signature: Method,
        // We need to produce a generic signature if the underlying Java method contains type parameters.
@@ -324,46 +323,7 @@ private class BridgeLowering(val context: JvmBackendContext) : FileLoweringPass,

    // Returns the special bridge overridden by the current methods if it exists.
    private val IrSimpleFunction.specialBridgeOrNull: SpecialBridge?
-        get() = specialBridgeCache.getOrPutNullable(symbol) {
-            specialBridgeMethods.getSpecialMethodInfo(this)?.let { specialMethodInfo ->
-                SpecialBridge(
-                    this, jvmMethod, methodInfo = specialMethodInfo, needsGenericSignature = specialMethodInfo.needsGenericSignature
-                )
-            } ?: specialBridgeMethods.getBuiltInWithDifferentJvmName(this)?.let { specialBuiltInInfo ->
-                SpecialBridge(this, jvmMethod, needsGenericSignature = specialBuiltInInfo.needsGenericSignature)
-            } ?: overriddenSymbols.asSequence().mapNotNull { it.owner.specialBridgeOrNull }.firstOrNull()?.let { specialBridge ->
-                if (specialBridge.needsGenericSignature) {
-                    // Compute the substituted signature.
-                    val erasedParameterCount = specialBridge.methodInfo?.argumentsToCheck ?: 0
-                    val substitutedParameterTypes = valueParameters.mapIndexed { index, param ->
-                        if (index < erasedParameterCount) context.irBuiltIns.anyNType else param.type
-                    }
-
-                    val substitutedOverride = context.irFactory.buildFun {
-                        updateFrom(specialBridge.overridden)
-                        name = Name.identifier(specialBridge.signature.name)
-                        returnType = this@specialBridgeOrNull.returnType
-                    }.apply {
-                        // All existing special bridges only have value parameter types.
-                        valueParameters = this@specialBridgeOrNull.valueParameters.zip(substitutedParameterTypes).map { (param, type) ->
-                            param.copyTo(this, IrDeclarationOrigin.BRIDGE, type = type)
-                        }
-                        overriddenSymbols = listOf(specialBridge.overridden.symbol)
-                        parent = this@specialBridgeOrNull.parent
-                    }
-
-                    specialBridge.copy(
-                        signature = substitutedOverride.jvmMethod,
-                        substitutedParameterTypes = substitutedParameterTypes,
-                        substitutedReturnType = returnType
-                    )
-                } else {
-                    specialBridge
-                }
-            }
-        }
-    private val specialBridgeMethods = SpecialBridgeMethods(context)
-    private val specialBridgeCache = mutableMapOf<IrSimpleFunctionSymbol, SpecialBridge?>()
+        get() = context.bridgeLoweringCache.computeSpecialBridge(this)

    private fun IrSimpleFunction.overriddenFinalSpecialBridges(): List<Method> = allOverridden().mapNotNullTo(mutableListOf()) {
        // Ignore special bridges in interfaces or Java classes. While we never generate special bridges in Java
@@ -565,10 +525,73 @@ private class BridgeLowering(val context: JvmBackendContext) : FileLoweringPass,
    private fun IrBuilderWithScope.irCastIfNeeded(expression: IrExpression, to: IrType): IrExpression =
        if (expression.type == to || to.isAny() || to.isNullableAny()) expression else irImplicitCast(expression, to)

-    private val signatureCache = mutableMapOf<IrFunctionSymbol, Method>()
-
    private val IrFunction.jvmMethod: Method
-        get() = signatureCache.getOrPut(symbol) { context.methodSignatureMapper.mapAsmMethod(this) }
+        get() = context.bridgeLoweringCache.computeJvmMethod(this)
+
+    internal class BridgeLoweringCache(private val context: JvmBackendContext) {
+        private val specialBridgeMethods = SpecialBridgeMethods(context)
+
+        // TODO: consider moving this cache out to the backend context and using it everywhere throughout the codegen.
+        // It might benefit performance, but can lead to confusing behavior if some declarations are changed along the way.
+        // For example, adding an override for a declaration whose signature is already cached can result in incorrect signature
+        // if its return type is a primitive type, and the new override's return type is an object type.
+        private val signatureCache = hashMapOf<IrFunctionSymbol, Method>()
+
+        fun computeJvmMethod(function: IrFunction): Method =
+            signatureCache.getOrPut(function.symbol) { context.methodSignatureMapper.mapAsmMethod(function) }
+
+        fun computeSpecialBridge(function: IrSimpleFunction): SpecialBridge? {
+            // Optimization: do not try to compute special bridge for irrelevant methods.
+            val correspondingProperty = function.correspondingPropertySymbol
+            if (correspondingProperty != null) {
+                if (correspondingProperty.owner.name !in specialBridgeMethods.specialPropertyNames) return null
+            } else {
+                if (function.name !in specialBridgeMethods.specialMethodNames) return null
+            }
+
+            val specialMethodInfo = specialBridgeMethods.getSpecialMethodInfo(function)
+            if (specialMethodInfo != null)
+                return SpecialBridge(
+                    function, computeJvmMethod(function), specialMethodInfo.needsGenericSignature, methodInfo = specialMethodInfo
+                )
+
+            val specialBuiltInInfo = specialBridgeMethods.getBuiltInWithDifferentJvmName(function)
+            if (specialBuiltInInfo != null)
+                return SpecialBridge(function, computeJvmMethod(function), specialBuiltInInfo.needsGenericSignature)
+
+            for (overridden in function.overriddenSymbols) {
+                val specialBridge = computeSpecialBridge(overridden.owner) ?: continue
+                if (!specialBridge.needsGenericSignature) return specialBridge
+
+                // Compute the substituted signature.
+                val erasedParameterCount = specialBridge.methodInfo?.argumentsToCheck ?: 0
+                val substitutedParameterTypes = function.valueParameters.mapIndexed { index, param ->
+                    if (index < erasedParameterCount) context.irBuiltIns.anyNType else param.type
+                }
+
+                val substitutedOverride = context.irFactory.buildFun {
+                    updateFrom(specialBridge.overridden)
+                    name = Name.identifier(specialBridge.signature.name)
+                    returnType = function.returnType
+                }.apply {
+                    // All existing special bridges only have value parameter types.
+                    valueParameters = function.valueParameters.zip(substitutedParameterTypes).map { (param, type) ->
+                        param.copyTo(this, IrDeclarationOrigin.BRIDGE, type = type)
+                    }
+                    overriddenSymbols = listOf(specialBridge.overridden.symbol)
+                    parent = function.parent
+                }
+
+                return specialBridge.copy(
+                    signature = computeJvmMethod(substitutedOverride),
+                    substitutedParameterTypes = substitutedParameterTypes,
+                    substitutedReturnType = function.returnType
+                )
+            }
+
+            return null
+        }
+    }
 }

 private fun IrDeclaration.comesFromJava() = parentAsClass.origin == IrDeclarationOrigin.IR_EXTERNAL_JAVA_DECLARATION_STUB