From 1952781e4b8e45552b58346d1e66f775528eb1ed Mon Sep 17 00:00:00 2001
From: Denis Zharkov <bintree@eblans.org>
Date: Sat, 5 Apr 2014 16:26:01 +0400
Subject: [PATCH] SwitchOptimization: generating `lookupswitch' opcode for
 'when' in case of integral types

---
 .../jet/codegen/ExpressionCodegen.java        | 22 +++++--
 .../codegen/box/when/switchOptimization.kt    | 62 ++++++++++---------
 2 files changed, 50 insertions(+), 34 deletions(-)

diff --git a/compiler/backend/src/org/jetbrains/jet/codegen/ExpressionCodegen.java b/compiler/backend/src/org/jetbrains/jet/codegen/ExpressionCodegen.java
index dd57fb6c3ea..82f880bffc3 100644
--- a/compiler/backend/src/org/jetbrains/jet/codegen/ExpressionCodegen.java
+++ b/compiler/backend/src/org/jetbrains/jet/codegen/ExpressionCodegen.java
@@ -26,8 +26,10 @@ import com.intellij.psi.tree.IElementType;
 import com.intellij.util.ArrayUtil;
 import com.intellij.util.Function;
 import com.intellij.util.containers.Stack;
+import com.sun.jdi.IntegerValue;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.jet.lang.resolve.constants.IntegerValueConstant;
 import org.jetbrains.jet.codegen.binding.CalculatedClosure;
 import org.jetbrains.jet.codegen.binding.CodegenBinding;
 import org.jetbrains.jet.codegen.binding.MutableClosure;
@@ -3755,9 +3757,9 @@ The "returned" value of try expression with no finally is either the last expres
                 assert conditionExpression != null;
 
                 CompileTimeConstant constant = getCompileTimeConstant(conditionExpression, bindingContext);
-                assert constant != null && constant instanceof IntValue;
+                assert doesConstantFitForSwitch(constant);
 
-                int value = ((IntValue) constant).getValue();
+                int value = (constant.getValue() instanceof Number) ? ((Number)constant.getValue()).intValue() : ((Character)constant.getValue()).charValue();
 
                 if (!transitions.containsKey(value)) {
                     transitions.put(value, entryLabel);
@@ -3828,8 +3830,20 @@ The "returned" value of try expression with no finally is either the last expres
         v.lookupswitch(defaultLabel, keys, labels);
     }
 
+    private boolean doesConstantFitForSwitch(CompileTimeConstant constant) {
+        if (constant == null || !(constant instanceof IntegerValueConstant)) {
+            return false;
+        }
+
+        long value = (constant.getValue() instanceof Number) ? ((Number)constant.getValue()).longValue() : ((Character)constant.getValue()).charValue();
+        return value >= Integer.MIN_VALUE &&
+               value <= Integer.MAX_VALUE;
+    }
+
     private boolean canSwitchBeUsedIn(JetWhenExpression expression, Type subjectType) {
-        if (subjectType.getSort() != Type.INT) {
+        int typeSort = subjectType.getSort();
+
+        if (typeSort != Type.INT && typeSort != Type.CHAR && typeSort != Type.SHORT && typeSort != Type.BYTE) {
             return false;
         }
 
@@ -3847,7 +3861,7 @@ The "returned" value of try expression with no finally is either the last expres
                 }
 
                 CompileTimeConstant constant = getCompileTimeConstant(patternExpression, bindingContext);
-                if (constant == null || !(constant instanceof IntValue)) {
+                if (!doesConstantFitForSwitch(constant)) {
                     return false;
                 }
             }
diff --git a/compiler/testData/codegen/box/when/switchOptimization.kt b/compiler/testData/codegen/box/when/switchOptimization.kt
index 26aba979fd6..f2de7abd966 100644
--- a/compiler/testData/codegen/box/when/switchOptimization.kt
+++ b/compiler/testData/codegen/box/when/switchOptimization.kt
@@ -1,3 +1,4 @@
+
 fun foo1(x : Int) : Int {
     return when (x % 10) {
         1,2,3 -> 1
@@ -7,43 +8,44 @@ fun foo1(x : Int) : Int {
     }
 }
 
-fun foo2(x : Int) : Int {
-    when (x % 10) {
-        1,2,3 -> return 1
-        4,5,6 -> return 2
-        7,8,9 -> return 3
+fun foo2(x : Short) : Int {
+    when ((x % 10).toShort()) {
+        1.toShort(),2.toShort(),3.toShort() -> return 1
+        4.toShort(),5.toShort(),6.toShort() -> return 2
+        7.toShort(),8.toShort(),9.toShort() -> return 3
     }
 
     return 4
 }
 
-fun foo3(x : Int) : Int {
-    when (x % 10) {
-        1 -> return 1
-        2 -> return 1
-        3 -> return 1
-        4 -> return 2
-        5 -> return 2
-        6 -> return 2
-        7 -> return 3
-        8 -> return 3
-        9 -> return 3
+fun foo3(x : Char) : Int {
+
+    when (x) {
+        '1' -> return 1
+        '2' -> return 1
+        '3' -> return 1
+        '4' -> return 2
+        '5' -> return 2
+        '6' -> return 2
+        '7' -> return 3
+        '8' -> return 3
+        '9' -> return 3
     }
 
     return 4
 }
 
-fun foo4(x : Int) : Int {
-    return when (x % 10) {
-        1 -> 1
-        2 -> 1
-        3 -> 1
-        4 -> 2
-        5 -> 2
-        6 -> 2
-        7 -> 3
-        8 -> 3
-        9 -> 3
+fun foo4(x : Byte) : Int {
+    return when ((x % 10).toByte()) {
+        1.toByte() -> 1
+        2.toByte() -> 1
+        3.toByte() -> 1
+        4.toByte() -> 2
+        5.toByte() -> 2
+        6.toByte() -> 2
+        7.toByte() -> 3
+        8.toByte() -> 3
+        9.toByte() -> 3
         else -> 4
     }
 }
@@ -66,9 +68,9 @@ fun test(foo : (Int) -> Int, name : String) : String {
 
 fun box(): String {
     val testResult = test({x -> foo1(x)}, "foo1") +
-                     test({x -> foo2(x)}, "foo2") +
-                     test({x -> foo3(x)}, "foo3") +
-                     test({x -> foo4(x)}, "foo4")
+                     test({x -> foo2(x.toShort())}, "foo2") +
+                     test({x -> foo3((x % 10 + '0'.toInt()).toChar())}, "foo3") +
+                     test({x -> foo4(x.toByte())}, "foo4")
     if (testResult != "") return testResult
     return "OK"
 }