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.KotlinBuiltIns;
023    import org.jetbrains.kotlin.builtins.PrimitiveType;
024    import org.jetbrains.kotlin.descriptors.*;
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.BindingContext;
030    import org.jetbrains.kotlin.resolve.DescriptorUtils;
031    import org.jetbrains.kotlin.resolve.calls.callUtil.CallUtilKt;
032    import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall;
033    import org.jetbrains.kotlin.types.KotlinType;
034    
035    import java.util.Arrays;
036    import java.util.List;
037    
038    import static org.jetbrains.kotlin.builtins.KotlinBuiltIns.RANGES_PACKAGE_FQ_NAME;
039    import static org.jetbrains.kotlin.codegen.AsmUtil.isPrimitiveNumberClassDescriptor;
040    
041    public class RangeCodegenUtil {
042        private static final ImmutableMap<FqName, PrimitiveType> RANGE_TO_ELEMENT_TYPE;
043        private static final ImmutableMap<FqName, PrimitiveType> PROGRESSION_TO_ELEMENT_TYPE;
044    
045        @NotNull
046        public static List<PrimitiveType> supportedRangeTypes() {
047            return Arrays.asList(PrimitiveType.CHAR, PrimitiveType.INT, PrimitiveType.LONG);
048        }
049    
050        static {
051            ImmutableMap.Builder<FqName, PrimitiveType> rangeBuilder = ImmutableMap.builder();
052            ImmutableMap.Builder<FqName, PrimitiveType> progressionBuilder = ImmutableMap.builder();
053            for (PrimitiveType primitiveType : supportedRangeTypes()) {
054                FqName rangeClassFqName = RANGES_PACKAGE_FQ_NAME.child(Name.identifier(primitiveType.getTypeName() + "Range"));
055                FqName progressionClassFqName = RANGES_PACKAGE_FQ_NAME.child(Name.identifier(primitiveType.getTypeName() + "Progression"));
056                rangeBuilder.put(rangeClassFqName, primitiveType);
057                progressionBuilder.put(progressionClassFqName, primitiveType);
058            }
059            RANGE_TO_ELEMENT_TYPE = rangeBuilder.build();
060            PROGRESSION_TO_ELEMENT_TYPE = progressionBuilder.build();
061        }
062    
063        private RangeCodegenUtil() {}
064    
065        public static boolean isRange(KotlinType rangeType) {
066            return !rangeType.isMarkedNullable() && getPrimitiveRangeElementType(rangeType) != null;
067        }
068    
069        public static boolean isProgression(KotlinType rangeType) {
070            return !rangeType.isMarkedNullable() && getPrimitiveProgressionElementType(rangeType) != null;
071        }
072    
073        @Nullable
074        public static BinaryCall getRangeAsBinaryCall(@NotNull KtForExpression forExpression) {
075            // We are looking for rangeTo() calls
076            // Other binary operations will succeed too, but will be filtered out later (by examining a resolvedCall)
077            KtExpression rangeExpression = forExpression.getLoopRange();
078            assert rangeExpression != null;
079            KtExpression loopRange = KtPsiUtil.deparenthesize(rangeExpression);
080            if (loopRange instanceof KtQualifiedExpression) {
081                // a.rangeTo(b)
082                KtQualifiedExpression qualifiedExpression = (KtQualifiedExpression) loopRange;
083                KtExpression selector = qualifiedExpression.getSelectorExpression();
084                if (selector instanceof KtCallExpression) {
085                    KtCallExpression callExpression = (KtCallExpression) selector;
086                    List<? extends ValueArgument> arguments = callExpression.getValueArguments();
087                    if (arguments.size() == 1) {
088                        return new BinaryCall(qualifiedExpression.getReceiverExpression(), callExpression.getCalleeExpression(),
089                                              arguments.get(0).getArgumentExpression());
090                    }
091                }
092            }
093            else if (loopRange instanceof KtBinaryExpression) {
094                // a rangeTo b
095                // a .. b
096                KtBinaryExpression binaryExpression = (KtBinaryExpression) loopRange;
097                return new BinaryCall(binaryExpression.getLeft(), binaryExpression.getOperationReference(), binaryExpression.getRight());
098    
099            }
100            return null;
101        }
102    
103        @Nullable
104        public static ResolvedCall<? extends CallableDescriptor> getLoopRangeResolvedCall(@NotNull KtForExpression forExpression, @NotNull BindingContext bindingContext) {
105            KtExpression loopRange = KtPsiUtil.deparenthesize(forExpression.getLoopRange());
106    
107            if (loopRange instanceof KtQualifiedExpression) {
108                KtQualifiedExpression qualifiedExpression = (KtQualifiedExpression) loopRange;
109                KtExpression selector = qualifiedExpression.getSelectorExpression();
110                if (selector instanceof KtCallExpression || selector instanceof KtSimpleNameExpression) {
111                    return CallUtilKt.getResolvedCall(selector, bindingContext);
112                }
113            }
114            else if (loopRange instanceof KtSimpleNameExpression) {
115                return CallUtilKt.getResolvedCall(loopRange, bindingContext);
116            }
117    
118            return null;
119        }
120    
121        @Nullable
122        private static PrimitiveType getPrimitiveRangeElementType(KotlinType rangeType) {
123            return getPrimitiveRangeOrProgressionElementType(rangeType, RANGE_TO_ELEMENT_TYPE);
124        }
125    
126        @Nullable
127        private static PrimitiveType getPrimitiveProgressionElementType(KotlinType rangeType) {
128            return getPrimitiveRangeOrProgressionElementType(rangeType, PROGRESSION_TO_ELEMENT_TYPE);
129        }
130    
131        @Nullable
132        private static PrimitiveType getPrimitiveRangeOrProgressionElementType(
133                @NotNull KotlinType rangeOrProgression,
134                @NotNull ImmutableMap<FqName, PrimitiveType> map
135        ) {
136            ClassifierDescriptor declarationDescriptor = rangeOrProgression.getConstructor().getDeclarationDescriptor();
137            if (declarationDescriptor == null) return null;
138            FqNameUnsafe fqName = DescriptorUtils.getFqName(declarationDescriptor);
139            if (!fqName.isSafe()) return null;
140            return map.get(fqName.toSafe());
141        }
142    
143        @Nullable
144        public static PrimitiveType getPrimitiveRangeOrProgressionElementType(@NotNull FqName rangeOrProgressionName) {
145            PrimitiveType result = RANGE_TO_ELEMENT_TYPE.get(rangeOrProgressionName);
146            return result != null ? result : PROGRESSION_TO_ELEMENT_TYPE.get(rangeOrProgressionName);
147        }
148    
149        public static boolean isRangeOrProgression(@NotNull FqName className) {
150            return getPrimitiveRangeOrProgressionElementType(className) != null;
151        }
152    
153        public static boolean isOptimizableRangeTo(CallableDescriptor rangeTo) {
154            if ("rangeTo".equals(rangeTo.getName().asString())) {
155                if (isPrimitiveNumberClassDescriptor(rangeTo.getContainingDeclaration())) {
156                    return true;
157                }
158            }
159            return false;
160        }
161    
162        public static boolean isArrayOrPrimitiveArrayIndices(@NotNull CallableDescriptor descriptor) {
163            if (!isTopLevelInPackage(descriptor, "indices", "kotlin.collections")) return false;
164    
165            ReceiverParameterDescriptor extensionReceiver = descriptor.getExtensionReceiverParameter();
166            if (extensionReceiver == null) return false;
167            KotlinType extensionReceiverType = extensionReceiver.getType();
168            if (!KotlinBuiltIns.isArray(extensionReceiverType) && !KotlinBuiltIns.isPrimitiveArray(extensionReceiverType)) return false;
169    
170            return true;
171        }
172    
173        public static boolean isCollectionIndices(@NotNull CallableDescriptor descriptor) {
174            if (!isTopLevelInPackage(descriptor, "indices", "kotlin.collections")) return false;
175    
176            ReceiverParameterDescriptor extensionReceiver = descriptor.getExtensionReceiverParameter();
177            if (extensionReceiver == null) return false;
178            KotlinType extensionReceiverType = extensionReceiver.getType();
179            if (!KotlinBuiltIns.isCollectionOrNullableCollection(extensionReceiverType)) return false;
180    
181            return true;
182        }
183    
184        private static boolean isTopLevelInPackage(@NotNull CallableDescriptor descriptor, @NotNull String name, @NotNull String packageName) {
185            if (!name.equals(descriptor.getName().asString())) return false;
186    
187            DeclarationDescriptor containingDeclaration = descriptor.getContainingDeclaration();
188            if (!(containingDeclaration instanceof PackageFragmentDescriptor)) return false;
189            String packageFqName = ((PackageFragmentDescriptor) containingDeclaration).getFqName().asString();
190            if (!packageName.equals(packageFqName)) return false;
191    
192            return true;
193        }
194    
195        public static class BinaryCall {
196            public final KtExpression left;
197            public final KtExpression op;
198            public final KtExpression right;
199    
200            private BinaryCall(KtExpression left, KtExpression op, KtExpression right) {
201                this.left = left;
202                this.op = op;
203                this.right = right;
204            }
205        }
206    }