JVM_IR: Generate args check in existing methods for special bridge methods.

Fix the number of arguments checked (1 for getOrDefault and 2 for
remove).

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/SpecialBridgeMethods.kt

@@ -17,6 +17,8 @@ import org.jetbrains.kotlin.ir.util.parentAsClass
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 
+data class SpecialMethodWithDefaultInfo(val defaultValueGenerator: (IrSimpleFunction) -> IrExpression, val argumentsToCheck: Int)
+
 class SpecialBridgeMethods(val context: CommonBackendContext) {
     private data class SpecialMethodDescription(val kotlinFqClassName: FqName?, val name: Name, val arity: Int)
 
@@ -48,20 +50,20 @@ 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<SpecialMethodDescription, (IrSimpleFunction) -> IrExpression>(
+    private val SPECIAL_METHODS_WITH_DEFAULTS_MAP = mapOf(
-        makeDescription("kotlin.collections.Collection", "contains", 1) to ::constFalse,
+        makeDescription("kotlin.collections.Collection", "contains", 1) to SpecialMethodWithDefaultInfo(::constFalse, 1),
-        makeDescription("kotlin.collections.MutableCollection", "remove", 1) to ::constFalse,
+        makeDescription("kotlin.collections.MutableCollection", "remove", 1) to SpecialMethodWithDefaultInfo(::constFalse, 1),
-        makeDescription("kotlin.collections.Map", "containsKey", 1) to ::constFalse,
+        makeDescription("kotlin.collections.Map", "containsKey", 1) to SpecialMethodWithDefaultInfo(::constFalse, 1),
-        makeDescription("kotlin.collections.Map", "containsValue", 1) to ::constFalse,
+        makeDescription("kotlin.collections.Map", "containsValue", 1) to SpecialMethodWithDefaultInfo(::constFalse, 1),
-        makeDescription("kotlin.collections.MutableMap", "remove", 2) to ::constFalse,
+        makeDescription("kotlin.collections.MutableMap", "remove", 2) to SpecialMethodWithDefaultInfo(::constFalse, 2),
-        makeDescription("kotlin.collections.Map", "getOrDefault", 1) to ::getSecondArg,
+        makeDescription("kotlin.collections.Map", "getOrDefault", 2) to SpecialMethodWithDefaultInfo(::getSecondArg, 1),
-        makeDescription("kotlin.collections.Map", "get", 1) to ::constNull,
+        makeDescription("kotlin.collections.Map", "get", 1) to SpecialMethodWithDefaultInfo(::constNull, 1),
-        makeDescription("kotlin.collections.MutableMap", "remove", 1) to ::constNull,
+        makeDescription("kotlin.collections.MutableMap", "remove", 1) to SpecialMethodWithDefaultInfo(::constNull, 1),
-        makeDescription("kotlin.collections.List", "indexOf", 1) to ::constMinusOne,
+        makeDescription("kotlin.collections.List", "indexOf", 1) to SpecialMethodWithDefaultInfo(::constMinusOne, 1),
-        makeDescription("kotlin.collections.List", "lastIndexOf", 1) to ::constMinusOne
+        makeDescription("kotlin.collections.List", "lastIndexOf", 1) to SpecialMethodWithDefaultInfo(::constMinusOne, 1)
     )
 
-    fun findSpecialWithOverride(irFunction: IrSimpleFunction): Pair<IrSimpleFunction, (IrSimpleFunction) -> IrExpression>? {
+    fun findSpecialWithOverride(irFunction: IrSimpleFunction): Pair<IrSimpleFunction, SpecialMethodWithDefaultInfo>? {
         irFunction.allOverridden().forEach { overridden ->
             val description = overridden.toDescription()
             SPECIAL_METHODS_WITH_DEFAULTS_MAP[description]?.let {

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/BridgesConstruction.kt

@@ -13,10 +13,7 @@ import org.jetbrains.kotlin.backend.common.ir.copyTo
 import org.jetbrains.kotlin.backend.common.ir.copyTypeParametersFrom
 import org.jetbrains.kotlin.backend.common.ir.isMethodOfAny
 import org.jetbrains.kotlin.backend.common.ir.isSuspend
-import org.jetbrains.kotlin.backend.common.lower.SpecialBridgeMethods
+import org.jetbrains.kotlin.backend.common.lower.*
-import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
-import org.jetbrains.kotlin.backend.common.lower.irBlockBody
-import org.jetbrains.kotlin.backend.common.lower.irNot
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.ir.backend.js.JsLoweredDeclarationOrigin
 import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
@@ -67,7 +64,7 @@ class BridgesConstruction(val context: CommonBackendContext) : ClassLoweringPass
         if (function.isMethodOfAny())
             return
 
-        val (specialOverride: IrSimpleFunction?, specialOverrideValueGenerator) =
+        val (specialOverride: IrSimpleFunction?, specialOverrideInfo) =
             specialBridgeMethods.findSpecialWithOverride(function) ?: Pair(null, null)
 
         val specialOverrideSignature = specialOverride?.let { FunctionAndSignature(it) }
@@ -93,7 +90,7 @@ class BridgesConstruction(val context: CommonBackendContext) : ClassLoweringPass
 
             val bridge: IrDeclaration = when {
                 specialOverrideSignature == from ->
-                    createBridge(function, from.function, to.function, specialOverrideValueGenerator)
+                    createBridge(function, from.function, to.function, specialOverrideInfo)
 
                 else ->
                     createBridge(function, from.function, to.function, null)
@@ -109,7 +106,7 @@ class BridgesConstruction(val context: CommonBackendContext) : ClassLoweringPass
         function: IrSimpleFunction,
         bridge: IrSimpleFunction,
         delegateTo: IrSimpleFunction,
-        defaultValueGenerator: ((IrSimpleFunction) -> IrExpression)?
+        specialMethodInfo: SpecialMethodWithDefaultInfo?
     ): IrFunction {
 
         val origin =
@@ -144,12 +141,12 @@ class BridgesConstruction(val context: CommonBackendContext) : ClassLoweringPass
         }
 
         context.createIrBuilder(irFunction.symbol).irBlockBody(irFunction) {
-            if (defaultValueGenerator != null) {
+            if (specialMethodInfo != null) {
-                irFunction.valueParameters.forEach {
+                irFunction.valueParameters.take(specialMethodInfo.argumentsToCheck).forEach {
                     +irIfThen(
                         context.irBuiltIns.unitType,
                         irNot(irIs(irGet(it), delegateTo.valueParameters[it.index].type)),
-                        irReturn(defaultValueGenerator(irFunction))
+                        irReturn(specialMethodInfo.defaultValueGenerator(irFunction))
                     )
                 }
             }

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/BridgeLowering.kt

@@ -11,10 +11,7 @@ import org.jetbrains.kotlin.backend.common.bridges.findAllReachableDeclarations
 import org.jetbrains.kotlin.backend.common.bridges.findConcreteSuperDeclaration
 import org.jetbrains.kotlin.backend.common.bridges.generateBridges
 import org.jetbrains.kotlin.backend.common.ir.*
-import org.jetbrains.kotlin.backend.common.lower.SpecialBridgeMethods
+import org.jetbrains.kotlin.backend.common.lower.*
-import org.jetbrains.kotlin.backend.common.lower.allOverridden
-import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
-import org.jetbrains.kotlin.backend.common.lower.irNot
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
@@ -31,11 +28,15 @@ import org.jetbrains.kotlin.ir.declarations.impl.IrValueParameterImpl
 import org.jetbrains.kotlin.ir.descriptors.WrappedReceiverParameterDescriptor
 import org.jetbrains.kotlin.ir.descriptors.WrappedSimpleFunctionDescriptor
 import org.jetbrains.kotlin.ir.descriptors.WrappedValueParameterDescriptor
+import org.jetbrains.kotlin.ir.expressions.IrBlockBody
 import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrExpressionBody
 import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
-import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
+import org.jetbrains.kotlin.ir.expressions.impl.*
 import org.jetbrains.kotlin.ir.symbols.impl.IrSimpleFunctionSymbolImpl
 import org.jetbrains.kotlin.ir.symbols.impl.IrValueParameterSymbolImpl
+import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.types.isNullable
 import org.jetbrains.kotlin.ir.util.defaultType
 import org.jetbrains.kotlin.ir.util.fqNameWhenAvailable
 import org.jetbrains.kotlin.ir.util.isInterface
@@ -85,16 +86,14 @@ private class BridgeLowering(val context: JvmBackendContext) : ClassLoweringPass
             return
         }
 
-
         val irClass = irFunction.parentAsClass
         val ourSignature = irFunction.getJvmSignature()
         val ourMethodName = ourSignature.name
 
-        val (specialOverride, specialOverrideValueGenerator) =
+        val (specialOverride, specialOverrideInfo) =
             specialBridgeMethods.findSpecialWithOverride(irFunction) ?: Pair(null, null)
         val specialOverrideSignature = specialOverride?.getJvmSignature()
 
-
         var targetForCommonBridges = irFunction
 
         // Special case: fake override redirecting to an implementation with a different JVM name,
@@ -139,7 +138,7 @@ private class BridgeLowering(val context: JvmBackendContext) : ClassLoweringPass
             if (firstOverridden == null || firstOverriddenSignature!!.name != ourMethodName || getRenamedOverridden(firstOverridden) == null) {
                 addBridge(
                     irClass, targetForCommonBridges, renamer, signaturesToSkip,
-                    defaultValueGenerator = null,
+                    specialOverrideInfo = null,
                     isSpecial = true
                 )
             } else {
@@ -153,11 +152,17 @@ private class BridgeLowering(val context: JvmBackendContext) : ClassLoweringPass
         ) {
             addBridge(
                 irClass, targetForCommonBridges, specialOverride, signaturesToSkip,
-                specialOverrideValueGenerator,
+                specialOverrideInfo,
                 isSpecial = true
             )
         }
 
+        // If there is an existing function that would conflict with a special bridge signature, insert the special bridge
+        // code directly as a prelude in the existing method.
+        if (!irFunction.isFakeOverride && specialOverride != null && specialOverrideSignature == ourSignature) {
+            irFunction.rewriteSpecialMethodBody(specialOverrideInfo!!)
+        }
+
         val bridgeSignatures = generateBridges(
             FunctionHandleForIrFunction(irFunction),
             { handle -> SignatureWithSource(handle.irFunction.getJvmSignature(), handle.irFunction) }
@@ -167,7 +172,7 @@ private class BridgeLowering(val context: JvmBackendContext) : ClassLoweringPass
             val method = bridgeSignature.from.source
             addBridge(
                 irClass, targetForCommonBridges, method, signaturesToSkip,
-                defaultValueGenerator = null,
+                specialOverrideInfo = null,
                 isSpecial = targetForCommonBridges.isCollectionStub()
             )
         }
@@ -178,14 +183,14 @@ private class BridgeLowering(val context: JvmBackendContext) : ClassLoweringPass
         target: IrSimpleFunction,
         method: IrSimpleFunction,
         signaturesToSkip: MutableSet<Method>,
-        defaultValueGenerator: ((IrSimpleFunction) -> IrExpression)?,
+        specialOverrideInfo: SpecialMethodWithDefaultInfo?,
         isSpecial: Boolean
     ) {
         val signature = method.getJvmSignature()
         if (signature in signaturesToSkip) return
 
         val bridge = createBridgeHeader(irClass, target, method, isSpecial = isSpecial, isSynthetic = !isSpecial)
-        bridge.createBridgeBody(target, defaultValueGenerator, isSpecial)
+        bridge.createBridgeBody(target, specialOverrideInfo, isSpecial)
         irClass.declarations.add(bridge)
 
         // For lambda classes, we move override from the `invoke` function to its bridge. This will allow us to avoid boxing
@@ -257,9 +262,61 @@ private class BridgeLowering(val context: JvmBackendContext) : ClassLoweringPass
         }
     }
 
+    private fun IrStatementsBuilder<IrBlockBody>.addParameterTypeCheck(
+        parameter: IrValueParameter,
+        type: IrType,
+        defaultValueGenerator: ((IrSimpleFunction) -> IrExpression),
+        function: IrSimpleFunction
+    ) {
+        +irIfThen(
+            context.irBuiltIns.unitType,
+            irNot(irIs(irGet(parameter), type)),
+            irReturn(defaultValueGenerator(function))
+        )
+    }
+
+    private fun IrSimpleFunction.rewriteSpecialMethodBody(specialOverrideInfo: SpecialMethodWithDefaultInfo) {
+        val argumentsToCheck = valueParameters.take(specialOverrideInfo.argumentsToCheck)
+        val shouldGenerateParameterChecks = argumentsToCheck.any { !it.type.isNullable() }
+        if (shouldGenerateParameterChecks) {
+            // Rewrite the body to check if arguments have wrong type. If so, return the default value, otherwise,
+            // use the existing function body.
+            context.createIrBuilder(symbol).run {
+                body = irBlockBody {
+                    // Change the parameter types to be Any? so that null checks are not generated. The checks
+                    // we insert here make them superfluous.
+                    val variableMap = mutableMapOf<IrValueParameter, IrValueParameter>()
+                    argumentsToCheck.forEach {
+                        val parameterType = it.type
+                        if (!parameterType.isNullable()) {
+                            val newParameter = it.copyTo(this@rewriteSpecialMethodBody, type = context.irBuiltIns.anyNType)
+                            variableMap.put(valueParameters[it.index], newParameter)
+                            valueParameters[it.index] = newParameter
+                            addParameterTypeCheck(
+                                newParameter,
+                                parameterType,
+                                specialOverrideInfo.defaultValueGenerator,
+                                this@rewriteSpecialMethodBody
+                            )
+                        }
+                    }
+
+                    // Map usages of rewritten parameters in old body to the new parameters and paste in the body after the explicit
+                    // null check.
+                    body?.transform(VariableRemapper(variableMap), null)
+                    if (body is IrExpressionBody) {
+                        +(body as IrExpressionBody).expression
+                    } else {
+                        (body as IrBlockBody).statements.forEach { +it }
+                    }
+                }
+            }
+        }
+    }
+
     private fun IrSimpleFunction.createBridgeBody(
         target: IrSimpleFunction,
-        defaultValueGenerator: ((IrSimpleFunction) -> IrExpression)?,
+        specialMethodInfo: SpecialMethodWithDefaultInfo?,
         isSpecial: Boolean,
         invokeStatically: Boolean = false
     ) {
@@ -270,12 +327,13 @@ private class BridgeLowering(val context: JvmBackendContext) : ClassLoweringPass
 
         context.createIrBuilder(symbol).run {
             body = irBlockBody {
-                if (defaultValueGenerator != null) {
+                if (specialMethodInfo != null) {
-                    valueParameters.forEach {
+                    valueParameters.take(specialMethodInfo.argumentsToCheck).forEach {
-                        +irIfThen(
+                        addParameterTypeCheck(
-                            context.irBuiltIns.unitType,
+                            it,
-                            irNot(irIs(irGet(it), maybeOrphanedTarget.valueParameters[it.index].type)),
+                            maybeOrphanedTarget.valueParameters[it.index].type,
-                            irReturn(defaultValueGenerator(this@createBridgeBody))
+                            specialMethodInfo.defaultValueGenerator,
+                            this@createBridgeBody
                         )
                     }
                 }

compiler/testData/codegen/box/builtinStubMethods/mapGetOrDefault/noTypeSafeBridge.kt

@@ -4,10 +4,6 @@
 // FULL_JDK
 // WITH_RUNTIME
 
-// Bridges are not generated because their signatures would conflict. The logic
-// should be inserted directly into existing methods, but this is not implemented.
-// IGNORE_BACKEND: JVM_IR
-
 class A : MutableMap<Any, Any?> {
     override val entries: MutableSet<MutableMap.MutableEntry<Any, Any?>>
         get() = throw UnsupportedOperationException()

compiler/testData/codegen/box/builtinStubMethods/mapGetOrDefault/typeSafeBridge.kt

@@ -1,6 +1,5 @@
 // IGNORE_BACKEND_FIR: JVM_IR
 // SKIP_JDK6
-// IGNORE_BACKEND: JVM_IR
 // TARGET_BACKEND: JVM
 // FULL_JDK
 // WITH_RUNTIME

compiler/testData/codegen/box/builtinStubMethods/mapGetOrDefault/typeSafeBridgeNotNullAny.kt

@@ -1,6 +1,5 @@
 // IGNORE_BACKEND_FIR: JVM_IR
 // SKIP_JDK6
-// IGNORE_BACKEND: JVM_IR
 // TARGET_BACKEND: JVM
 // FULL_JDK
 // WITH_RUNTIME

compiler/testData/codegen/box/builtinStubMethods/mapRemove/typeSafeBridgeNotNullAny.kt

@@ -1,6 +1,5 @@
 // IGNORE_BACKEND_FIR: JVM_IR
 // SKIP_JDK6
-// IGNORE_BACKEND: JVM_IR
 // TARGET_BACKEND: JVM
 // FULL_JDK
 // WITH_RUNTIME

compiler/testData/codegen/box/specialBuiltins/emptyMap.kt

@@ -1,7 +1,4 @@
 // IGNORE_BACKEND_FIR: JVM_IR
-// Bridges are not generated because their signatures would conflict. The logic
-// should be inserted directly into existing methods, but this is not implemented.
-// IGNORE_BACKEND: JVM_IR
 
 private object EmptyMap : Map<Any, Nothing> {
     override val size: Int get() = 0

compiler/testData/codegen/box/specialBuiltins/notEmptyListAny.kt

@@ -1,5 +1,4 @@
 // IGNORE_BACKEND_FIR: JVM_IR
-// IGNORE_BACKEND: JVM_IR
 // IGNORE_BACKEND: JS
 
 private object NotEmptyList : MutableList<Any> {

compiler/testData/codegen/box/specialBuiltins/notEmptyMap.kt

@@ -1,5 +1,4 @@
 // IGNORE_BACKEND_FIR: JVM_IR
-// IGNORE_BACKEND: JVM_IR
 // IGNORE_BACKEND: JS
 
 private object NotEmptyMap : MutableMap<Any, Any> {

compiler/testData/writeSignature/java8/mutableMapRemove.kt

@@ -1,5 +1,8 @@
 //FULL_JDK
 
+// In BridgeLowering, no concrete overwrite is found for getOrDefault.
+// IGNORE_BACKEND: JVM_IR
+
 class KotlinMap1<K, V> : java.util.AbstractMap<K, V>() {
     override val entries: MutableSet<MutableMap.MutableEntry<K, V>>
         get() = throw UnsupportedOperationException()