[JVM] Emit lookupswitch/tableswitch for Short/Byte if possible

^KT-76589

Author: scaventz

compiler/ir/backend.jvm/codegen/src/org/jetbrains/kotlin/backend/jvm/codegen/SwitchGenerator.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.backend.common.IrWhenUtils
 import org.jetbrains.kotlin.codegen.`when`.SwitchCodegen.Companion.preferLookupOverSwitch
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.types.*
+import org.jetbrains.kotlin.ir.util.fqNameWhenAvailable
 import org.jetbrains.kotlin.ir.util.hasAnnotation
 import org.jetbrains.kotlin.load.java.JvmAnnotationNames
 import org.jetbrains.org.objectweb.asm.Label
@@ -62,11 +63,13 @@ class SwitchGenerator(private val expression: IrWhen, private val data: BlockInf
 
         val firstCondition = callToLabels[0].call
         if (firstCondition.symbol != context.irBuiltIns.eqeqSymbol) return null
-        val subject = firstCondition.arguments[0]
+        val subject = firstCondition.arguments[0] ?: return null
         return when {
             subject is IrCall && subject.isCoerceFromUIntToInt() ->
                 generateUIntSwitch(subject.arguments[0]!! as? IrGetValue, calls, callToLabels, expressionToLabels, elseExpression)
-            subject is IrGetValue || subject is IrConst && subject.type.isString() -> // also generate tableswitch for literal string subject
+            subject is IrGetValue
+                    || subject is IrConst && subject.type.isString()  // also generate tableswitch for literal string subject
+                    || subject.isByteOrShortToInt() ->
                 generatePrimitiveSwitch(subject, calls, callToLabels, expressionToLabels, elseExpression)
             else ->
                 null
@@ -111,6 +114,12 @@ class SwitchGenerator(private val expression: IrWhen, private val data: BlockInf
         )
     }
 
+    fun IrExpression.isByteOrShortToInt(): Boolean {
+        if (this !is IrCall) return false
+        val fqName = symbol.owner.fqNameWhenAvailable?.asString()
+        return fqName == "kotlin.Byte.toInt" || fqName == "kotlin.Short.toInt"
+    }
+
     private fun generatePrimitiveSwitch(
         subject: IrExpression,
         conditions: List<IrCall>,
@@ -122,7 +131,7 @@ class SwitchGenerator(private val expression: IrWhen, private val data: BlockInf
         if (!areConstantComparisons(conditions)) return null
 
         return when {
-            subject is IrGetValue && areConstIntComparisons(conditions) -> {
+            (subject is IrGetValue || subject.isByteOrShortToInt()) && areConstIntComparisons(conditions) -> {
                 val cases = extractSwitchCasesAndFilterUnreachableLabels(callToLabels, expressionToLabels)
                 IntSwitch(
                     subject,
@@ -187,8 +196,17 @@ class SwitchGenerator(private val expression: IrWhen, private val data: BlockInf
             return lhs.all { it != null && it.value == lhs[0]!!.value }
         }
 
+        fun isValidCallToIntLHS(): Boolean {
+            val lhs = conditions.map {
+                val call = it.takeIf { it.symbol == context.irBuiltIns.eqeqSymbol }?.arguments[0] ?: return false
+                if (!call.isByteOrShortToInt()) return false
+                (call as IrCall).arguments[0] as? IrGetValue
+            }
+            return lhs.all { it != null && it.symbol.owner == lhs[0]!!.symbol.owner }
+        }
+
         // All conditions are equality checks && all LHS refer to the same tmp variable.
-        if (!isValidIrGetValueTypeLHS() && !isValidIrConstTypeLHS())
+        if (!isValidIrGetValueTypeLHS() && !isValidIrConstTypeLHS() && !isValidCallToIntLHS())
             return false
 
         // All RHS are constants
@@ -218,7 +236,7 @@ class SwitchGenerator(private val expression: IrWhen, private val data: BlockInf
         subjectTypePredicate: (IrType) -> Boolean,
         irConstPredicate: (IrConst) -> Boolean
     ): Boolean {
-        val lhs = conditions.map { it.arguments[0] as? IrGetValue ?: it.arguments[0] as IrConst }
+        val lhs = conditions.map { it.arguments[0] as? IrGetValue ?: it.arguments[0] as? IrConst ?: it.arguments[0] as IrCall }
         if (lhs.any { !subjectTypePredicate(it.type) })
             return false
 
@@ -314,7 +332,7 @@ class SwitchGenerator(private val expression: IrWhen, private val data: BlockInf
     }
 
     inner class IntSwitch(
-        subject: IrGetValue,
+        subject: IrExpression,
         elseExpression: IrExpression?,
         expressionToLabels: ArrayList<ExpressionToLabel>,
         private val cases: List<ValueToLabel>

compiler/testData/codegen/box/when/lookupSwitchOverByte.kt

@@ -1,14 +1,13 @@
-val byte10: Byte = 10
-val byte30: Byte = 30
-val byte50: Byte = 50
-val byte42: Byte = 42
+const val byte10: Byte = 10
+const val byte30: Byte = 30
+const val byte50: Byte = 50
+const val byte42: Byte = 42
 
 fun foo(p: Byte): String {
-    val localByte50: Byte = 50
     return when (p) {
         byte10 -> "10"
         byte30 -> "30"
-        localByte50 -> "50"
+        byte50 -> "50"
         else -> "else"
     }
 }

compiler/testData/codegen/box/when/lookupSwitchOverMixTypes.kt

@@ -2,10 +2,9 @@ const val byte10: Byte = 10
 const val short50: Short = 50
 
 fun foo(p: Any): String {
-    val newByte10: Byte = 10
     return when (p) {
         0 -> "0"
-        newByte10 -> "byte10"
+        byte10 -> "byte10"
         short50 -> "short50"
         'z' -> "z"
         else -> "else"
@@ -20,8 +19,4 @@ fun box(): String {
         && foo("else") == "else"
     ) "OK"
     else "FAIL"
-}
-
-// CHECK_BYTECODE_TEXT
-// 1 LOOKUPSWITCH
-// 0 TABLESWITCH
+}

compiler/testData/codegen/box/when/lookupSwitchOverShort.kt

@@ -1,7 +1,7 @@
-val short100: Short = 100
-val short200: Short = 200
-val short300: Short = 300
-val short400: Short = 400
+const val short100: Short = 100
+const val short200: Short = 200
+const val short300: Short = 300
+const val short400: Short = 400
 
 fun foo(p: Short): String {
     return when (p) {

compiler/testData/codegen/box/when/tableSwitchOverByte.kt

@@ -1,7 +1,7 @@
-val byte101: Byte = 101
-val byte102: Byte = 102
-val byte103: Byte = 103
-val byte104: Byte = 104
+const val byte101: Byte = 101
+const val byte102: Byte = 102
+const val byte103: Byte = 103
+const val byte104: Byte = 104
 
 fun foo(p: Byte): String {
     return when (p) {

compiler/testData/codegen/box/when/tableSwitchOverShort.kt

@@ -1,7 +1,7 @@
-val short101: Short = 101
-val short102: Short = 102
-val short103: Short = 103
-val short104: Short = 104
+const val short101: Short = 101
+const val short102: Short = 102
+const val short103: Short = 103
+const val short104: Short = 104
 
 fun foo(p: Short): String {
     return when (p) {