001    /*
002     * Copyright 2010-2015 JetBrains s.r.o.
003     *
004     * Licensed under the Apache License, Version 2.0 (the "License");
005     * you may not use this file except in compliance with the License.
006     * You may obtain a copy of the License at
007     *
008     * http://www.apache.org/licenses/LICENSE-2.0
009     *
010     * Unless required by applicable law or agreed to in writing, software
011     * distributed under the License is distributed on an "AS IS" BASIS,
012     * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013     * See the License for the specific language governing permissions and
014     * limitations under the License.
015     */
016    
017    package org.jetbrains.kotlin.codegen.when;
018    
019    import kotlin.jvm.functions.Function1;
020    import org.jetbrains.annotations.NotNull;
021    import org.jetbrains.annotations.Nullable;
022    import org.jetbrains.kotlin.codegen.ExpressionCodegen;
023    import org.jetbrains.kotlin.codegen.binding.CodegenBinding;
024    import org.jetbrains.kotlin.psi.*;
025    import org.jetbrains.kotlin.resolve.BindingContext;
026    import org.jetbrains.kotlin.resolve.constants.CompileTimeConstant;
027    import org.jetbrains.kotlin.resolve.constants.IntegerValueConstant;
028    import org.jetbrains.kotlin.resolve.constants.NullValue;
029    import org.jetbrains.kotlin.resolve.constants.StringValue;
030    import org.jetbrains.org.objectweb.asm.Type;
031    
032    import java.util.ArrayList;
033    import java.util.List;
034    
035    public class SwitchCodegenUtil {
036        public static boolean checkAllItemsAreConstantsSatisfying(
037                @NotNull JetWhenExpression expression,
038                @NotNull BindingContext bindingContext,
039                Function1<CompileTimeConstant, Boolean> predicate
040        ) {
041            for (JetWhenEntry entry : expression.getEntries()) {
042                for (JetWhenCondition condition : entry.getConditions()) {
043                    if (!(condition instanceof JetWhenConditionWithExpression)) {
044                        return false;
045                    }
046    
047                    // ensure that expression is constant
048                    JetExpression patternExpression = ((JetWhenConditionWithExpression) condition).getExpression();
049    
050                    if (patternExpression == null) return false;
051    
052                    CompileTimeConstant constant = ExpressionCodegen.getCompileTimeConstant(patternExpression, bindingContext);
053                    if (constant == null || !predicate.invoke(constant)) {
054                        return false;
055                    }
056                }
057            }
058    
059            return true;
060        }
061    
062        @NotNull
063        public static Iterable<CompileTimeConstant> getAllConstants(
064                @NotNull JetWhenExpression expression,
065                @NotNull BindingContext bindingContext
066        ) {
067            List<CompileTimeConstant> result = new ArrayList<CompileTimeConstant>();
068    
069            for (JetWhenEntry entry : expression.getEntries()) {
070                addConstantsFromEntry(result, entry, bindingContext);
071            }
072    
073            return result;
074        }
075    
076        private static void addConstantsFromEntry(
077                @NotNull List<CompileTimeConstant> result,
078                @NotNull JetWhenEntry entry,
079                @NotNull BindingContext bindingContext
080        ) {
081            for (JetWhenCondition condition : entry.getConditions()) {
082                if (!(condition instanceof JetWhenConditionWithExpression)) continue;
083    
084                JetExpression patternExpression = ((JetWhenConditionWithExpression) condition).getExpression();
085    
086                assert patternExpression != null : "expression in when should not be null";
087                result.add(ExpressionCodegen.getCompileTimeConstant(patternExpression, bindingContext));
088            }
089        }
090    
091        @NotNull
092        public static Iterable<CompileTimeConstant> getConstantsFromEntry(
093                @NotNull JetWhenEntry entry,
094                @NotNull BindingContext bindingContext
095        ) {
096            List<CompileTimeConstant> result = new ArrayList<CompileTimeConstant>();
097            addConstantsFromEntry(result, entry, bindingContext);
098            return result;
099        }
100    
101        @Nullable
102        public static SwitchCodegen buildAppropriateSwitchCodegenIfPossible(
103                @NotNull JetWhenExpression expression,
104                boolean isStatement,
105                @NotNull ExpressionCodegen codegen
106        ) {
107            BindingContext bindingContext = codegen.getBindingContext();
108            if (!isThereConstantEntriesButNulls(expression, bindingContext)) {
109                return null;
110            }
111    
112            Type subjectType = codegen.expressionType(expression.getSubjectExpression());
113    
114            WhenByEnumsMapping mapping = codegen.getBindingContext().get(CodegenBinding.MAPPING_FOR_WHEN_BY_ENUM, expression);
115    
116            if (mapping != null) {
117                return new EnumSwitchCodegen(expression, isStatement, codegen, mapping);
118            }
119    
120            if (isIntegralConstantsSwitch(expression, subjectType, bindingContext)) {
121                return new IntegralConstantsSwitchCodegen(expression, isStatement, codegen);
122            }
123    
124            if (isStringConstantsSwitch(expression, subjectType, bindingContext)) {
125                return new StringSwitchCodegen(expression, isStatement, codegen);
126            }
127    
128            return null;
129        }
130    
131        private static boolean isThereConstantEntriesButNulls(
132                @NotNull JetWhenExpression expression,
133                @NotNull BindingContext bindingContext
134        ) {
135            for (CompileTimeConstant constant : getAllConstants(expression, bindingContext)) {
136                if (constant != null && !(constant instanceof NullValue)) return true;
137            }
138    
139            return false;
140        }
141    
142        private static boolean isIntegralConstantsSwitch(
143                @NotNull JetWhenExpression expression,
144                @NotNull Type subjectType,
145                @NotNull BindingContext bindingContext
146        ) {
147            int typeSort = subjectType.getSort();
148    
149            if (typeSort != Type.INT && typeSort != Type.CHAR && typeSort != Type.SHORT && typeSort != Type.BYTE) {
150                return false;
151            }
152    
153            return checkAllItemsAreConstantsSatisfying(expression, bindingContext, new Function1<CompileTimeConstant, Boolean>() {
154                @Override
155                public Boolean invoke(
156                        @NotNull CompileTimeConstant constant
157                ) {
158                    return constant instanceof IntegerValueConstant;
159                }
160            });
161        }
162    
163        private static boolean isStringConstantsSwitch(
164                @NotNull JetWhenExpression expression,
165                @NotNull Type subjectType,
166                @NotNull BindingContext bindingContext
167        ) {
168    
169            if (!subjectType.getClassName().equals(String.class.getName())) {
170                return false;
171            }
172    
173            return checkAllItemsAreConstantsSatisfying(expression, bindingContext, new Function1<CompileTimeConstant, Boolean>() {
174                @Override
175                public Boolean invoke(
176                        @NotNull CompileTimeConstant constant
177                ) {
178                    return constant instanceof StringValue || constant instanceof NullValue;
179                }
180            });
181        }
182    }