From 8d8c3d2b9e6d404d73de9b06676f580481f9ee2d Mon Sep 17 00:00:00 2001
From: Denis Zharkov <denis.zharkov@jetbrains.com>
Date: Thu, 17 Jul 2014 11:59:58 +0400
Subject: [PATCH] When2Switch: generate ifnonnull check for nullable values
 before *switch-opcode

---
 .../binding/CodegenAnnotatingVisitor.java     |  9 ++--
 .../jet/codegen/when/EnumSwitchCodegen.java   |  5 +-
 .../jet/codegen/when/StringSwitchCodegen.java |  3 ++
 .../jet/codegen/when/SwitchCodegen.java       | 48 +++++++++++++++++--
 .../jet/codegen/when/SwitchCodegenUtil.java   | 21 +++++---
 .../whenEnumOptimization/nullability.kt       | 40 ++++++++++++++++
 .../whenStringOptimization/nullability.kt     | 41 ++++++++++++++++
 .../whenEnumOptimization/nullability.kt       | 26 ++++++++++
 .../whenStringOptimization/nullability.kt     | 18 +++++++
 .../codegen/BytecodeTextTestGenerated.java    | 10 ++++
 ...lackBoxWithStdlibCodegenTestGenerated.java | 10 ++++
 11 files changed, 218 insertions(+), 13 deletions(-)
 create mode 100644 compiler/testData/codegen/boxWithStdlib/whenEnumOptimization/nullability.kt
 create mode 100644 compiler/testData/codegen/boxWithStdlib/whenStringOptimization/nullability.kt
 create mode 100644 compiler/testData/codegen/bytecodeText/whenEnumOptimization/nullability.kt
 create mode 100644 compiler/testData/codegen/bytecodeText/whenStringOptimization/nullability.kt

diff --git a/compiler/backend/src/org/jetbrains/jet/codegen/binding/CodegenAnnotatingVisitor.java b/compiler/backend/src/org/jetbrains/jet/codegen/binding/CodegenAnnotatingVisitor.java
index 608ce19731b..1ea38075945 100644
--- a/compiler/backend/src/org/jetbrains/jet/codegen/binding/CodegenAnnotatingVisitor.java
+++ b/compiler/backend/src/org/jetbrains/jet/codegen/binding/CodegenAnnotatingVisitor.java
@@ -40,6 +40,7 @@ import org.jetbrains.jet.lang.resolve.calls.model.ResolvedCall;
 import org.jetbrains.jet.lang.resolve.calls.model.ResolvedValueArgument;
 import org.jetbrains.jet.lang.resolve.constants.CompileTimeConstant;
 import org.jetbrains.jet.lang.resolve.constants.EnumValue;
+import org.jetbrains.jet.lang.resolve.constants.NullValue;
 import org.jetbrains.jet.lang.resolve.java.JvmAbi;
 import org.jetbrains.jet.lang.resolve.java.PackageClassUtils;
 import org.jetbrains.jet.lang.resolve.kotlin.PackagePartClassUtils;
@@ -525,6 +526,8 @@ class CodegenAnnotatingVisitor extends JetVisitorVoid {
             );
 
             for (CompileTimeConstant constant : SwitchCodegenUtil.getAllConstants(expression, bindingContext)) {
+                if (constant instanceof NullValue) continue;
+
                 assert constant instanceof EnumValue : "expression in when should be EnumValue";
                 mapping.putFirstTime((EnumValue) constant, mapping.size() + 1);
             }
@@ -543,11 +546,11 @@ class CodegenAnnotatingVisitor extends JetVisitorVoid {
                     new Function1<CompileTimeConstant, Boolean>() {
                         @Override
                         public Boolean invoke(
-                             @NotNull CompileTimeConstant constant
+                                @NotNull CompileTimeConstant constant
                         ) {
-                         return constant instanceof EnumValue;
+                            return constant instanceof EnumValue || constant instanceof NullValue;
                         }
-                }
+                    }
         );
     }
 
diff --git a/compiler/backend/src/org/jetbrains/jet/codegen/when/EnumSwitchCodegen.java b/compiler/backend/src/org/jetbrains/jet/codegen/when/EnumSwitchCodegen.java
index 082f397ed42..ff2241c3913 100644
--- a/compiler/backend/src/org/jetbrains/jet/codegen/when/EnumSwitchCodegen.java
+++ b/compiler/backend/src/org/jetbrains/jet/codegen/when/EnumSwitchCodegen.java
@@ -25,13 +25,16 @@ public class EnumSwitchCodegen extends SwitchCodegen {
     protected void generateSubject() {
         codegen.getState().getMappingsClassesForWhenByEnum().generateMappingsClassForExpression(expression);
 
+        super.generateSubject();
+        generateNullCheckIfNeeded();
+
         v.getstatic(
                 mapping.getMappingsClassInternalName(),
                 mapping.getFieldName(),
                 MappingClassesForWhenByEnumCodegen.MAPPINGS_FIELD_DESCRIPTOR
         );
 
-        super.generateSubject();
+        v.swap();
 
         v.invokevirtual(
                 mapping.getEnumClassInternalName(),
diff --git a/compiler/backend/src/org/jetbrains/jet/codegen/when/StringSwitchCodegen.java b/compiler/backend/src/org/jetbrains/jet/codegen/when/StringSwitchCodegen.java
index db5b334c104..78324811957 100644
--- a/compiler/backend/src/org/jetbrains/jet/codegen/when/StringSwitchCodegen.java
+++ b/compiler/backend/src/org/jetbrains/jet/codegen/when/StringSwitchCodegen.java
@@ -75,6 +75,9 @@ public class StringSwitchCodegen extends SwitchCodegen {
         v.store(tempVarIndex, subjectType);
 
         v.load(tempVarIndex, subjectType);
+
+        generateNullCheckIfNeeded();
+
         v.invokevirtual(
                 subjectType.getInternalName(),
                 "hashCode", HASH_CODE_METHOD_DESC, false
diff --git a/compiler/backend/src/org/jetbrains/jet/codegen/when/SwitchCodegen.java b/compiler/backend/src/org/jetbrains/jet/codegen/when/SwitchCodegen.java
index d4ffa1158b9..9040a091e8d 100644
--- a/compiler/backend/src/org/jetbrains/jet/codegen/when/SwitchCodegen.java
+++ b/compiler/backend/src/org/jetbrains/jet/codegen/when/SwitchCodegen.java
@@ -23,12 +23,16 @@ import org.jetbrains.jet.lang.psi.JetWhenEntry;
 import org.jetbrains.jet.lang.psi.JetWhenExpression;
 import org.jetbrains.jet.lang.resolve.BindingContext;
 import org.jetbrains.jet.lang.resolve.constants.CompileTimeConstant;
+import org.jetbrains.jet.lang.resolve.constants.NullValue;
+import org.jetbrains.jet.lang.types.JetType;
 import org.jetbrains.org.objectweb.asm.Label;
 import org.jetbrains.org.objectweb.asm.Type;
 import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter;
 
 import java.util.*;
 
+import static org.jetbrains.jet.lang.resolve.BindingContext.EXPRESSION_TYPE;
+
 abstract public class SwitchCodegen {
     protected final JetWhenExpression expression;
     protected final boolean isStatement;
@@ -39,7 +43,7 @@ abstract public class SwitchCodegen {
     protected final InstructionAdapter v;
 
     protected final NavigableMap<Integer, Label> transitionsTable = new TreeMap<Integer, Label>();
-    protected final Collection<Label> entryLabels = new ArrayList<Label>();
+    protected final List<Label> entryLabels = new ArrayList<Label>();
     protected Label elseLabel = new Label();
     protected Label endLabel = new Label();
     protected Label defaultLabel;
@@ -66,10 +70,11 @@ abstract public class SwitchCodegen {
 
         boolean hasElse = expression.getElseExpression() != null;
 
-        generateSubject();
-
         // if there is no else-entry and it's statement then default --- endLabel
         defaultLabel = (hasElse || !isStatement) ? elseLabel : endLabel;
+
+        generateSubject();
+
         generateSwitchInstructionByTransitionsTable();
 
         generateEntries();
@@ -93,6 +98,7 @@ abstract public class SwitchCodegen {
             Label entryLabel = new Label();
 
             for (CompileTimeConstant constant : SwitchCodegenUtil.getConstantsFromEntry(entry, bindingContext)) {
+                if (constant instanceof NullValue) continue;
                 processConstant(constant, entryLabel);
             }
 
@@ -124,6 +130,42 @@ abstract public class SwitchCodegen {
         codegen.gen(expression.getSubjectExpression(), subjectType);
     }
 
+    protected void generateNullCheckIfNeeded() {
+        JetType subjectJetType = bindingContext.get(EXPRESSION_TYPE, expression.getSubjectExpression());
+
+        assert subjectJetType != null : "subject type can't be null (i.e. void)";
+
+        if (subjectJetType.isNullable()) {
+            int nullEntryIndex = findNullEntryIndex(expression);
+            Label nullLabel = nullEntryIndex == -1 ? defaultLabel : entryLabels.get(nullEntryIndex);
+            Label notNullLabel = new Label();
+
+            v.dup();
+            v.ifnonnull(notNullLabel);
+
+            v.pop();
+
+            v.goTo(nullLabel);
+
+            v.visitLabel(notNullLabel);
+        }
+    }
+
+    private int findNullEntryIndex(@NotNull JetWhenExpression expression) {
+        int entryIndex = 0;
+        for (JetWhenEntry entry : expression.getEntries()) {
+            for (CompileTimeConstant constant : SwitchCodegenUtil.getConstantsFromEntry(entry, bindingContext)) {
+                if (constant instanceof NullValue) {
+                    return entryIndex;
+                }
+            }
+
+            entryIndex++;
+        }
+
+        return -1;
+    }
+
     private void generateSwitchInstructionByTransitionsTable() {
         int[] keys = new int[transitionsTable.size()];
         Label[] labels = new Label[transitionsTable.size()];
diff --git a/compiler/backend/src/org/jetbrains/jet/codegen/when/SwitchCodegenUtil.java b/compiler/backend/src/org/jetbrains/jet/codegen/when/SwitchCodegenUtil.java
index 74730d893ab..bf28b6097e9 100644
--- a/compiler/backend/src/org/jetbrains/jet/codegen/when/SwitchCodegenUtil.java
+++ b/compiler/backend/src/org/jetbrains/jet/codegen/when/SwitchCodegenUtil.java
@@ -25,6 +25,7 @@ import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lang.resolve.BindingContext;
 import org.jetbrains.jet.lang.resolve.constants.CompileTimeConstant;
 import org.jetbrains.jet.lang.resolve.constants.IntegerValueConstant;
+import org.jetbrains.jet.lang.resolve.constants.NullValue;
 import org.jetbrains.jet.lang.resolve.constants.StringValue;
 import org.jetbrains.org.objectweb.asm.Type;
 
@@ -78,6 +79,8 @@ public class SwitchCodegenUtil {
             @NotNull BindingContext bindingContext
     ) {
         for (JetWhenCondition condition : entry.getConditions()) {
+            if (!(condition instanceof JetWhenConditionWithExpression)) continue;
+
             JetExpression patternExpression = ((JetWhenConditionWithExpression) condition).getExpression();
 
             assert patternExpression != null : "expression in when should not be null";
@@ -101,12 +104,12 @@ public class SwitchCodegenUtil {
             boolean isStatement,
             @NotNull ExpressionCodegen codegen
     ) {
-        if (!isThereEntriesButElse(expression)) {
+        BindingContext bindingContext = codegen.getBindingContext();
+        if (!isThereConstantEntriesButNulls(expression, bindingContext)) {
             return null;
         }
 
         Type subjectType = codegen.expressionType(expression.getSubjectExpression());
-        BindingContext bindingContext = codegen.getBindingContext();
 
         WhenByEnumsMapping mapping = codegen.getBindingContext().get(CodegenBinding.MAPPING_FOR_WHEN_BY_ENUM, expression);
 
@@ -125,9 +128,15 @@ public class SwitchCodegenUtil {
         return null;
     }
 
-    private static boolean isThereEntriesButElse(@NotNull JetWhenExpression expression) {
-        List<JetWhenEntry> entries = expression.getEntries();
-        return !entries.isEmpty() && (entries.size() > 1 || !entries.get(0).isElse());
+    private static boolean isThereConstantEntriesButNulls(
+            @NotNull JetWhenExpression expression,
+            @NotNull BindingContext bindingContext
+    ) {
+        for (CompileTimeConstant constant : getAllConstants(expression, bindingContext)) {
+            if (constant != null && !(constant instanceof NullValue)) return true;
+        }
+
+        return false;
     }
 
     private static boolean isIntegralConstantsSwitch(
@@ -166,7 +175,7 @@ public class SwitchCodegenUtil {
             public Boolean invoke(
                     @NotNull CompileTimeConstant constant
             ) {
-                return constant instanceof StringValue;
+                return constant instanceof StringValue || constant instanceof NullValue;
             }
         });
     }
diff --git a/compiler/testData/codegen/boxWithStdlib/whenEnumOptimization/nullability.kt b/compiler/testData/codegen/boxWithStdlib/whenEnumOptimization/nullability.kt
new file mode 100644
index 00000000000..9fa02467f08
--- /dev/null
+++ b/compiler/testData/codegen/boxWithStdlib/whenEnumOptimization/nullability.kt
@@ -0,0 +1,40 @@
+import kotlin.test.assertEquals
+
+enum class Season {
+    WINTER
+    SPRING
+    SUMMER
+    AUTUMN
+}
+
+fun foo1(x : Season?) : String {
+    when(x) {
+        Season.AUTUMN, Season.SPRING -> return "autumn_or_spring";
+        Season.SUMMER, null -> return "summer_or_null"
+    }
+
+    return "other"
+}
+
+fun foo2(x : Season?) : String {
+    when(x) {
+        Season.AUTUMN, Season.SPRING -> return "autumn_or_spring";
+        Season.SUMMER -> return "summer"
+    }
+
+    return "other"
+}
+
+fun box() : String {
+    assertEquals("autumn_or_spring", foo1(Season.AUTUMN))
+    assertEquals("autumn_or_spring", foo1(Season.SPRING))
+    assertEquals("summer_or_null", foo1(Season.SUMMER))
+    assertEquals("summer_or_null", foo1(null))
+
+    assertEquals("autumn_or_spring", foo2(Season.AUTUMN))
+    assertEquals("autumn_or_spring", foo2(Season.SPRING))
+    assertEquals("summer", foo2(Season.SUMMER))
+    assertEquals("other", foo2(null))
+
+    return "OK"
+}
diff --git a/compiler/testData/codegen/boxWithStdlib/whenStringOptimization/nullability.kt b/compiler/testData/codegen/boxWithStdlib/whenStringOptimization/nullability.kt
new file mode 100644
index 00000000000..3219824d1cb
--- /dev/null
+++ b/compiler/testData/codegen/boxWithStdlib/whenStringOptimization/nullability.kt
@@ -0,0 +1,41 @@
+import kotlin.test.assertEquals
+
+fun foo1(x : String?) : String {
+    when (x) {
+        "abc", "cde" -> return "abc_cde"
+        "efg", "ghi", null -> return "efg_ghi"
+    }
+
+    return "other"
+}
+
+fun foo2(x : String?) : String {
+    when (x) {
+        "abc", "cde" -> return "abc_cde"
+        "efg", "ghi" -> return "efg_ghi"
+        else -> return "other"
+    }
+}
+
+fun box() : String {
+    //foo1
+    assertEquals("abc_cde", foo1("abc"))
+    assertEquals("abc_cde", foo1("cde"))
+    assertEquals("efg_ghi", foo1("efg"))
+    assertEquals("efg_ghi", foo1("ghi"))
+    assertEquals("efg_ghi", foo1(null))
+
+    assertEquals("other",   foo1("xyz"))
+
+    //foo2
+    assertEquals("abc_cde", foo2("abc"))
+    assertEquals("abc_cde", foo2("cde"))
+    assertEquals("efg_ghi", foo2("efg"))
+    assertEquals("efg_ghi", foo2("ghi"))
+
+
+    assertEquals("other",   foo2("xyz"))
+    assertEquals("other",   foo2(null))
+
+    return "OK"
+}
diff --git a/compiler/testData/codegen/bytecodeText/whenEnumOptimization/nullability.kt b/compiler/testData/codegen/bytecodeText/whenEnumOptimization/nullability.kt
new file mode 100644
index 00000000000..ab9a6185018
--- /dev/null
+++ b/compiler/testData/codegen/bytecodeText/whenEnumOptimization/nullability.kt
@@ -0,0 +1,26 @@
+enum class Season {
+    WINTER
+    SPRING
+    SUMMER
+    AUTUMN
+}
+
+fun foo1(x : Season?) : String {
+    when(x) {
+        Season.AUTUMN, Season.SPRING -> return "autumn_or_spring";
+        Season.SUMMER, null -> return "summer_or_null"
+    }
+
+    return "other"
+}
+
+fun foo2(x : Season?) : String {
+    when(x) {
+        Season.AUTUMN, Season.SPRING -> return "autumn_or_spring";
+        Season.SUMMER -> return "summer"
+    }
+
+    return "other"
+}
+
+// 2 TABLESWITCH
diff --git a/compiler/testData/codegen/bytecodeText/whenStringOptimization/nullability.kt b/compiler/testData/codegen/bytecodeText/whenStringOptimization/nullability.kt
new file mode 100644
index 00000000000..52905d3a900
--- /dev/null
+++ b/compiler/testData/codegen/bytecodeText/whenStringOptimization/nullability.kt
@@ -0,0 +1,18 @@
+fun foo1(x : String?) : String {
+    when (x) {
+        "abc", "cde" -> return "abc_cde"
+        "efg", "ghi", null -> return "efg_ghi"
+    }
+
+    return "other"
+}
+
+fun foo2(x : String?) : String {
+    when (x) {
+        "abc", "cde" -> return "abc_cde"
+        "efg", "ghi" -> return "efg_ghi"
+        else -> return "other"
+    }
+}
+
+// 2 LOOKUPSWITCH
diff --git a/compiler/tests/org/jetbrains/jet/codegen/BytecodeTextTestGenerated.java b/compiler/tests/org/jetbrains/jet/codegen/BytecodeTextTestGenerated.java
index 52541da43b4..62f744e921d 100644
--- a/compiler/tests/org/jetbrains/jet/codegen/BytecodeTextTestGenerated.java
+++ b/compiler/tests/org/jetbrains/jet/codegen/BytecodeTextTestGenerated.java
@@ -340,6 +340,11 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
             doTest("compiler/testData/codegen/bytecodeText/whenEnumOptimization/nonConstantEnum.kt");
         }
         
+        @TestMetadata("nullability.kt")
+        public void testNullability() throws Exception {
+            doTest("compiler/testData/codegen/bytecodeText/whenEnumOptimization/nullability.kt");
+        }
+        
         @TestMetadata("subjectAny.kt")
         public void testSubjectAny() throws Exception {
             doTest("compiler/testData/codegen/bytecodeText/whenEnumOptimization/subjectAny.kt");
@@ -383,6 +388,11 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
             doTest("compiler/testData/codegen/bytecodeText/whenStringOptimization/nonInlinedConst.kt");
         }
         
+        @TestMetadata("nullability.kt")
+        public void testNullability() throws Exception {
+            doTest("compiler/testData/codegen/bytecodeText/whenStringOptimization/nullability.kt");
+        }
+        
         @TestMetadata("sameHashCode.kt")
         public void testSameHashCode() throws Exception {
             doTest("compiler/testData/codegen/bytecodeText/whenStringOptimization/sameHashCode.kt");
diff --git a/compiler/tests/org/jetbrains/jet/codegen/generated/BlackBoxWithStdlibCodegenTestGenerated.java b/compiler/tests/org/jetbrains/jet/codegen/generated/BlackBoxWithStdlibCodegenTestGenerated.java
index 5d97188a6f0..7631aa249a5 100644
--- a/compiler/tests/org/jetbrains/jet/codegen/generated/BlackBoxWithStdlibCodegenTestGenerated.java
+++ b/compiler/tests/org/jetbrains/jet/codegen/generated/BlackBoxWithStdlibCodegenTestGenerated.java
@@ -1993,6 +1993,11 @@ public class BlackBoxWithStdlibCodegenTestGenerated extends AbstractBlackBoxCode
             doTestWithStdlib("compiler/testData/codegen/boxWithStdlib/whenEnumOptimization/nonConstantEnum.kt");
         }
         
+        @TestMetadata("nullability.kt")
+        public void testNullability() throws Exception {
+            doTestWithStdlib("compiler/testData/codegen/boxWithStdlib/whenEnumOptimization/nullability.kt");
+        }
+        
         @TestMetadata("subjectAny.kt")
         public void testSubjectAny() throws Exception {
             doTestWithStdlib("compiler/testData/codegen/boxWithStdlib/whenEnumOptimization/subjectAny.kt");
@@ -2026,6 +2031,11 @@ public class BlackBoxWithStdlibCodegenTestGenerated extends AbstractBlackBoxCode
             doTestWithStdlib("compiler/testData/codegen/boxWithStdlib/whenStringOptimization/expression.kt");
         }
         
+        @TestMetadata("nullability.kt")
+        public void testNullability() throws Exception {
+            doTestWithStdlib("compiler/testData/codegen/boxWithStdlib/whenStringOptimization/nullability.kt");
+        }
+        
         @TestMetadata("sameHashCode.kt")
         public void testSameHashCode() throws Exception {
             doTestWithStdlib("compiler/testData/codegen/boxWithStdlib/whenStringOptimization/sameHashCode.kt");