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.lang.resolve.lazy.declarations;
018    
019    import com.google.common.collect.ArrayListMultimap;
020    import com.google.common.collect.HashMultimap;
021    import com.google.common.collect.Lists;
022    import com.google.common.collect.Multimap;
023    import kotlin.Function0;
024    import org.jetbrains.annotations.NotNull;
025    import org.jetbrains.jet.lang.psi.*;
026    import org.jetbrains.jet.lang.resolve.lazy.data.JetClassInfoUtil;
027    import org.jetbrains.jet.lang.resolve.lazy.data.JetClassLikeInfo;
028    import org.jetbrains.jet.lang.resolve.lazy.data.JetScriptInfo;
029    import org.jetbrains.jet.lang.resolve.name.Name;
030    import org.jetbrains.jet.storage.NotNullLazyValue;
031    import org.jetbrains.jet.storage.StorageManager;
032    
033    import java.util.Collection;
034    import java.util.List;
035    
036    import static org.jetbrains.jet.lang.resolve.lazy.ResolveSessionUtils.safeNameForLazyResolve;
037    
038    public abstract class AbstractPsiBasedDeclarationProvider implements DeclarationProvider {
039    
040        protected static class Index {
041            // This mutable state is only modified under inside the computable
042            private final List<JetDeclaration> allDeclarations = Lists.newArrayList();
043            private final Multimap<Name, JetNamedFunction> functions = HashMultimap.create();
044            private final Multimap<Name, JetProperty> properties = HashMultimap.create();
045            private final Multimap<Name, JetClassLikeInfo> classesAndObjects = ArrayListMultimap.create(); // order matters here
046    
047            public void putToIndex(@NotNull JetDeclaration declaration) {
048                if (declaration instanceof JetClassInitializer) {
049                    return;
050                }
051                allDeclarations.add(declaration);
052                if (declaration instanceof JetNamedFunction) {
053                    JetNamedFunction namedFunction = (JetNamedFunction) declaration;
054                    functions.put(safeNameForLazyResolve(namedFunction), namedFunction);
055                }
056                else if (declaration instanceof JetProperty) {
057                    JetProperty property = (JetProperty) declaration;
058                    properties.put(safeNameForLazyResolve(property), property);
059                }
060                else if (declaration instanceof JetClassOrObject) {
061                    JetClassOrObject classOrObject = (JetClassOrObject) declaration;
062                    classesAndObjects.put(
063                            safeNameForLazyResolve(classOrObject.getNameAsName()),
064                            JetClassInfoUtil.createClassLikeInfo(classOrObject)
065                    );
066                }
067                else if (declaration instanceof JetScript) {
068                    JetScript script = (JetScript) declaration;
069                    JetScriptInfo scriptInfo = new JetScriptInfo(script);
070                    classesAndObjects.put(
071                            scriptInfo.getFqName().shortName(),
072                            scriptInfo
073                    );
074                }
075                else if (declaration instanceof JetParameter ||
076                         declaration instanceof JetTypedef ||
077                         declaration instanceof JetMultiDeclaration) {
078                    // Do nothing, just put it into allDeclarations is enough
079                }
080                else {
081                    throw new IllegalArgumentException("Unknown declaration: " + declaration);
082                }
083            }
084        }
085    
086        private final NotNullLazyValue<Index> index;
087    
088        public AbstractPsiBasedDeclarationProvider(@NotNull StorageManager storageManager) {
089            index = storageManager.createLazyValue(new Function0<Index>() {
090                @Override
091                public Index invoke() {
092                    Index index = new Index();
093                    doCreateIndex(index);
094                    return index;
095                }
096            });
097        }
098    
099        protected abstract void doCreateIndex(@NotNull Index index);
100    
101        @NotNull
102        @Override
103        public List<JetDeclaration> getAllDeclarations() {
104            return index.invoke().allDeclarations;
105        }
106    
107        @NotNull
108        @Override
109        public List<JetNamedFunction> getFunctionDeclarations(@NotNull Name name) {
110            return Lists.newArrayList(index.invoke().functions.get(name));
111        }
112    
113        @NotNull
114        @Override
115        public List<JetProperty> getPropertyDeclarations(@NotNull Name name) {
116            return Lists.newArrayList(index.invoke().properties.get(name));
117        }
118    
119        @NotNull
120        @Override
121        public Collection<JetClassLikeInfo> getClassOrObjectDeclarations(@NotNull Name name) {
122            return index.invoke().classesAndObjects.get(name);
123        }
124    }