When by enum:

Generate TABLESWITCH/LOOKUPSWITCH bytecode command in case of "when" by enum entries

compiler/backend/src/org/jetbrains/jet/codegen/AsmUtil.java

@@ -788,4 +788,9 @@ public class AsmUtil {
             v.aconst(type);
         }
     }
+
+    @NotNull
+    public static Type getArrayOf(@NotNull String internalClassName) {
+        return Type.getType("[L" + internalClassName + ";");
+    }
 }

compiler/backend/src/org/jetbrains/jet/codegen/binding/CodegenAnnotatingVisitor.java

@@ -20,13 +20,14 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.tree.IElementType;
 import com.intellij.psi.tree.TokenSet;
 import com.intellij.util.containers.Stack;
+import kotlin.Function1;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
-import org.jetbrains.jet.codegen.AsmUtil;
-import org.jetbrains.jet.codegen.JvmRuntimeTypes;
-import org.jetbrains.jet.codegen.SamCodegenUtil;
-import org.jetbrains.jet.codegen.SamType;
+import org.jetbrains.jet.codegen.*;
 import org.jetbrains.jet.codegen.state.GenerationState;
+import org.jetbrains.jet.codegen.when.SwitchCodegenUtil;
+import org.jetbrains.jet.codegen.when.WhenByEnumsMapping;
+import org.jetbrains.jet.lang.cfg.WhenChecker;
 import org.jetbrains.jet.lang.descriptors.*;
 import org.jetbrains.jet.lang.descriptors.impl.ClassDescriptorImpl;
 import org.jetbrains.jet.lang.psi.*;
@@ -37,8 +38,11 @@ import org.jetbrains.jet.lang.resolve.bindingContextUtil.BindingContextUtilPacka
 import org.jetbrains.jet.lang.resolve.calls.model.ExpressionValueArgument;
 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.java.JvmAbi;
 import org.jetbrains.jet.lang.resolve.java.PackageClassUtils;
+import org.jetbrains.jet.lang.resolve.kotlin.PackagePartClassUtils;
 import org.jetbrains.jet.lang.resolve.name.FqName;
 import org.jetbrains.jet.lang.resolve.name.Name;
 import org.jetbrains.jet.lang.resolve.scopes.JetScope;
@@ -489,4 +493,74 @@ class CodegenAnnotatingVisitor extends JetVisitorVoid {
             }
         }
     }
+
+    @Override
+    public void visitWhenExpression(@NotNull JetWhenExpression expression) {
+        super.visitWhenExpression(expression);
+        if (isWhenWithEnums(expression)) {
+            String currentClassName = getCurrentTopLevelClassOrPackagePartInternalName(expression.getContainingJetFile());
+
+            if (bindingContext.get(MAPPINGS_FOR_WHENS_BY_ENUM_IN_CLASS_FILE, currentClassName) == null) {
+                bindingTrace.record(
+                        MAPPINGS_FOR_WHENS_BY_ENUM_IN_CLASS_FILE,
+                        currentClassName,
+                        new ArrayList<WhenByEnumsMapping>()
+                );
+            }
+
+            List<WhenByEnumsMapping> mappings = bindingContext.get(MAPPINGS_FOR_WHENS_BY_ENUM_IN_CLASS_FILE, currentClassName);
+            assert mappings != null : "guaranteed by contract";
+
+            int fieldNumber = mappings.size();
+
+            JetType type = bindingContext.get(BindingContext.EXPRESSION_TYPE, expression.getSubjectExpression());
+            assert type != null : "should not be null in a valid when by enums";
+            ClassDescriptor classDescriptor = (ClassDescriptor) type.getConstructor().getDeclarationDescriptor();
+            assert classDescriptor != null : "because it's enum";
+
+            WhenByEnumsMapping mapping = new WhenByEnumsMapping(
+                    CodegenBinding.getAsmType(bindingContext, classDescriptor).getInternalName(),
+                    currentClassName,
+                    fieldNumber
+            );
+
+            for (CompileTimeConstant constant : SwitchCodegenUtil.getAllConstants(expression, bindingContext)) {
+                assert constant instanceof EnumValue : "expression in when should be EnumValue";
+                mapping.putFirstTime((EnumValue) constant, mapping.size() + 1);
+            }
+
+            mappings.add(mapping);
+
+            bindingTrace.record(MAPPING_FOR_WHEN_BY_ENUM, expression, mapping);
+        }
+    }
+
+    private boolean isWhenWithEnums(@NotNull JetWhenExpression expression) {
+        return WhenChecker.isWhenByEnum(expression, bindingContext) &&
+                SwitchCodegenUtil.checkAllItemsAreConstantsSatisfying(
+                    expression,
+                    bindingContext,
+                    new Function1<CompileTimeConstant, Boolean>() {
+                        @Override
+                        public Boolean invoke(
+                             @NotNull CompileTimeConstant constant
+                        ) {
+                         return constant instanceof EnumValue;
+                        }
+                }
+        );
+    }
+
+    @NotNull
+    private String getCurrentTopLevelClassOrPackagePartInternalName(@NotNull JetFile file) {
+        ListIterator<ClassDescriptorWithState> iterator = classStack.listIterator(classStack.size());
+        while(iterator.hasPrevious()) {
+            ClassDescriptor previous = iterator.previous().getDescriptor();
+            if (DescriptorUtils.isTopLevelOrInnerClass(previous)) {
+                return CodegenBinding.getAsmType(bindingContext, previous).getInternalName();
+            }
+        }
+
+        return PackagePartClassUtils.getPackagePartInternalName(file);
+    }
 }

compiler/backend/src/org/jetbrains/jet/codegen/binding/CodegenBinding.java

@@ -21,6 +21,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.codegen.SamType;
 import org.jetbrains.jet.codegen.state.GenerationState;
+import org.jetbrains.jet.codegen.when.WhenByEnumsMapping;
 import org.jetbrains.jet.lang.descriptors.*;
 import org.jetbrains.jet.lang.descriptors.impl.ClassDescriptorImpl;
 import org.jetbrains.jet.lang.psi.*;
@@ -60,6 +61,12 @@ public class CodegenBinding {
 
     public static final WritableSlice<JetExpression, SamType> SAM_VALUE = Slices.createSimpleSlice();
 
+    public static final WritableSlice<JetWhenExpression, WhenByEnumsMapping> MAPPING_FOR_WHEN_BY_ENUM = Slices.
+            <JetWhenExpression, WhenByEnumsMapping>sliceBuilder().build();
+
+    public static final WritableSlice<String, List<WhenByEnumsMapping>> MAPPINGS_FOR_WHENS_BY_ENUM_IN_CLASS_FILE =
+            Slices.<String, List<WhenByEnumsMapping>>sliceBuilder().build();
+
     static {
         BasicWritableSlice.initSliceDebugNames(CodegenBinding.class);
     }

compiler/backend/src/org/jetbrains/jet/codegen/state/GenerationState.java

@@ -25,6 +25,7 @@ import org.jetbrains.jet.codegen.inline.InlineCodegenUtil;
 import org.jetbrains.jet.codegen.intrinsics.IntrinsicMethods;
 import org.jetbrains.jet.codegen.optimization.OptimizationClassBuilderFactory;
 import org.jetbrains.jet.codegen.optimization.OptimizationUtils;
+import org.jetbrains.jet.codegen.when.MappingsClassesForWhenByEnum;
 import org.jetbrains.jet.lang.descriptors.ModuleDescriptor;
 import org.jetbrains.jet.lang.descriptors.ScriptDescriptor;
 import org.jetbrains.jet.lang.diagnostics.DiagnosticHolder;
@@ -79,6 +80,9 @@ public class GenerationState {
     @NotNull
     private final SamWrapperClasses samWrapperClasses = new SamWrapperClasses(this);
 
+    @NotNull
+    private final MappingsClassesForWhenByEnum mappingsClassesForWhenByEnum = new MappingsClassesForWhenByEnum(this);
+
     @NotNull
     private final BindingTrace bindingTrace;
 
@@ -222,6 +226,11 @@ public class GenerationState {
         return samWrapperClasses;
     }
 
+    @NotNull
+    public MappingsClassesForWhenByEnum getMappingsClassesForWhenByEnum() {
+        return mappingsClassesForWhenByEnum;
+    }
+
     public boolean isGenerateNotNullAssertions() {
         return generateNotNullAssertions;
     }

compiler/backend/src/org/jetbrains/jet/codegen/when/EnumSwitchCodegen.java

@@ -0,0 +1,53 @@
+package org.jetbrains.jet.codegen.when;
+
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.jet.codegen.ExpressionCodegen;
+import org.jetbrains.jet.lang.psi.JetWhenExpression;
+import org.jetbrains.jet.lang.resolve.constants.CompileTimeConstant;
+import org.jetbrains.jet.lang.resolve.constants.EnumValue;
+import org.jetbrains.org.objectweb.asm.Label;
+import org.jetbrains.org.objectweb.asm.Type;
+
+public class EnumSwitchCodegen extends SwitchCodegen {
+    private final WhenByEnumsMapping mapping;
+
+    public EnumSwitchCodegen(
+            @NotNull JetWhenExpression expression,
+            boolean isStatement,
+            @NotNull ExpressionCodegen codegen,
+            @NotNull WhenByEnumsMapping mapping
+    ) {
+        super(expression, isStatement, codegen);
+        this.mapping = mapping;
+    }
+
+    @Override
+    protected void generateSubject() {
+        codegen.getState().getMappingsClassesForWhenByEnum().generateMappingsClassForExpression(expression);
+
+        v.getstatic(
+                mapping.getMappingsClassInternalName(),
+                mapping.getFieldName(),
+                MappingClassesForWhenByEnumCodegen.MAPPINGS_FIELD_DESCRIPTOR
+        );
+
+        super.generateSubject();
+
+        v.invokevirtual(
+                mapping.getEnumClassInternalName(),
+                "ordinal",
+                Type.getMethodDescriptor(Type.INT_TYPE),
+                false
+        );
+        v.aload(Type.INT_TYPE);
+    }
+
+    @Override
+    protected void processConstant(
+            @NotNull CompileTimeConstant constant,
+            @NotNull Label entryLabel
+    ) {
+        assert constant instanceof EnumValue : "guaranteed by usage contract";
+        putTransitionOnce(mapping.getIndexByEntry((EnumValue) constant), entryLabel);
+    }
+}

compiler/backend/src/org/jetbrains/jet/codegen/when/MappingClassesForWhenByEnumCodegen.java

@@ -0,0 +1,147 @@
+/*
+ * Copyright 2010-2014 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.jet.codegen.when;
+
+import com.intellij.util.ArrayUtil;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.jet.codegen.AsmUtil;
+import org.jetbrains.jet.codegen.ClassBuilder;
+import org.jetbrains.jet.codegen.state.GenerationState;
+import org.jetbrains.jet.lang.psi.JetFile;
+import org.jetbrains.jet.lang.resolve.constants.EnumValue;
+import org.jetbrains.jet.lang.resolve.java.diagnostics.JvmDeclarationOrigin;
+import org.jetbrains.org.objectweb.asm.MethodVisitor;
+import org.jetbrains.org.objectweb.asm.Type;
+import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter;
+
+import java.util.List;
+import java.util.Map;
+
+import static org.jetbrains.jet.lang.resolve.java.AsmTypeConstants.OBJECT_TYPE;
+import static org.jetbrains.org.objectweb.asm.Opcodes.*;
+
+public class MappingClassesForWhenByEnumCodegen {
+    public static final String MAPPINGS_FIELD_DESCRIPTOR = Type.getDescriptor(int[].class);
+    private final GenerationState state;
+
+    public MappingClassesForWhenByEnumCodegen(@NotNull GenerationState state) {
+        this.state = state;
+    }
+
+    public void generate(
+            @NotNull List<WhenByEnumsMapping> mappings,
+            @NotNull Type mappingsClass,
+            @NotNull JetFile srcFile
+    ) {
+        ClassBuilder cb = state.getFactory().newVisitor(
+                JvmDeclarationOrigin.NO_ORIGIN,
+                mappingsClass,
+                srcFile
+        );
+        cb.defineClass(
+                srcFile,
+                V1_6,
+                ACC_FINAL | ACC_SYNTHETIC,
+                mappingsClass.getInternalName(),
+                null,
+                OBJECT_TYPE.getInternalName(),
+                ArrayUtil.EMPTY_STRING_ARRAY
+        );
+
+        generateFields(cb, mappings);
+        generateInitialization(cb, mappings);
+
+        cb.done();
+    }
+
+    private static void generateFields(
+            @NotNull ClassBuilder cb,
+            @NotNull List<WhenByEnumsMapping> mappings
+    ) {
+        for (WhenByEnumsMapping mapping : mappings) {
+            cb.newField(
+                    JvmDeclarationOrigin.NO_ORIGIN,
+                    ACC_STATIC | ACC_PUBLIC | ACC_FINAL | ACC_SYNTHETIC,
+                    mapping.getFieldName(),
+                    MAPPINGS_FIELD_DESCRIPTOR,
+                    null, null
+            );
+        }
+    }
+
+    private static void generateInitialization(
+            @NotNull ClassBuilder cb,
+            @NotNull List<WhenByEnumsMapping> mappings
+    ) {
+        MethodVisitor mv = cb.newMethod(
+                JvmDeclarationOrigin.NO_ORIGIN,
+                ACC_STATIC | ACC_SYNTHETIC, "<clinit>", "()V", null, ArrayUtil.EMPTY_STRING_ARRAY
+        );
+
+        mv.visitCode();
+
+        InstructionAdapter v = new InstructionAdapter(mv);
+
+        for (WhenByEnumsMapping mapping : mappings) {
+            generateInitializationForMapping(cb, v, mapping);
+        }
+
+        v.areturn(Type.VOID_TYPE);
+
+        mv.visitMaxs(-1, -1);
+        mv.visitEnd();
+    }
+
+    private static void generateInitializationForMapping(
+            @NotNull ClassBuilder cb,
+            @NotNull InstructionAdapter v,
+            @NotNull WhenByEnumsMapping mapping
+    ) {
+        v.invokestatic(
+                mapping.getEnumClassInternalName(), "values",
+                Type.getMethodDescriptor(
+                        AsmUtil.getArrayOf(mapping.getEnumClassInternalName())
+                ),
+                false
+        );
+        v.arraylength();
+
+        v.newarray(Type.INT_TYPE);
+        v.putstatic(cb.getThisName(), mapping.getFieldName(), MAPPINGS_FIELD_DESCRIPTOR);
+
+        String enumClassDesc = Type.getObjectType(mapping.getEnumClassInternalName()).getDescriptor();
+
+        for (Map.Entry<EnumValue, Integer> item : mapping.enumValuesToIntMapping()) {
+            EnumValue enumEntry = item.getKey();
+            int mappedValue = item.getValue();
+
+            v.getstatic(cb.getThisName(), mapping.getFieldName(), MAPPINGS_FIELD_DESCRIPTOR);
+            v.getstatic(
+                    mapping.getEnumClassInternalName(),
+                    enumEntry.getValue().getName().asString(),
+                    enumClassDesc
+            );
+            v.invokevirtual(
+                    mapping.getEnumClassInternalName(), "ordinal",
+                    Type.getMethodDescriptor(Type.INT_TYPE),
+                    false
+            );
+            v.iconst(mappedValue);
+            v.astore(Type.INT_TYPE);
+        }
+    }
+}

compiler/backend/src/org/jetbrains/jet/codegen/when/MappingsClassesForWhenByEnum.java

@@ -0,0 +1,58 @@
+/*
+ * Copyright 2010-2014 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.jet.codegen.when;
+
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.jet.codegen.binding.CodegenBinding;
+import org.jetbrains.jet.codegen.state.GenerationState;
+import org.jetbrains.jet.lang.psi.JetWhenExpression;
+import org.jetbrains.org.objectweb.asm.Type;
+
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+public class MappingsClassesForWhenByEnum {
+    private final GenerationState state;
+    private final Set<String> generatedMappingClasses = new HashSet<String>();
+    private final MappingClassesForWhenByEnumCodegen mappingsCodegen;
+
+    public MappingsClassesForWhenByEnum(@NotNull GenerationState state) {
+        this.state = state;
+        this.mappingsCodegen = new MappingClassesForWhenByEnumCodegen(state);
+    }
+
+    public void generateMappingsClassForExpression(@NotNull JetWhenExpression expression) {
+        WhenByEnumsMapping mapping = state.getBindingContext().get(CodegenBinding.MAPPING_FOR_WHEN_BY_ENUM, expression);
+
+        assert mapping != null : "mapping class should not be requested for non enum when";
+
+        if (!generatedMappingClasses.contains(mapping.getMappingsClassInternalName())) {
+            List<WhenByEnumsMapping> mappings = state.getBindingContext().get(
+                    CodegenBinding.MAPPINGS_FOR_WHENS_BY_ENUM_IN_CLASS_FILE,
+                    mapping.getOuterClassInternalNameForExpression()
+            );
+
+            assert mappings != null : "guaranteed by usage contract of EnumSwitchCodegen";
+
+            Type mappingsClassType = Type.getObjectType(mapping.getMappingsClassInternalName());
+
+            mappingsCodegen.generate(mappings, mappingsClassType, expression.getContainingJetFile());
+            generatedMappingClasses.add(mapping.getMappingsClassInternalName());
+        }
+    }
+}

compiler/backend/src/org/jetbrains/jet/codegen/when/SwitchCodegenUtil.java

@@ -20,6 +20,7 @@ import kotlin.Function1;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.codegen.ExpressionCodegen;
+import org.jetbrains.jet.codegen.binding.CodegenBinding;
 import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lang.resolve.BindingContext;
 import org.jetbrains.jet.lang.resolve.constants.CompileTimeConstant;
@@ -35,6 +36,12 @@ public class SwitchCodegenUtil {
             @NotNull Type subjectType,
             @NotNull BindingContext bindingContext
     ) {
+        // *switch opcode can be used if each item is enum entry or integral constant
+        // in case of enum CodegenAnnotationVisitor should put mappings into bindingContext
+        if (bindingContext.get(CodegenBinding.MAPPING_FOR_WHEN_BY_ENUM, expression) != null) {
+            return true;
+        }
+
         int typeSort = subjectType.getSort();
 
         if (typeSort != Type.INT && typeSort != Type.CHAR && typeSort != Type.SHORT && typeSort != Type.BYTE) {
@@ -120,6 +127,12 @@ public class SwitchCodegenUtil {
             boolean isStatement,
             @NotNull ExpressionCodegen codegen
     ) {
+        WhenByEnumsMapping mapping = codegen.getBindingContext().get(CodegenBinding.MAPPING_FOR_WHEN_BY_ENUM, expression);
+
+        if (mapping != null) {
+            return new EnumSwitchCodegen(expression, isStatement, codegen, mapping);
+        }
+
         return new IntegralConstantsSwitchCodegen(expression, isStatement, codegen);
     }
 }

compiler/backend/src/org/jetbrains/jet/codegen/when/WhenByEnumsMapping.java

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2010-2014 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.jet.codegen.when;
+
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.jet.lang.resolve.constants.EnumValue;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class WhenByEnumsMapping {
+    private static final String MAPPING_ARRAY_FIELD_PREFIX = "$EnumSwitchMapping$";
+    private static final String MAPPINGS_CLASS_NAME_POSTFIX = "$WhenMappings";
+
+    private final Map<EnumValue, Integer> map = new HashMap<EnumValue, Integer>();
+    private final String enumClassInternalName;
+    private final String outerClassInternalNameForExpression;
+    private final String mappingsClassInternalName;
+    private final int fieldNumber;
+
+    public WhenByEnumsMapping(String enumClassInternalName, String outerClassInternalNameForExpression, int fieldNumber) {
+        this.enumClassInternalName = enumClassInternalName;
+        this.outerClassInternalNameForExpression = outerClassInternalNameForExpression;
+        this.mappingsClassInternalName = outerClassInternalNameForExpression + MAPPINGS_CLASS_NAME_POSTFIX;
+        this.fieldNumber = fieldNumber;
+    }
+
+    public int getIndexByEntry(@NotNull EnumValue value) {
+        Integer result = map.get(value);
+        assert result != null : "entry " + value + " has no mapping";
+        return result;
+    }
+
+    public void putFirstTime(EnumValue value, int index) {
+        if (!map.containsKey(value)) {
+            map.put(value, index);
+        }
+    }
+
+    public int size() {
+        return map.size();
+    }
+
+    public String getFieldName() {
+        return MAPPING_ARRAY_FIELD_PREFIX + fieldNumber;
+    }
+
+    public String getEnumClassInternalName() {
+        return enumClassInternalName;
+    }
+
+    public String getOuterClassInternalNameForExpression() {
+        return outerClassInternalNameForExpression;
+    }
+
+    public String getMappingsClassInternalName() {
+        return mappingsClassInternalName;
+    }
+
+    public Iterable<Map.Entry<EnumValue, Integer>> enumValuesToIntMapping() {
+        return map.entrySet();
+    }
+}