When2Switch: generate ifnonnull check for nullable values before *switch-opcode

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;
                         }
-                }
+                    }
         );
     }
 

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(),

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

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()];

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) {
+    private static boolean isThereConstantEntriesButNulls(
-        List<JetWhenEntry> entries = expression.getEntries();
+            @NotNull JetWhenExpression expression,
-        return !entries.isEmpty() && (entries.size() > 1 || !entries.get(0).isElse());
+            @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;
             }
         });
     }

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"
+}

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"
+}

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

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

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");

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");