Fuse primitive equality with safe call to avoid boxing

In code like 'a?.b == 42', we can immediately generate equality
comparison result when receiver is null (false for '==', true for '!='),
since the primitive value is definitely non-null.
Otherwise unnecessary boxing/unboxing is generated to handle possibly
null result of 'a?.b'.

compiler/backend/src/org/jetbrains/kotlin/codegen/BranchedValue.kt

@@ -209,33 +209,9 @@ class NumberCompare(
         right: StackValue
 ) : BranchedValue(left, right, operandType, NumberCompare.getNumberCompareOpcode(opToken)) {
 
-    override fun patchOpcode(opcode: Int, v: InstructionAdapter): Int {
+    override fun patchOpcode(opcode: Int, v: InstructionAdapter): Int =
-        // Opcode takes one int operand from the stack
+            patchOpcode(opcode, v, opToken, operandType)
-        assert(opcode in IFEQ..IFLE) {
-            "Opcode for comparing must be in range ${IFEQ..IFLE}, but $opcode was found"
-        }
 
-        return when (operandType) {
-            Type.FLOAT_TYPE, Type.DOUBLE_TYPE -> {
-                if (opToken == KtTokens.GT || opToken == KtTokens.GTEQ) {
-                    v.cmpl(operandType)
-                }
-                else {
-                    v.cmpg(operandType)
-                }
-
-                opcode
-            }
-            Type.LONG_TYPE -> {
-                v.lcmp()
-
-                opcode
-            }
-            else -> {
-                opcode + (IF_ICMPEQ - IFEQ)
-            }
-        }
-    }
     companion object {
         fun getNumberCompareOpcode(opToken: IElementType): Int {
             return when (opToken) {
@@ -250,6 +226,27 @@ class NumberCompare(
                 }
             }
         }
+
+        fun patchOpcode(opcode: Int, v: InstructionAdapter, opToken: IElementType, operandType: Type): Int {
+            assert(opcode in IFEQ..IFLE) {
+                "Opcode for comparing must be in range ${IFEQ..IFLE}, but $opcode was found"
+            }
+            return when (operandType) {
+                Type.FLOAT_TYPE, Type.DOUBLE_TYPE -> {
+                    if (opToken == KtTokens.GT || opToken == KtTokens.GTEQ)
+                        v.cmpl(operandType)
+                    else
+                        v.cmpg(operandType)
+                    opcode
+                }
+                Type.LONG_TYPE -> {
+                    v.lcmp()
+                    opcode
+                }
+                else ->
+                    opcode + (IF_ICMPEQ - IFEQ)
+            }
+        }
     }
 }
 
@@ -270,3 +267,91 @@ class ObjectCompare(
         }
     }
 }
+
+abstract class SafeCallFusedWithPrimitiveEqualityBase(
+        val opToken: IElementType,
+        operandType: Type,
+        left: StackValue,
+        right: StackValue
+) : BranchedValue(left, right, operandType, NumberCompare.getNumberCompareOpcode(opToken)) {
+    private val trueIfEqual = opToken == KtTokens.EQEQ || opToken == KtTokens.EQEQEQ
+
+    protected abstract fun cleanupOnNullReceiver(v: InstructionAdapter)
+
+    override fun patchOpcode(opcode: Int, v: InstructionAdapter): Int =
+            NumberCompare.patchOpcode(opcode, v, opToken, operandType)
+
+    override fun condJump(jumpLabel: Label, v: InstructionAdapter, jumpIfFalse: Boolean) {
+        val endLabel = Label()
+
+        arg1.put(operandType, v)
+        arg2!!.put(operandType, v)
+        v.visitJumpInsn(patchOpcode(if (jumpIfFalse) opcode else negatedOperations[opcode]!!, v), jumpLabel)
+        v.goTo(endLabel)
+
+        cleanupOnNullReceiver(v)
+        if (jumpIfFalse == trueIfEqual) {
+            v.goTo(jumpLabel)
+        }
+
+        v.mark(endLabel)
+    }
+
+    override fun putSelector(type: Type, v: InstructionAdapter) {
+        val falseLabel = Label()
+        val endLabel = Label()
+
+        arg1.put(operandType, v)
+        arg2!!.put(operandType, v)
+        v.visitJumpInsn(patchOpcode(opcode, v), falseLabel)
+
+        if (!trueIfEqual) {
+            val trueLabel = Label()
+            v.goTo(trueLabel)
+            cleanupOnNullReceiver(v)
+            v.mark(trueLabel)
+        }
+
+        v.iconst(1)
+        v.goTo(endLabel)
+
+        if (trueIfEqual) {
+            cleanupOnNullReceiver(v)
+        }
+
+        v.mark(falseLabel)
+        v.iconst(0)
+
+        v.mark(endLabel)
+        coerceTo(type, v)
+    }
+}
+
+class SafeCallToPrimitiveEquality(
+        opToken: IElementType,
+        operandType: Type,
+        left: StackValue,
+        right: StackValue,
+        val safeReceiverType: Type,
+        val safeReceiverIsNull: Label
+) : SafeCallFusedWithPrimitiveEqualityBase(opToken, operandType, left, right) {
+    override fun cleanupOnNullReceiver(v: InstructionAdapter) {
+        v.mark(safeReceiverIsNull)
+        AsmUtil.pop(v, safeReceiverType)
+    }
+}
+
+class PrimitiveToSafeCallEquality(
+        opToken: IElementType,
+        operandType: Type,
+        left: StackValue,
+        right: StackValue,
+        val safeReceiverType: Type,
+        val safeReceiverIsNull: Label
+) : SafeCallFusedWithPrimitiveEqualityBase(opToken, operandType, left, right) {
+    override fun cleanupOnNullReceiver(v: InstructionAdapter) {
+        v.mark(safeReceiverIsNull)
+        AsmUtil.pop(v, safeReceiverType)
+        AsmUtil.pop(v, arg1.type)
+    }
+}

compiler/backend/src/org/jetbrains/kotlin/codegen/ExpressionCodegen.java

@@ -2927,6 +2927,13 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
             return genCmpWithZero(left, opToken);
         }
 
+        if (left instanceof KtSafeQualifiedExpression && isPrimitive(rightType)) {
+            return genCmpSafeCallToPrimitive((KtSafeQualifiedExpression) left, right, rightType, opToken);
+        }
+        if (isPrimitive(leftType) && right instanceof KtSafeQualifiedExpression) {
+            return genCmpPrimitiveToSafeCall(left, leftType, (KtSafeQualifiedExpression) right, opToken);
+        }
+
         if (isPrimitive(leftType) != isPrimitive(rightType)) {
             leftType = boxType(leftType);
             rightType = boxType(rightType);
@@ -2941,6 +2948,40 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
         return genEqualsForExpressionsPreferIEEE754Arithmetic(left, right, opToken, leftType, rightType, null);
     }
 
+    private StackValue genCmpPrimitiveToSafeCall(
+            @NotNull KtExpression left,
+            @NotNull Type leftType,
+            @NotNull KtSafeQualifiedExpression right,
+            @NotNull IElementType opToken
+    ) {
+        Label rightIsNull = new Label();
+        return new PrimitiveToSafeCallEquality(
+                opToken,
+                leftType,
+                genLazy(left, leftType),
+                generateSafeQualifiedExpression(right, rightIsNull),
+                expressionType(right.getReceiverExpression()),
+                rightIsNull
+        );
+    }
+
+    private StackValue genCmpSafeCallToPrimitive(
+            @NotNull KtSafeQualifiedExpression left,
+            @NotNull KtExpression right,
+            @NotNull Type rightType,
+            @NotNull IElementType opToken
+    ) {
+        Label leftIsNull = new Label();
+        return new SafeCallToPrimitiveEquality(
+                opToken,
+                rightType,
+                generateSafeQualifiedExpression(left, leftIsNull),
+                genLazy(right, rightType),
+                expressionType(left.getReceiverExpression()),
+                leftIsNull
+        );
+    }
+
     /*tries to use IEEE 754 arithmetic*/
     private StackValue genEqualsForExpressionsPreferIEEE754Arithmetic(
             @Nullable KtExpression left,

compiler/testData/codegen/box/safeCall/primitiveEqSafeCall.kt

@@ -0,0 +1,56 @@
+fun Long.id() = this
+
+fun String.drop2() = if (length >= 2) subSequence(2, length) else null
+
+fun String.anyLength(): Any = length
+
+
+fun doSimple(s: String?) = 3 == s?.length
+
+fun doLongReceiver(x: Long) = 3L == x?.id()
+
+fun doChain(s: String?) = 1 == s?.drop2()?.length
+
+fun doIf(s: String?) =
+        if (1 == s?.length) "A" else "B"
+
+fun doCmpWithAny(s: String?) =
+        3 == s?.anyLength()
+
+fun doIfNot(s: String?) =
+        if (!(1 == s?.length)) "A" else "B"
+
+fun doIfNotNot(s: String?) =
+        if (!!(1 == s?.length)) "A" else "B"
+
+
+fun box(): String = when {
+    doSimple(null) -> "failed 1"
+    doSimple("1") -> "failed 2"
+    !doSimple("123") -> "failed 3"
+
+    doLongReceiver(2L) -> "failed 4"
+    !doLongReceiver(3L) -> "failed 5"
+
+    doChain(null) -> "failed 6"
+    doChain("1") -> "failed 7"
+    !doChain("123") -> "failed 7"
+
+    doIf("1") != "A" -> "failed 8"
+    doIf("123") != "B" -> "failed 9"
+    doIf(null) != "B" -> "failed 10"
+
+    doCmpWithAny(null) -> "failed 11"
+    doCmpWithAny("1") -> "failed 12"
+    !doCmpWithAny("123") -> "failed 13"
+
+    doIfNot("1") != "B" -> "failed 8"
+    doIfNot("123") != "A" -> "failed 9"
+    doIfNot(null) != "A" -> "failed 10"
+
+    doIfNotNot("1") != "A" -> "failed 8"
+    doIfNotNot("123") != "B" -> "failed 9"
+    doIfNotNot(null) != "B" -> "failed 10"
+
+    else -> "OK"
+}

compiler/testData/codegen/box/safeCall/primitiveNotEqSafeCall.kt

@@ -0,0 +1,43 @@
+fun Long.id() = this
+
+fun String.drop2() = if (length >= 2) subSequence(2, length) else null
+
+fun String.anyLength(): Any = length
+
+
+fun doSimple(s: String?) = 3 != s?.length
+
+fun doLongReceiver(x: Long) = 3L != x?.id()
+
+fun doChain(s: String?) = 1 != s?.drop2()?.length
+
+fun doIf(s: String?) =
+        if (1 != s?.length) "A" else "B"
+
+fun doCmpWithAny(s: String?) =
+        3 != s?.anyLength()
+
+
+fun box(): String = when {
+    !doSimple(null) -> "failed 1"
+    !doSimple("1") -> "failed 2"
+    doSimple("123") -> "failed 3"
+
+    !doLongReceiver(2L) -> "failed 4"
+    doLongReceiver(3L) -> "failed 5"
+
+    !doChain(null) -> "failed 6"
+    !doChain("1") -> "failed 7"
+    doChain("123") -> "failed 7"
+
+    doIf("1") == "A" -> "failed 8"
+    doIf("123") == "B" -> "failed 9"
+    doIf(null) == "B" -> "failed 10"
+
+    !doCmpWithAny(null) -> "failed 11"
+    !doCmpWithAny("1") -> "failed 12"
+    doCmpWithAny("123") -> "failed 13"
+
+
+    else -> "OK"
+}

compiler/testData/codegen/box/safeCall/safeCallEqPrimitive.kt

@@ -0,0 +1,55 @@
+fun Long.id() = this
+
+fun String.drop2() = if (length >= 2) subSequence(2, length) else null
+
+fun String.anyLength(): Any = length
+
+
+fun doSimple(s: String?) = s?.length == 3
+
+fun doLongReceiver(x: Long) = x?.id() == 3L
+
+fun doChain(s: String?) = s?.drop2()?.length == 1
+
+fun doIf(s: String?) =
+        if (s?.length == 1) "A" else "B"
+
+fun doCmpWithAny(s: String?) =
+        s?.anyLength() == 3
+
+fun doIfNot(s: String?) =
+        if (!(s?.length == 1)) "A" else "B"
+
+fun doIfNotNot(s: String?) =
+        if (!!(s?.length == 1)) "A" else "B"
+
+fun box(): String = when {
+    doSimple(null) -> "failed 1"
+    doSimple("1") -> "failed 2"
+    !doSimple("123") -> "failed 3"
+
+    doLongReceiver(2L) -> "failed 4"
+    !doLongReceiver(3L) -> "failed 5"
+
+    doChain(null) -> "failed 6"
+    doChain("1") -> "failed 7"
+    !doChain("123") -> "failed 7"
+
+    doIf("1") != "A" -> "failed 8"
+    doIf("123") != "B" -> "failed 9"
+    doIf(null) != "B" -> "failed 10"
+
+    doCmpWithAny(null) -> "failed 11"
+    doCmpWithAny("1") -> "failed 12"
+    !doCmpWithAny("123") -> "failed 13"
+
+    doIfNot("1") != "B" -> "failed 8"
+    doIfNot("123") != "A" -> "failed 9"
+    doIfNot(null) != "A" -> "failed 10"
+
+    doIfNotNot("1") != "A" -> "failed 8"
+    doIfNotNot("123") != "B" -> "failed 9"
+    doIfNotNot(null) != "B" -> "failed 10"
+
+    else -> "OK"
+}

compiler/testData/codegen/box/safeCall/safeCallNotEqPrimitive.kt

@@ -0,0 +1,43 @@
+fun Long.id() = this
+
+fun String.drop2() = if (length >= 2) subSequence(2, length) else null
+
+fun String.anyLength(): Any = length
+
+
+fun doSimple(s: String?) = s?.length != 3
+
+fun doLongReceiver(x: Long) = x?.id() != 3L
+
+fun doChain(s: String?) = s?.drop2()?.length != 1
+
+fun doIf(s: String?) =
+        if (s?.length != 1) "A" else "B"
+
+fun doCmpWithAny(s: String?) =
+        s?.anyLength() != 3
+
+
+fun box(): String = when {
+    !doSimple(null) -> "failed 1"
+    !doSimple("1") -> "failed 2"
+    doSimple("123") -> "failed 3"
+
+    !doLongReceiver(2L) -> "failed 4"
+    doLongReceiver(3L) -> "failed 5"
+
+    !doChain(null) -> "failed 6"
+    !doChain("1") -> "failed 7"
+    doChain("123") -> "failed 7"
+
+    doIf("1") == "A" -> "failed 8"
+    doIf("123") == "B" -> "failed 9"
+    doIf(null) == "B" -> "failed 10"
+
+    !doCmpWithAny(null) -> "failed 11"
+    !doCmpWithAny("1") -> "failed 12"
+    doCmpWithAny("123") -> "failed 13"
+
+
+    else -> "OK"
+}

compiler/testData/codegen/bytecodeText/boxingOptimization/safeCallToPrimitiveEquality.kt

@@ -0,0 +1,23 @@
+fun Long.id() = this
+
+fun String.drop2() = if (length >= 2) subSequence(2, length) else null
+
+fun doSimple1(s: String?) = s?.length == 3
+
+fun doLongReceiver1(x: Long) = x?.id() == 3L
+
+fun doChain1(s: String?) = s?.drop2()?.length == 1
+
+fun doIf1(s: String?) =
+        if (s?.length == 1) "A" else "B"
+
+fun doSimple2(s: String?) = 3 == s?.length
+
+fun doLongReceiver2(x: Long) = 3L == x?.id()
+
+fun doChain2(s: String?) = 1 == s?.drop2()?.length
+
+fun doIf2(s: String?) =
+        if (1 == s?.length) "A" else "B"
+
+// 0 valueOf

compiler/tests-ir-jvm/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -16616,6 +16616,30 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             doTest(fileName);
         }
 
+        @TestMetadata("primitiveEqSafeCall.kt")
+        public void testPrimitiveEqSafeCall() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/safeCall/primitiveEqSafeCall.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("primitiveNotEqSafeCall.kt")
+        public void testPrimitiveNotEqSafeCall() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/safeCall/primitiveNotEqSafeCall.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("safeCallEqPrimitive.kt")
+        public void testSafeCallEqPrimitive() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/safeCall/safeCallEqPrimitive.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("safeCallNotEqPrimitive.kt")
+        public void testSafeCallNotEqPrimitive() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/safeCall/safeCallNotEqPrimitive.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("safeCallOnLong.kt")
         public void testSafeCallOnLong() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/safeCall/safeCallOnLong.kt");

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -16616,6 +16616,30 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("primitiveEqSafeCall.kt")
+        public void testPrimitiveEqSafeCall() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/safeCall/primitiveEqSafeCall.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("primitiveNotEqSafeCall.kt")
+        public void testPrimitiveNotEqSafeCall() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/safeCall/primitiveNotEqSafeCall.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("safeCallEqPrimitive.kt")
+        public void testSafeCallEqPrimitive() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/safeCall/safeCallEqPrimitive.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("safeCallNotEqPrimitive.kt")
+        public void testSafeCallNotEqPrimitive() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/safeCall/safeCallNotEqPrimitive.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("safeCallOnLong.kt")
         public void testSafeCallOnLong() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/safeCall/safeCallOnLong.kt");

compiler/tests/org/jetbrains/kotlin/codegen/BytecodeTextTestGenerated.java

@@ -527,6 +527,12 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
             doTest(fileName);
         }
 
+        @TestMetadata("safeCallToPrimitiveEquality.kt")
+        public void testSafeCallToPrimitiveEquality() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/boxingOptimization/safeCallToPrimitiveEquality.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("safeCallWithElvis.kt")
         public void testSafeCallWithElvis() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/boxingOptimization/safeCallWithElvis.kt");

compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -16616,6 +16616,30 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             doTest(fileName);
         }
 
+        @TestMetadata("primitiveEqSafeCall.kt")
+        public void testPrimitiveEqSafeCall() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/safeCall/primitiveEqSafeCall.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("primitiveNotEqSafeCall.kt")
+        public void testPrimitiveNotEqSafeCall() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/safeCall/primitiveNotEqSafeCall.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("safeCallEqPrimitive.kt")
+        public void testSafeCallEqPrimitive() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/safeCall/safeCallEqPrimitive.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("safeCallNotEqPrimitive.kt")
+        public void testSafeCallNotEqPrimitive() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/safeCall/safeCallNotEqPrimitive.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("safeCallOnLong.kt")
         public void testSafeCallOnLong() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/safeCall/safeCallOnLong.kt");

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -20626,6 +20626,30 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             doTest(fileName);
         }
 
+        @TestMetadata("primitiveEqSafeCall.kt")
+        public void testPrimitiveEqSafeCall() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/safeCall/primitiveEqSafeCall.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("primitiveNotEqSafeCall.kt")
+        public void testPrimitiveNotEqSafeCall() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/safeCall/primitiveNotEqSafeCall.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("safeCallEqPrimitive.kt")
+        public void testSafeCallEqPrimitive() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/safeCall/safeCallEqPrimitive.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("safeCallNotEqPrimitive.kt")
+        public void testSafeCallNotEqPrimitive() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/safeCall/safeCallNotEqPrimitive.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("safeCallOnLong.kt")
         public void testSafeCallOnLong() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/safeCall/safeCallOnLong.kt");