JVM IR: Optimize inline class equality with nullable arguments

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

@@ -32,6 +32,7 @@ import org.jetbrains.kotlin.ir.symbols.IrVariableSymbol
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.types.classOrNull
 import org.jetbrains.kotlin.ir.types.isNullable
+import org.jetbrains.kotlin.ir.types.makeNotNull
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.resolve.DescriptorUtils
@@ -253,55 +254,64 @@ private class JvmInlineClassLowering(private val context: JvmBackendContext) : F
     private fun IrExpression.coerceToUnboxed() =
         coerceInlineClasses(this, this.type, this.type.unboxInlineClass())
 
-    private fun IrClass.getEqualsImpl(): IrFunction {
+    // Precondition: left has an inline class type, but may not be unboxed
-        val equals = functions.single { it.name.asString() == "equals" && it.overriddenSymbols.isNotEmpty() }
-        return manager.getReplacementFunction(equals)!!.function
-    }
-
     private fun IrBuilderWithScope.specializeEqualsCall(left: IrExpression, right: IrExpression): IrExpression? {
         // There's already special handling for null-comparisons in the Equals intrinsic.
-        // We cannot specialize calls for which the first argument is already boxed.
+        if (left.isNullConst() || right.isNullConst())
-        if (left.isNullConst() || right.isNullConst() || left.type.unboxInlineClass() == left.type)
             return null
 
-        // Unsigned types use primitive comparisons.
+        // We don't specialize calls when both arguments are boxed.
-        val rightIsUnboxed = right.type.unboxInlineClass() !== right.type
+        val leftIsUnboxed = left.type.unboxInlineClass() != left.type
+        val rightIsUnboxed = right.type.unboxInlineClass() != right.type
+        if (!leftIsUnboxed && !rightIsUnboxed)
+            return null
+
+        // Precondition: left is an unboxed inline class type
+        fun equals(left: IrExpression, right: IrExpression): IrExpression {
+            // Unsigned types use primitive comparisons
             if (left.type.isUnsigned() && right.type.isUnsigned() && rightIsUnboxed)
                 return irEquals(left.coerceToUnboxed(), right.coerceToUnboxed())
 
-        val equalsMethod = if (rightIsUnboxed)
+            val klass = left.type.classOrNull!!.owner
-            manager.getSpecializedEqualsMethod(left.type.classOrNull!!.owner, context.irBuiltIns)
+            val equalsMethod = if (rightIsUnboxed) {
-        else
+                manager.getSpecializedEqualsMethod(klass, context.irBuiltIns)
-            left.type.classOrNull!!.owner.getEqualsImpl()
+            } else {
+                val equals = klass.functions.single { it.name.asString() == "equals" && it.overriddenSymbols.isNotEmpty() }
+                manager.getReplacementFunction(equals)!!.function
+            }
+
+            return irCall(equalsMethod).apply {
+                putValueArgument(0, left)
+                putValueArgument(1, right)
+            }
+        }
 
         val leftNullCheck = left.type.isNullable()
         val rightNullCheck = rightIsUnboxed && right.type.isNullable() // equals-impl has a nullable second argument
         return if (leftNullCheck || rightNullCheck) {
-            irLetS(left) { leftVal ->
+            irBlock {
-                irLetS(right) { rightVal ->
+                val leftVal = if (left is IrGetValue) left.symbol.owner else irTemporary(left)
-                    val equalsCall = irCall(equalsMethod).apply {
+                val rightVal = if (right is IrGetValue) right.symbol.owner else irTemporary(right)
-                        putValueArgument(0, irGet(leftVal.owner))
+
-                        putValueArgument(1, irGet(rightVal.owner))
+                val equalsCall = equals(
-                    }
+                    if (leftNullCheck) irImplicitCast(irGet(leftVal), left.type.makeNotNull()) else irGet(leftVal),
+                    if (rightNullCheck) irImplicitCast(irGet(rightVal), right.type.makeNotNull()) else irGet(rightVal)
+                )
 
                 val equalsRight = if (rightNullCheck) {
-                        irIfNull(context.irBuiltIns.booleanType, irGet(rightVal.owner), irFalse(), equalsCall)
+                    irIfNull(context.irBuiltIns.booleanType, irGet(rightVal), irFalse(), equalsCall)
                 } else {
                     equalsCall
                 }
 
                 if (leftNullCheck) {
-                        irIfNull(context.irBuiltIns.booleanType, irGet(leftVal.owner), irEqualsNull(irGet(rightVal.owner)), equalsRight)
+                    +irIfNull(context.irBuiltIns.booleanType, irGet(leftVal), irEqualsNull(irGet(rightVal)), equalsRight)
                 } else {
-                        equalsRight
+                    +equalsRight
-                    }
                 }
             }
         } else {
-            irCall(equalsMethod).apply {
+            equals(left, right)
-                putValueArgument(0, left)
-                putValueArgument(1, right)
-            }
         }
     }