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;
018    
019    import com.google.common.collect.ImmutableMap;
020    import org.jetbrains.annotations.NotNull;
021    import org.jetbrains.annotations.Nullable;
022    import org.jetbrains.kotlin.builtins.PrimitiveType;
023    import org.jetbrains.kotlin.descriptors.CallableDescriptor;
024    import org.jetbrains.kotlin.descriptors.ClassifierDescriptor;
025    import org.jetbrains.kotlin.name.FqName;
026    import org.jetbrains.kotlin.name.FqNameUnsafe;
027    import org.jetbrains.kotlin.name.Name;
028    import org.jetbrains.kotlin.psi.*;
029    import org.jetbrains.kotlin.resolve.DescriptorUtils;
030    import org.jetbrains.kotlin.types.KotlinType;
031    
032    import java.util.Arrays;
033    import java.util.List;
034    
035    import static org.jetbrains.kotlin.builtins.KotlinBuiltIns.RANGES_PACKAGE_FQ_NAME;
036    import static org.jetbrains.kotlin.codegen.AsmUtil.isPrimitiveNumberClassDescriptor;
037    
038    public class RangeCodegenUtil {
039        private static final ImmutableMap<FqName, PrimitiveType> RANGE_TO_ELEMENT_TYPE;
040        private static final ImmutableMap<FqName, PrimitiveType> PROGRESSION_TO_ELEMENT_TYPE;
041    
042        @NotNull
043        public static List<PrimitiveType> supportedRangeTypes() {
044            return Arrays.asList(PrimitiveType.CHAR, PrimitiveType.INT, PrimitiveType.LONG);
045        }
046    
047        static {
048            ImmutableMap.Builder<FqName, PrimitiveType> rangeBuilder = ImmutableMap.builder();
049            ImmutableMap.Builder<FqName, PrimitiveType> progressionBuilder = ImmutableMap.builder();
050            for (PrimitiveType primitiveType : supportedRangeTypes()) {
051                FqName rangeClassFqName = RANGES_PACKAGE_FQ_NAME.child(Name.identifier(primitiveType.getTypeName() + "Range"));
052                FqName progressionClassFqName = RANGES_PACKAGE_FQ_NAME.child(Name.identifier(primitiveType.getTypeName() + "Progression"));
053                rangeBuilder.put(rangeClassFqName, primitiveType);
054                progressionBuilder.put(progressionClassFqName, primitiveType);
055            }
056            RANGE_TO_ELEMENT_TYPE = rangeBuilder.build();
057            PROGRESSION_TO_ELEMENT_TYPE = progressionBuilder.build();
058        }
059    
060        private RangeCodegenUtil() {}
061    
062        public static boolean isRange(KotlinType rangeType) {
063            return !rangeType.isMarkedNullable() && getPrimitiveRangeElementType(rangeType) != null;
064        }
065    
066        public static boolean isProgression(KotlinType rangeType) {
067            return !rangeType.isMarkedNullable() && getPrimitiveProgressionElementType(rangeType) != null;
068        }
069    
070        @Nullable
071        public static BinaryCall getRangeAsBinaryCall(@NotNull KtForExpression forExpression) {
072            // We are looking for rangeTo() calls
073            // Other binary operations will succeed too, but will be filtered out later (by examining a resolvedCall)
074            KtExpression rangeExpression = forExpression.getLoopRange();
075            assert rangeExpression != null;
076            KtExpression loopRange = KtPsiUtil.deparenthesize(rangeExpression);
077            if (loopRange instanceof KtQualifiedExpression) {
078                // a.rangeTo(b)
079                KtQualifiedExpression qualifiedExpression = (KtQualifiedExpression) loopRange;
080                KtExpression selector = qualifiedExpression.getSelectorExpression();
081                if (selector instanceof KtCallExpression) {
082                    KtCallExpression callExpression = (KtCallExpression) selector;
083                    List<? extends ValueArgument> arguments = callExpression.getValueArguments();
084                    if (arguments.size() == 1) {
085                        return new BinaryCall(qualifiedExpression.getReceiverExpression(), callExpression.getCalleeExpression(),
086                                              arguments.get(0).getArgumentExpression());
087                    }
088                }
089            }
090            else if (loopRange instanceof KtBinaryExpression) {
091                // a rangeTo b
092                // a .. b
093                KtBinaryExpression binaryExpression = (KtBinaryExpression) loopRange;
094                return new BinaryCall(binaryExpression.getLeft(), binaryExpression.getOperationReference(), binaryExpression.getRight());
095    
096            }
097            return null;
098        }
099    
100        @Nullable
101        private static PrimitiveType getPrimitiveRangeElementType(KotlinType rangeType) {
102            return getPrimitiveRangeOrProgressionElementType(rangeType, RANGE_TO_ELEMENT_TYPE);
103        }
104    
105        @Nullable
106        private static PrimitiveType getPrimitiveProgressionElementType(KotlinType rangeType) {
107            return getPrimitiveRangeOrProgressionElementType(rangeType, PROGRESSION_TO_ELEMENT_TYPE);
108        }
109    
110        @Nullable
111        private static PrimitiveType getPrimitiveRangeOrProgressionElementType(
112                @NotNull KotlinType rangeOrProgression,
113                @NotNull ImmutableMap<FqName, PrimitiveType> map
114        ) {
115            ClassifierDescriptor declarationDescriptor = rangeOrProgression.getConstructor().getDeclarationDescriptor();
116            if (declarationDescriptor == null) return null;
117            FqNameUnsafe fqName = DescriptorUtils.getFqName(declarationDescriptor);
118            if (!fqName.isSafe()) return null;
119            return map.get(fqName.toSafe());
120        }
121    
122        @Nullable
123        public static PrimitiveType getPrimitiveRangeOrProgressionElementType(@NotNull FqName rangeOrProgressionName) {
124            PrimitiveType result = RANGE_TO_ELEMENT_TYPE.get(rangeOrProgressionName);
125            return result != null ? result : PROGRESSION_TO_ELEMENT_TYPE.get(rangeOrProgressionName);
126        }
127    
128        public static boolean isRangeOrProgression(@NotNull FqName className) {
129            return getPrimitiveRangeOrProgressionElementType(className) != null;
130        }
131    
132        public static boolean isOptimizableRangeTo(CallableDescriptor rangeTo) {
133            if ("rangeTo".equals(rangeTo.getName().asString())) {
134                if (isPrimitiveNumberClassDescriptor(rangeTo.getContainingDeclaration())) {
135                    return true;
136                }
137            }
138            return false;
139        }
140    
141        public static class BinaryCall {
142            public final KtExpression left;
143            public final KtExpression op;
144            public final KtExpression right;
145    
146            private BinaryCall(KtExpression left, KtExpression op, KtExpression right) {
147                this.left = left;
148                this.op = op;
149                this.right = right;
150            }
151        }
152    }