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.JetDeclaration;
025    import org.jetbrains.kotlin.psi.JetFile;
026    import org.jetbrains.kotlin.psi.JetNamedFunction;
027    import org.jetbrains.kotlin.resolve.BindingContext;
028    import org.jetbrains.kotlin.resolve.DescriptorUtils;
029    import org.jetbrains.kotlin.resolve.annotations.AnnotationsPackage;
030    import org.jetbrains.kotlin.types.JetType;
031    import org.jetbrains.kotlin.types.TypeProjection;
032    import org.jetbrains.kotlin.types.checker.JetTypeChecker;
033    
034    import java.util.Collection;
035    import java.util.List;
036    
037    import static org.jetbrains.kotlin.resolve.descriptorUtil.DescriptorUtilPackage.getBuiltIns;
038    
039    public class MainFunctionDetector {
040        private final NotNullFunction<JetNamedFunction, FunctionDescriptor> getFunctionDescriptor;
041    
042        /** Assumes that the function declaration is already resolved and the descriptor can be found in the {@code bindingContext}. */
043        public MainFunctionDetector(@NotNull final BindingContext bindingContext) {
044            this.getFunctionDescriptor = new NotNullFunction<JetNamedFunction, FunctionDescriptor>() {
045                @NotNull
046                @Override
047                public FunctionDescriptor fun(JetNamedFunction function) {
048                    SimpleFunctionDescriptor functionDescriptor = bindingContext.get(BindingContext.FUNCTION, function);
049                    if (functionDescriptor == null) {
050                        throw new IllegalStateException("No descriptor resolved for " + function + " " + function.getText());
051                    }
052                    return functionDescriptor;
053                }
054            };
055        }
056    
057        public MainFunctionDetector(@NotNull NotNullFunction<JetNamedFunction, FunctionDescriptor> functionResolver) {
058            this.getFunctionDescriptor = functionResolver;
059        }
060    
061        public boolean hasMain(@NotNull List<JetDeclaration> declarations) {
062            return findMainFunction(declarations) != null;
063        }
064    
065        public boolean isMain(@NotNull JetNamedFunction function) {
066            if (!"main".equals(function.getName())) return false;
067    
068            FunctionDescriptor functionDescriptor = getFunctionDescriptor.fun(function);
069            List<ValueParameterDescriptor> parameters = functionDescriptor.getValueParameters();
070            if (parameters.size() != 1) return false;
071    
072            ValueParameterDescriptor parameter = parameters.get(0);
073            JetType parameterType = parameter.getType();
074            if (!KotlinBuiltIns.isArray(parameterType)) return false;
075    
076            List<TypeProjection> typeArguments = parameterType.getArguments();
077            if (typeArguments.size() != 1) return false;
078    
079            JetType typeArgument = typeArguments.get(0).getType();
080            if (!JetTypeChecker.DEFAULT.equalTypes(typeArgument, getBuiltIns(functionDescriptor).getStringType())) return false;
081    
082            if (DescriptorUtils.isTopLevelDeclaration(functionDescriptor)) return true;
083    
084            DeclarationDescriptor containingDeclaration = functionDescriptor.getContainingDeclaration();
085            return containingDeclaration instanceof ClassDescriptor
086                   && ((ClassDescriptor) containingDeclaration).getKind().isSingleton()
087                   && AnnotationsPackage.hasPlatformStaticAnnotation(functionDescriptor);
088        }
089    
090        @Nullable
091        public JetNamedFunction getMainFunction(@NotNull Collection<JetFile> files) {
092            for (JetFile file : files) {
093                JetNamedFunction mainFunction = findMainFunction(file.getDeclarations());
094                if (mainFunction != null) {
095                    return mainFunction;
096                }
097            }
098            return null;
099        }
100    
101        @Nullable
102        private JetNamedFunction findMainFunction(@NotNull List<JetDeclaration> declarations) {
103            for (JetDeclaration declaration : declarations) {
104                if (declaration instanceof JetNamedFunction) {
105                    JetNamedFunction candidateFunction = (JetNamedFunction) declaration;
106                    if (isMain(candidateFunction)) {
107                        return candidateFunction;
108                    }
109                }
110            }
111            return null;
112        }
113    }