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.idea;
018    
019    import com.intellij.util.NotNullFunction;
020    import org.jetbrains.annotations.NotNull;
021    import org.jetbrains.annotations.Nullable;
022    import org.jetbrains.kotlin.builtins.KotlinBuiltIns;
023    import org.jetbrains.kotlin.descriptors.*;
024    import org.jetbrains.kotlin.psi.KtAnnotationEntry;
025    import org.jetbrains.kotlin.psi.KtDeclaration;
026    import org.jetbrains.kotlin.psi.KtFile;
027    import org.jetbrains.kotlin.psi.KtNamedFunction;
028    import org.jetbrains.kotlin.resolve.BindingContext;
029    import org.jetbrains.kotlin.resolve.DescriptorUtils;
030    import org.jetbrains.kotlin.resolve.annotations.AnnotationUtilKt;
031    import org.jetbrains.kotlin.types.KotlinType;
032    import org.jetbrains.kotlin.types.TypeProjection;
033    import org.jetbrains.kotlin.types.Variance;
034    
035    import java.util.Collection;
036    import java.util.List;
037    
038    public class MainFunctionDetector {
039        private final NotNullFunction<KtNamedFunction, FunctionDescriptor> getFunctionDescriptor;
040    
041        /** Assumes that the function declaration is already resolved and the descriptor can be found in the {@code bindingContext}. */
042        public MainFunctionDetector(@NotNull final BindingContext bindingContext) {
043            this.getFunctionDescriptor = new NotNullFunction<KtNamedFunction, FunctionDescriptor>() {
044                @NotNull
045                @Override
046                public FunctionDescriptor fun(KtNamedFunction function) {
047                    SimpleFunctionDescriptor functionDescriptor = bindingContext.get(BindingContext.FUNCTION, function);
048                    if (functionDescriptor == null) {
049                        throw new IllegalStateException("No descriptor resolved for " + function + " " + function.getText());
050                    }
051                    return functionDescriptor;
052                }
053            };
054        }
055    
056        public MainFunctionDetector(@NotNull NotNullFunction<KtNamedFunction, FunctionDescriptor> functionResolver) {
057            this.getFunctionDescriptor = functionResolver;
058        }
059    
060        public boolean hasMain(@NotNull List<KtDeclaration> declarations) {
061            return findMainFunction(declarations) != null;
062        }
063    
064        public boolean isMain(@NotNull KtNamedFunction function) {
065            if (function.isLocal()) {
066                return false;
067            }
068    
069            if (function.getValueParameters().size() != 1 || !function.getTypeParameters().isEmpty()) {
070                return false;
071            }
072    
073            /* Psi only check for kotlin.jvm.jvmName annotation */
074            if (!"main".equals(function.getName()) && !hasAnnotationWithExactNumberOfArguments(function, 1)) {
075                return false;
076            }
077    
078            /* Psi only check for kotlin.jvm.jvmStatic annotation */
079            if (!function.isTopLevel() && !hasAnnotationWithExactNumberOfArguments(function, 0)) {
080                return false;
081            }
082    
083            return isMain(getFunctionDescriptor.fun(function));
084        }
085    
086        public static boolean isMain(@NotNull DeclarationDescriptor descriptor) {
087            if (!(descriptor instanceof FunctionDescriptor)) return false;
088    
089            FunctionDescriptor functionDescriptor = (FunctionDescriptor) descriptor;
090            if (!getJVMFunctionName(functionDescriptor).equals("main")) {
091                return false;
092            }
093    
094            List<ValueParameterDescriptor> parameters = functionDescriptor.getValueParameters();
095            if (parameters.size() != 1 || !functionDescriptor.getTypeParameters().isEmpty()) return false;
096    
097            ValueParameterDescriptor parameter = parameters.get(0);
098            KotlinType parameterType = parameter.getType();
099            if (!KotlinBuiltIns.isArray(parameterType)) return false;
100    
101            List<TypeProjection> typeArguments = parameterType.getArguments();
102            if (typeArguments.size() != 1) return false;
103    
104            KotlinType typeArgument = typeArguments.get(0).getType();
105            if (!KotlinBuiltIns.isString(typeArgument)) {
106                return false;
107            }
108            if (typeArguments.get(0).getProjectionKind() == Variance.IN_VARIANCE) {
109                return false;
110            }
111    
112            if (DescriptorUtils.isTopLevelDeclaration(functionDescriptor)) return true;
113    
114            DeclarationDescriptor containingDeclaration = functionDescriptor.getContainingDeclaration();
115            return containingDeclaration instanceof ClassDescriptor
116                   && ((ClassDescriptor) containingDeclaration).getKind().isSingleton()
117                   && AnnotationUtilKt.hasJvmStaticAnnotation(functionDescriptor);
118        }
119    
120        @Nullable
121        public KtNamedFunction getMainFunction(@NotNull Collection<KtFile> files) {
122            for (KtFile file : files) {
123                KtNamedFunction mainFunction = findMainFunction(file.getDeclarations());
124                if (mainFunction != null) {
125                    return mainFunction;
126                }
127            }
128            return null;
129        }
130    
131        @Nullable
132        private KtNamedFunction findMainFunction(@NotNull List<KtDeclaration> declarations) {
133            for (KtDeclaration declaration : declarations) {
134                if (declaration instanceof KtNamedFunction) {
135                    KtNamedFunction candidateFunction = (KtNamedFunction) declaration;
136                    if (isMain(candidateFunction)) {
137                        return candidateFunction;
138                    }
139                }
140            }
141            return null;
142        }
143    
144        @NotNull
145        private static String getJVMFunctionName(FunctionDescriptor functionDescriptor) {
146            String platformName = DescriptorUtils.getJvmName(functionDescriptor);
147            if (platformName != null) {
148                return platformName;
149            }
150    
151            return functionDescriptor.getName().asString();
152        }
153    
154        private static boolean hasAnnotationWithExactNumberOfArguments(@NotNull KtNamedFunction function, int number) {
155            for (KtAnnotationEntry entry : function.getAnnotationEntries()) {
156                if (entry.getValueArguments().size() == number) {
157                    return true;
158                }
159            }
160    
161            return false;
162        }
163    }