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