001    /*
002     * Copyright 2010-2013 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.jet.plugin;
018    
019    import com.intellij.util.NotNullFunction;
020    import org.jetbrains.annotations.NotNull;
021    import org.jetbrains.annotations.Nullable;
022    import org.jetbrains.jet.lang.descriptors.FunctionDescriptor;
023    import org.jetbrains.jet.lang.descriptors.SimpleFunctionDescriptor;
024    import org.jetbrains.jet.lang.descriptors.ValueParameterDescriptor;
025    import org.jetbrains.jet.lang.psi.JetDeclaration;
026    import org.jetbrains.jet.lang.psi.JetFile;
027    import org.jetbrains.jet.lang.psi.JetNamedFunction;
028    import org.jetbrains.jet.lang.resolve.BindingContext;
029    import org.jetbrains.jet.lang.resolve.lazy.ResolveSession;
030    import org.jetbrains.jet.lang.types.JetType;
031    import org.jetbrains.jet.lang.types.TypeProjection;
032    import org.jetbrains.jet.lang.types.checker.JetTypeChecker;
033    import org.jetbrains.jet.lang.types.lang.KotlinBuiltIns;
034    
035    import java.util.Collection;
036    import java.util.List;
037    
038    public class MainFunctionDetector {
039        private final NotNullFunction<JetNamedFunction, 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<JetNamedFunction, FunctionDescriptor>() {
044                @NotNull
045                @Override
046                public FunctionDescriptor fun(JetNamedFunction 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        /** Uses the {@code resolveSession} to resolve the function declaration. Suitable when the function declaration is not resolved yet. */
057        public MainFunctionDetector(@NotNull final ResolveSession resolveSession) {
058            this.getFunctionDescriptor = new NotNullFunction<JetNamedFunction, FunctionDescriptor>() {
059                @NotNull
060                @Override
061                public FunctionDescriptor fun(JetNamedFunction function) {
062                    return (FunctionDescriptor) resolveSession.resolveToDescriptor(function);
063                }
064            };
065        }
066    
067        public boolean hasMain(@NotNull List<JetDeclaration> declarations) {
068            return findMainFunction(declarations) != null;
069        }
070    
071        public boolean isMain(@NotNull JetNamedFunction function) {
072            if ("main".equals(function.getName())) {
073                FunctionDescriptor functionDescriptor = getFunctionDescriptor.fun(function);
074                List<ValueParameterDescriptor> parameters = functionDescriptor.getValueParameters();
075                if (parameters.size() == 1) {
076                    ValueParameterDescriptor parameter = parameters.get(0);
077                    JetType parameterType = parameter.getType();
078                    KotlinBuiltIns kotlinBuiltIns = KotlinBuiltIns.getInstance();
079                    if (kotlinBuiltIns.isArray(parameterType)) {
080                        List<TypeProjection> typeArguments = parameterType.getArguments();
081                        if (typeArguments.size() == 1) {
082                            JetType typeArgument = typeArguments.get(0).getType();
083                            if (JetTypeChecker.DEFAULT.equalTypes(typeArgument, kotlinBuiltIns.getStringType())) {
084                                return true;
085                            }
086                        }
087                    }
088                }
089            }
090            return false;
091        }
092    
093        @Nullable
094        public JetNamedFunction getMainFunction(@NotNull Collection<JetFile> files) {
095            for (JetFile file : files) {
096                JetNamedFunction mainFunction = findMainFunction(file.getDeclarations());
097                if (mainFunction != null) {
098                    return mainFunction;
099                }
100            }
101            return null;
102        }
103    
104        @Nullable
105        private JetNamedFunction findMainFunction(@NotNull List<JetDeclaration> declarations) {
106            for (JetDeclaration declaration : declarations) {
107                if (declaration instanceof JetNamedFunction) {
108                    JetNamedFunction candidateFunction = (JetNamedFunction) declaration;
109                    if (isMain(candidateFunction)) {
110                        return candidateFunction;
111                    }
112                }
113            }
114            return null;
115        }
116    }