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.ConstantValue;
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 KtWhenExpression expression,
038                @NotNull BindingContext bindingContext,
039                boolean shouldInlineConstVals,
040                Function1<ConstantValue<?>, Boolean> predicate
041        ) {
042            for (KtWhenEntry entry : expression.getEntries()) {
043                for (KtWhenCondition condition : entry.getConditions()) {
044                    if (!(condition instanceof KtWhenConditionWithExpression)) {
045                        return false;
046                    }
047    
048                    // ensure that expression is constant
049                    KtExpression patternExpression = ((KtWhenConditionWithExpression) condition).getExpression();
050    
051                    if (patternExpression == null) return false;
052    
053                    ConstantValue<?> constant = ExpressionCodegen.getCompileTimeConstant(patternExpression, bindingContext, shouldInlineConstVals);
054                    if (constant == null || !predicate.invoke(constant)) {
055                        return false;
056                    }
057                }
058            }
059    
060            return true;
061        }
062    
063        @NotNull
064        public static Iterable<ConstantValue<?>> getAllConstants(
065                @NotNull KtWhenExpression expression,
066                @NotNull BindingContext bindingContext,
067                boolean shouldInlineConstVals
068        ) {
069            List<ConstantValue<?>> result = new ArrayList<ConstantValue<?>>();
070    
071            for (KtWhenEntry entry : expression.getEntries()) {
072                addConstantsFromEntry(result, entry, bindingContext, shouldInlineConstVals);
073            }
074    
075            return result;
076        }
077    
078        private static void addConstantsFromEntry(
079                @NotNull List<ConstantValue<?>> result,
080                @NotNull KtWhenEntry entry,
081                @NotNull BindingContext bindingContext,
082                boolean shouldInlineConstVals
083        ) {
084            for (KtWhenCondition condition : entry.getConditions()) {
085                if (!(condition instanceof KtWhenConditionWithExpression)) continue;
086    
087                KtExpression patternExpression = ((KtWhenConditionWithExpression) condition).getExpression();
088    
089                assert patternExpression != null : "expression in when should not be null";
090                result.add(ExpressionCodegen.getCompileTimeConstant(patternExpression, bindingContext, shouldInlineConstVals));
091            }
092        }
093    
094        @NotNull
095        public static Iterable<ConstantValue<?>> getConstantsFromEntry(
096                @NotNull KtWhenEntry entry,
097                @NotNull BindingContext bindingContext,
098                boolean shouldInlineConstVals
099        ) {
100            List<ConstantValue<?>> result = new ArrayList<ConstantValue<?>>();
101            addConstantsFromEntry(result, entry, bindingContext, shouldInlineConstVals);
102            return result;
103        }
104    
105        @Nullable
106        public static SwitchCodegen buildAppropriateSwitchCodegenIfPossible(
107                @NotNull KtWhenExpression expression,
108                boolean isStatement,
109                boolean isExhaustive,
110                @NotNull ExpressionCodegen codegen
111        ) {
112            BindingContext bindingContext = codegen.getBindingContext();
113            boolean shouldInlineConstVals = codegen.getState().getShouldInlineConstVals();
114            if (!isThereConstantEntriesButNulls(expression, bindingContext, shouldInlineConstVals)) {
115                return null;
116            }
117    
118            Type subjectType = codegen.expressionType(expression.getSubjectExpression());
119    
120            WhenByEnumsMapping mapping = codegen.getBindingContext().get(CodegenBinding.MAPPING_FOR_WHEN_BY_ENUM, expression);
121    
122            if (mapping != null) {
123                return new EnumSwitchCodegen(expression, isStatement, isExhaustive, codegen, mapping);
124            }
125    
126            if (isIntegralConstantsSwitch(expression, subjectType, bindingContext, shouldInlineConstVals)) {
127                return new IntegralConstantsSwitchCodegen(expression, isStatement, isExhaustive, codegen);
128            }
129    
130            if (isStringConstantsSwitch(expression, subjectType, bindingContext, shouldInlineConstVals)) {
131                return new StringSwitchCodegen(expression, isStatement, isExhaustive, codegen);
132            }
133    
134            return null;
135        }
136    
137        private static boolean isThereConstantEntriesButNulls(
138                @NotNull KtWhenExpression expression,
139                @NotNull BindingContext bindingContext,
140                boolean shouldInlineConstVals
141        ) {
142            for (ConstantValue<?> constant : getAllConstants(expression, bindingContext, shouldInlineConstVals)) {
143                if (constant != null && !(constant instanceof NullValue)) return true;
144            }
145    
146            return false;
147        }
148    
149        private static boolean isIntegralConstantsSwitch(
150                @NotNull KtWhenExpression expression,
151                @NotNull Type subjectType,
152                @NotNull BindingContext bindingContext,
153                boolean shouldInlineConstVals
154        ) {
155            int typeSort = subjectType.getSort();
156    
157            if (typeSort != Type.INT && typeSort != Type.CHAR && typeSort != Type.SHORT && typeSort != Type.BYTE) {
158                return false;
159            }
160    
161            return checkAllItemsAreConstantsSatisfying(expression, bindingContext, shouldInlineConstVals, new Function1<ConstantValue<?>, Boolean>() {
162                @Override
163                public Boolean invoke(
164                        @NotNull ConstantValue<?> constant
165                ) {
166                    return constant instanceof IntegerValueConstant;
167                }
168            });
169        }
170    
171        private static boolean isStringConstantsSwitch(
172                @NotNull KtWhenExpression expression,
173                @NotNull Type subjectType,
174                @NotNull BindingContext bindingContext,
175                boolean shouldInlineConstVals
176        ) {
177    
178            if (!subjectType.getClassName().equals(String.class.getName())) {
179                return false;
180            }
181    
182            return checkAllItemsAreConstantsSatisfying(expression, bindingContext, shouldInlineConstVals, new Function1<ConstantValue<?>, Boolean>() {
183                @Override
184                public Boolean invoke(
185                        @NotNull ConstantValue<?> constant
186                ) {
187                    return constant instanceof StringValue || constant instanceof NullValue;
188                }
189            });
190        }
191    }