JVM IR: Avoid boxing in inline class equality

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/intrinsics/Equals.kt

@@ -9,12 +9,15 @@ import com.intellij.psi.tree.IElementType
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.codegen.*
 import org.jetbrains.kotlin.codegen.AsmUtil
+import org.jetbrains.kotlin.codegen.AsmUtil.genAreEqualCall
 import org.jetbrains.kotlin.codegen.AsmUtil.isPrimitive
 import org.jetbrains.kotlin.codegen.StackValue
 import org.jetbrains.kotlin.codegen.intrinsics.IntrinsicMethods
+import org.jetbrains.kotlin.codegen.pseudoInsns.fakeAlwaysFalseIfeq
+import org.jetbrains.kotlin.codegen.pseudoInsns.fakeAlwaysTrueIfeq
 import org.jetbrains.kotlin.ir.expressions.IrFunctionAccessExpression
 import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
-import org.jetbrains.kotlin.ir.types.classOrNull
+import org.jetbrains.kotlin.ir.types.isNullable
 import org.jetbrains.kotlin.ir.types.toKotlinType
 import org.jetbrains.kotlin.ir.util.isNullConst
 import org.jetbrains.kotlin.lexer.KtTokens
@@ -28,6 +31,12 @@ import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter
 
 class Equals(val operator: IElementType) : IntrinsicMethod() {
+    private class BooleanConstantFalseCheck(val value: PromisedValue) : BooleanValue(value.codegen) {
+        override fun materialize() = value.discard().let { mv.iconst(0) }
+        override fun jumpIfFalse(target: Label) = value.discard().let { mv.fakeAlwaysTrueIfeq(target) }
+        override fun jumpIfTrue(target: Label) = value.discard().let { mv.fakeAlwaysFalseIfeq(target) }
+    }
+
     private class BooleanNullCheck(val value: PromisedValue) : BooleanValue(value.codegen) {
         override fun jumpIfFalse(target: Label) = value.materialize().also { mv.ifnonnull(target) }
         override fun jumpIfTrue(target: Label) = value.materialize().also { mv.ifnull(target) }
@@ -36,25 +45,31 @@ class Equals(val operator: IElementType) : IntrinsicMethod() {
     override fun invoke(expression: IrFunctionAccessExpression, codegen: ExpressionCodegen, data: BlockInfo): PromisedValue? {
         val (a, b) = expression.receiverAndArgs()
         if (a.isNullConst() || b.isNullConst()) {
-            return BooleanNullCheck(if (a.isNullConst()) b.accept(codegen, data) else a.accept(codegen, data))
+            val value = if (a.isNullConst()) b.accept(codegen, data) else a.accept(codegen, data)
+            return if (a.type.isNullable() && b.type.isNullable())
+                BooleanNullCheck(value)
+            else
+                BooleanConstantFalseCheck(value)
         }
 
         val leftType = with(codegen) { a.asmType }
         val rightType = with(codegen) { b.asmType }
         val opToken = expression.origin
         val useEquals = opToken !== IrStatementOrigin.EQEQEQ && opToken !== IrStatementOrigin.EXCLEQEQ &&
-                (a.type.classOrNull?.owner?.isInline == true || b.type.classOrNull?.owner?.isInline == true ||
                 // `==` is `equals` for objects and floating-point numbers. In the latter case, the difference
                 // is that `equals` is a total order (-0 < +0 and NaN == NaN) and `===` is IEEE754-compliant.
-                !isPrimitive(leftType) || leftType != rightType || leftType == Type.FLOAT_TYPE || leftType == Type.DOUBLE_TYPE)
+                (!isPrimitive(leftType) || leftType != rightType || leftType == Type.FLOAT_TYPE || leftType == Type.DOUBLE_TYPE)
-        val operandType = if (!isPrimitive(leftType) || useEquals) AsmTypes.OBJECT_TYPE else leftType
+        return if (useEquals) {
-        val aValue = a.accept(codegen, data).coerce(operandType, a.type).materialized
+            a.accept(codegen, data).coerce(AsmTypes.OBJECT_TYPE, codegen.context.irBuiltIns.anyNType).materialize()
-        val bValue = b.accept(codegen, data).coerce(operandType, b.type).materialized
+            b.accept(codegen, data).coerce(AsmTypes.OBJECT_TYPE, codegen.context.irBuiltIns.anyNType).materialize()
-        if (useEquals) {
+            genAreEqualCall(codegen.mv)
-            AsmUtil.genAreEqualCall(codegen.mv)
+            MaterialValue(codegen, Type.BOOLEAN_TYPE, codegen.context.irBuiltIns.booleanType)
-            return MaterialValue(codegen, Type.BOOLEAN_TYPE, codegen.context.irBuiltIns.booleanType)
+        } else {
+            val operandType = if (!isPrimitive(leftType)) AsmTypes.OBJECT_TYPE else leftType
+            val aValue = a.accept(codegen, data).coerce(operandType, a.type).materialized
+            val bValue = b.accept(codegen, data).coerce(operandType, b.type).materialized
+            BooleanComparison(operator, aValue, bValue)
         }
-        return BooleanComparison(operator, aValue, bValue)
     }
 }
 

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

@@ -29,6 +29,8 @@ import org.jetbrains.kotlin.ir.expressions.impl.IrSetVariableImpl
 import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
 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.util.*
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
@@ -246,14 +248,85 @@ private class JvmInlineClassLowering(private val context: JvmBackendContext) : F
             putValueArgument(0, argument)
         }
 
-    override fun visitCall(expression: IrCall): IrExpression {
+    private fun IrExpression.coerceToUnboxed() =
-        val function = expression.symbol.owner
+        coerceInlineClasses(this, this.type, this.type.unboxInlineClass())
-        if (!function.isInlineClassFieldGetter)
+
-            return super.visitCall(expression)
+    private fun IrClass.getEqualsImpl(): IrFunction {
-        val arg = expression.dispatchReceiver!!.transform(this, null)
+        val equals = functions.single { it.name.asString() == "equals" && it.overriddenSymbols.isNotEmpty() }
-        return coerceInlineClasses(arg, function.dispatchReceiverParameter!!.type, expression.type)
+        return manager.getReplacementFunction(equals)!!.function
     }
 
+    private fun IrBuilderWithScope.specializeEqualsCall(left: IrExpression, right: IrExpression): IrExpression? {
+        // There's already special handling for null-comparisons in Equals.kt.
+        // We cannot specialize calls for which the first argument is already boxed.
+        if (left.isNullConst() || right.isNullConst() || left.type.unboxInlineClass() == left.type)
+            return null
+
+        // Unsigned types use primitive comparisons.
+        val rightIsUnboxed = right.type.unboxInlineClass() !== right.type
+        if (left.type.isUnsigned() && right.type.isUnsigned() && rightIsUnboxed)
+            return irEquals(left.coerceToUnboxed(), right.coerceToUnboxed())
+
+        val equalsMethod = if (rightIsUnboxed)
+            manager.getSpecializedEqualsMethod(left.type.classOrNull!!.owner, context.irBuiltIns)
+        else
+            left.type.classOrNull!!.owner.getEqualsImpl()
+
+        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 ->
+                irLetS(right) { rightVal ->
+                    val equalsCall = irCall(equalsMethod).apply {
+                        putValueArgument(0, irGet(leftVal.owner))
+                        putValueArgument(1, irGet(rightVal.owner))
+                    }
+
+                    val equalsRight = if (rightNullCheck) {
+                        irIfNull(context.irBuiltIns.booleanType, irGet(rightVal.owner), irFalse(), equalsCall)
+                    } else {
+                        equalsCall
+                    }
+
+                    if (leftNullCheck) {
+                        irIfNull(context.irBuiltIns.booleanType, irGet(leftVal.owner), irEqualsNull(irGet(rightVal.owner)), equalsRight)
+                    } else {
+                        equalsRight
+                    }
+                }
+            }
+        } else {
+            irCall(equalsMethod).apply {
+                putValueArgument(0, left)
+                putValueArgument(1, right)
+            }
+        }
+    }
+
+    override fun visitCall(expression: IrCall): IrExpression =
+        when {
+            // Getting the underlying field of an inline class merely changes the IR type,
+            // since the underlying representations are the same.
+            expression.symbol.owner.isInlineClassFieldGetter -> {
+                val arg = expression.dispatchReceiver!!.transform(this, null)
+                coerceInlineClasses(arg, expression.symbol.owner.dispatchReceiverParameter!!.type, expression.type)
+            }
+            // Specialize calls to equals with at least one inline class argument to avoid boxing.
+            expression.isInlineClassEqEq -> {
+                expression.transformChildrenVoid()
+                context.createIrBuilder(currentScope!!.scope.scopeOwnerSymbol, expression.startOffset, expression.endOffset)
+                    .specializeEqualsCall(expression.getValueArgument(0)!!, expression.getValueArgument(1)!!)
+                    ?: expression
+            }
+            else ->
+                super.visitCall(expression)
+        }
+
+    private val IrCall.isInlineClassEqEq: Boolean
+        get() = symbol == context.irBuiltIns.eqeqSymbol &&
+                (getValueArgument(0)?.type?.classOrNull?.owner?.isInline == true ||
+                        getValueArgument(1)?.type?.classOrNull?.owner?.isInline == true)
+
     override fun visitGetField(expression: IrGetField): IrExpression {
         val field = expression.symbol.owner
         val parent = field.parent

compiler/testData/codegen/box/inlineClasses/inlineClassWithCustomEquals.kt

@@ -1,3 +1,4 @@
+// IGNORE_BACKEND: JVM_IR
 // !LANGUAGE: +InlineClasses
 
 @file:Suppress("RESERVED_MEMBER_INSIDE_INLINE_CLASS")