IDEA-52789: Groovy map literals should have LinkedHashMap type
[idea/community.git] / plugins / groovy / src / org / jetbrains / plugins / groovy / lang / psi / impl / auxiliary / GrListOrMapImpl.java
1 /*
2  * Copyright 2000-2009 JetBrains s.r.o.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16
17 package org.jetbrains.plugins.groovy.lang.psi.impl.auxiliary;
18
19 import com.intellij.lang.ASTNode;
20 import com.intellij.psi.*;
21 import com.intellij.psi.search.GlobalSearchScope;
22 import com.intellij.psi.tree.TokenSet;
23 import com.intellij.util.Function;
24 import org.jetbrains.annotations.NotNull;
25 import org.jetbrains.annotations.Nullable;
26 import org.jetbrains.plugins.groovy.lang.lexer.GroovyTokenTypes;
27 import org.jetbrains.plugins.groovy.lang.parser.GroovyElementTypes;
28 import org.jetbrains.plugins.groovy.lang.psi.GroovyElementVisitor;
29 import org.jetbrains.plugins.groovy.lang.psi.api.auxiliary.GrListOrMap;
30 import org.jetbrains.plugins.groovy.lang.psi.api.statements.GrVariableDeclaration;
31 import org.jetbrains.plugins.groovy.lang.psi.api.statements.arguments.GrArgumentLabel;
32 import org.jetbrains.plugins.groovy.lang.psi.api.statements.arguments.GrNamedArgument;
33 import org.jetbrains.plugins.groovy.lang.psi.api.statements.expressions.GrExpression;
34 import org.jetbrains.plugins.groovy.lang.psi.api.statements.expressions.GrParenthesizedExpression;
35 import org.jetbrains.plugins.groovy.lang.psi.api.types.GrTypeElement;
36 import org.jetbrains.plugins.groovy.lang.psi.impl.GrTupleType;
37 import org.jetbrains.plugins.groovy.lang.psi.impl.GroovyPsiManager;
38 import org.jetbrains.plugins.groovy.lang.psi.impl.statements.expressions.GrExpressionImpl;
39 import org.jetbrains.plugins.groovy.lang.psi.impl.statements.expressions.TypesUtil;
40 import org.jetbrains.plugins.groovy.lang.psi.util.PsiUtil;
41
42 import static org.jetbrains.plugins.groovy.lang.lexer.GroovyTokenTypes.mCOMMA;
43
44 /**
45  * @author ilyas
46  */
47 public class GrListOrMapImpl extends GrExpressionImpl implements GrListOrMap {
48   private static final TokenSet MAP_LITERAL_TOKEN_SET = TokenSet.create(GroovyElementTypes.ARGUMENT, GroovyTokenTypes.mCOLON);
49   private static final Function<GrListOrMapImpl, PsiType> TYPES_CALCULATOR = new MyTypesCalculator();
50
51   public GrListOrMapImpl(@NotNull ASTNode node) {
52     super(node);
53   }
54
55   public void accept(GroovyElementVisitor visitor) {
56     visitor.visitListOrMap(this);
57   }
58
59   public String toString() {
60     return "Generalized list";
61   }
62
63   @Override
64   public ASTNode addInternal(ASTNode first, ASTNode last, ASTNode anchor, Boolean before) {
65     final GrExpression[] initializers = getInitializers();
66     if (initializers.length == 0) {
67       return super.addInternal(first, last, getNode().getFirstChildNode(), false);
68     }
69     final ASTNode lastChild = getNode().getLastChildNode();
70     getNode().addLeaf(mCOMMA, ",", lastChild);
71     return super.addInternal(first, last, lastChild.getTreePrev(), false);
72   }
73
74   public PsiType getType() {
75     return GroovyPsiManager.getInstance(getProject()).getType(this, TYPES_CALCULATOR);
76   }
77
78   public boolean isMap() {
79     return findChildByType(MAP_LITERAL_TOKEN_SET) != null;
80   }
81
82   @NotNull
83   public GrExpression[] getInitializers() {
84     return findChildrenByClass(GrExpression.class);
85   }
86
87   @NotNull
88   public GrNamedArgument[] getNamedArguments() {
89     return findChildrenByClass(GrNamedArgument.class);
90   }
91
92   private static class MyTypesCalculator implements Function<GrListOrMapImpl, PsiType> {
93     @Nullable
94     public PsiType fun(GrListOrMapImpl listOrMap) {
95       final GlobalSearchScope scope = listOrMap.getResolveScope();
96       if (listOrMap.isMap()) {
97         return inferMapInitializerType(listOrMap, JavaPsiFacade.getInstance(listOrMap.getProject()), scope);
98       }
99
100       PsiElement parent = listOrMap.getParent();
101       if (parent.getParent() instanceof GrVariableDeclaration) {
102         GrTypeElement typeElement = ((GrVariableDeclaration)parent.getParent()).getTypeElementGroovy();
103         if (typeElement != null) {
104           PsiType declaredType = typeElement.getType();
105           if (declaredType instanceof PsiArrayType) return declaredType;
106         }
107       }
108
109       return getTupleType(listOrMap.getInitializers(), listOrMap);
110     }
111
112     @Nullable
113     private static PsiClassType inferMapInitializerType(GrListOrMapImpl listOrMap, JavaPsiFacade facade, GlobalSearchScope scope) {
114       PsiClass mapClass = facade.findClass("java.util.LinkedHashMap", scope);
115       if (mapClass == null) {
116         mapClass = facade.findClass(CommonClassNames.JAVA_UTIL_MAP, scope);
117       }
118       PsiElementFactory factory = facade.getElementFactory();
119       if (mapClass != null) {
120         PsiTypeParameter[] typeParameters = mapClass.getTypeParameters();
121         if (typeParameters.length == 2) {
122           GrNamedArgument[] namedArgs = listOrMap.getNamedArguments();
123           GrExpression[] values = new GrExpression[namedArgs.length];
124           GrArgumentLabel[] labels = new GrArgumentLabel[namedArgs.length];
125
126           for (int i = 0; i < values.length; i++) {
127             GrExpression expr = namedArgs[i].getExpression();
128             if (expr == null) return null;
129             values[i] = expr;
130             GrArgumentLabel label = namedArgs[i].getLabel();
131             if (label == null) return null;
132             labels[i] = label;
133           }
134
135           PsiType initializerType = getInitializerType(values);
136           PsiType labelType = getLabelsType(labels);
137           PsiSubstitutor substitutor = PsiSubstitutor.EMPTY.
138             put(typeParameters[0], labelType).
139             put(typeParameters[1], initializerType);
140           return factory.createType(mapClass, substitutor);
141         }
142         else {
143           return facade.getElementFactory().createType(mapClass);
144         }
145       }
146       return null;
147     }
148
149     @Nullable
150     private static PsiType getLabelsType(GrArgumentLabel[] labels) {
151       if (labels.length == 0) return null;
152       PsiType result = null;
153       PsiManager manager = labels[0].getManager();
154       final PsiElementFactory factory = JavaPsiFacade.getElementFactory(labels[0].getProject());
155       for (GrArgumentLabel label : labels) {
156         PsiElement el = label.getNameElement();
157         PsiType other;
158         if (el instanceof GrParenthesizedExpression) {
159           other = ((GrParenthesizedExpression)el).getType();
160         }
161         else {
162           final ASTNode node = el.getNode();
163           if (node != null) {
164             other = TypesUtil.getPsiType(el, node.getElementType());
165             if (other == null) {
166               other = factory.createTypeByFQClassName(CommonClassNames.JAVA_LANG_STRING, el.getResolveScope());
167             }
168           }
169           else {
170             other = null;
171           }
172         }
173         result = getLeastUpperBound(result, other, manager);
174       }
175       return result;
176     }
177
178     private static PsiClassType getTupleType(GrExpression[] initializers, GrListOrMap listOrMap) {
179       PsiType[] result = new PsiType[initializers.length];
180       boolean isLValue = PsiUtil.isLValue(listOrMap);
181       for (int i = 0; i < result.length; i++) {
182         result[i] = isLValue ? initializers[i].getNominalType() : initializers[i].getType();
183       }
184       return new GrTupleType(result, JavaPsiFacade.getInstance(listOrMap.getProject()), listOrMap.getResolveScope());
185     }
186
187     @Nullable
188     private static PsiType getInitializerType(GrExpression[] initializers) {
189       if (initializers.length == 0) return null;
190       PsiManager manager = initializers[0].getManager();
191       PsiType result = initializers[0].getType();
192       for (int i = 1; i < initializers.length; i++) {
193         result = getLeastUpperBound(result, initializers[i].getType(), manager);
194       }
195
196       return result;
197     }
198
199     @Nullable
200     private static PsiType getLeastUpperBound(PsiType result, PsiType other, PsiManager manager) {
201       if (other == null) return result;
202       if (result == null) result = other;
203       if (result.isAssignableFrom(other)) return result;
204       if (other.isAssignableFrom(result)) result = other;
205
206       result = TypesUtil.getLeastUpperBound(result, other, manager);
207       return result;
208     }
209
210   }
211 }