c12e64fdf25b7913722397e87862a87faf0f26e7
[idea/community.git] / java / java-impl / src / com / intellij / codeInsight / daemon / impl / quickfix / AddVariableInitializerFix.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 package com.intellij.codeInsight.daemon.impl.quickfix;
17
18 import com.intellij.codeInsight.CodeInsightBundle;
19 import com.intellij.codeInsight.FileModificationService;
20 import com.intellij.codeInsight.completion.*;
21 import com.intellij.codeInsight.intention.IntentionAction;
22 import com.intellij.codeInsight.lookup.*;
23 import com.intellij.codeInsight.template.*;
24 import com.intellij.openapi.diagnostic.Logger;
25 import com.intellij.openapi.editor.Editor;
26 import com.intellij.openapi.project.Project;
27 import com.intellij.psi.*;
28 import com.intellij.psi.util.InheritanceUtil;
29 import com.intellij.psi.util.PsiTreeUtil;
30 import com.intellij.psi.util.PsiTypesUtil;
31 import com.intellij.psi.util.PsiUtil;
32 import com.intellij.refactoring.util.RefactoringUtil;
33 import com.intellij.util.Consumer;
34 import com.intellij.util.Function;
35 import com.intellij.util.IncorrectOperationException;
36 import com.intellij.util.SmartList;
37 import com.intellij.util.containers.ContainerUtil;
38 import org.jetbrains.annotations.NotNull;
39 import org.jetbrains.annotations.Nullable;
40
41 import java.util.Collections;
42 import java.util.Comparator;
43 import java.util.List;
44
45 import static com.intellij.util.containers.ContainerUtil.*;
46
47 public class AddVariableInitializerFix implements IntentionAction {
48   private static final Logger LOG = Logger.getInstance("#com.intellij.codeInsight.daemon.impl.quickfix.AddReturnFix");
49   private final PsiVariable myVariable;
50
51   public AddVariableInitializerFix(@NotNull PsiVariable variable) {
52     myVariable = variable;
53   }
54
55   @Override
56   @NotNull
57   public String getText() {
58     return CodeInsightBundle.message("quickfix.add.variable.text", myVariable.getName());
59   }
60
61   @Override
62   @NotNull
63   public String getFamilyName() {
64     return CodeInsightBundle.message("quickfix.add.variable.family.name");
65   }
66
67   @Override
68   public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) {
69     return myVariable.isValid() &&
70            myVariable.getManager().isInProject(myVariable) &&
71            !myVariable.hasInitializer() &&
72            !(myVariable instanceof PsiParameter)
73         ;
74   }
75
76   @Override
77   public void invoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException {
78     if (!FileModificationService.getInstance().prepareFileForWrite(myVariable.getContainingFile())) return;
79
80     final LookupElement[] suggestedInitializers = suggestInitializer(myVariable);
81     LOG.assertTrue(suggestedInitializers.length > 0);
82     LOG.assertTrue(suggestedInitializers[0] instanceof ExpressionLookupItem);
83     final PsiExpression initializer = (PsiExpression) suggestedInitializers[0].getObject();
84     if (myVariable instanceof PsiLocalVariable) {
85       ((PsiLocalVariable)myVariable).setInitializer(initializer);
86     }
87     else if (myVariable instanceof PsiField) {
88       ((PsiField)myVariable).setInitializer(initializer);
89     }
90     else {
91       LOG.error("Unknown variable type: "+myVariable);
92     }
93     runAssignmentTemplate(Collections.singletonList(myVariable.getInitializer()), suggestedInitializers, editor);
94   }
95
96   public static void runAssignmentTemplate(@NotNull final List<PsiExpression> initializers,
97                                            @NotNull final LookupElement[] suggestedInitializers,
98                                            @Nullable Editor editor) {
99     if (editor == null) return;
100     LOG.assertTrue(!initializers.isEmpty());
101     final PsiExpression initializer = ContainerUtil.getFirstItem(initializers);
102     PsiElement context = initializers.size() == 1 ? initializer : PsiTreeUtil.findCommonParent(initializers);
103     PsiDocumentManager.getInstance(initializer.getProject()).doPostponedOperationsAndUnblockDocument(editor.getDocument());
104     final TemplateBuilderImpl builder = (TemplateBuilderImpl)TemplateBuilderFactory.getInstance().createTemplateBuilder(context);
105     for (PsiExpression e : initializers) {
106       builder.replaceElement(e, new Expression() {
107         @Nullable
108         @Override
109         public Result calculateResult(ExpressionContext context1) {
110           return calculateQuickResult(context1);
111         }
112
113         @Nullable
114         @Override
115         public Result calculateQuickResult(ExpressionContext context1) {
116           return new PsiElementResult(suggestedInitializers[0].getPsiElement());
117         }
118
119         @Nullable
120         @Override
121         public LookupElement[] calculateLookupItems(ExpressionContext context1) {
122           return suggestedInitializers;
123         }
124       });
125     }
126     builder.run(editor, false);
127   }
128
129   @NotNull
130   public static LookupElement[] suggestInitializer(final PsiVariable variable) {
131     PsiType type = variable.getType();
132     final PsiElementFactory elementFactory = JavaPsiFacade.getElementFactory(variable.getProject());
133     if (type instanceof PsiClassType) {
134       final PsiClass aClass = PsiTypesUtil.getPsiClass(type);
135       if (aClass != null) {
136         final LookupElement nullLookupItem = new ExpressionLookupItem(elementFactory.createExpressionFromText(PsiKeyword.NULL, variable));
137         if (InheritanceUtil.isInheritor(aClass, CommonClassNames.JAVA_LANG_ITERABLE) ||
138             InheritanceUtil.isInheritor(aClass, CommonClassNames.JAVA_UTIL_MAP)) {
139           final List<PsiType> suggestedTypes = new SmartList<PsiType>();
140           JavaInheritorsGetter.processInheritors(variable.getContainingFile(), variable, Collections.singleton((PsiClassType) type), PrefixMatcher.ALWAYS_TRUE, new Consumer<PsiType>() {
141             @Override
142             public void consume(PsiType type) {
143               LOG.assertTrue(type instanceof PsiClassType);
144               final PsiClass psiClass = PsiTypesUtil.getPsiClass(type);
145               if (psiClass != null &&
146                   !psiClass.isInterface() &&
147                   !psiClass.hasModifierProperty(PsiModifier.ABSTRACT) &&
148                   PsiUtil.hasDefaultConstructor(psiClass)) {
149                 suggestedTypes.add(type);
150               }
151             }
152           });
153
154           List<LookupElement> sortedLookups = map(sorted(map(suggestedTypes, new Function<PsiType, LookupElement>() {
155             @Override
156             public LookupElement fun(PsiType type) {
157               return PsiTypeLookupItem.createLookupItem(type, variable);
158             }
159           }), new Comparator<LookupElement>() {
160             @Override
161             public int compare(LookupElement o1, LookupElement o2) {
162               final int count1 = StatisticsWeigher.getBaseStatisticsInfo(o1, null).getUseCount();
163               final int count2 = StatisticsWeigher.getBaseStatisticsInfo(o2, null).getUseCount();
164               return count2 - count1;
165             }
166           }), new Function<LookupElement, LookupElement>() {
167             @Override
168             public LookupElement fun(LookupElement element) {
169               final LookupElementDecorator<LookupElement> constructorLookupElement =
170                 LookupElementDecorator.withInsertHandler(element, ConstructorInsertHandler.BASIC_INSTANCE);
171               return new LookupElementDecorator<LookupElement>(constructorLookupElement) {
172                 @Override
173                 public void renderElement(LookupElementPresentation presentation) {
174                   super.renderElement(presentation);
175                   presentation.setTailText("");
176                   presentation.setItemText(PsiKeyword.NEW + " " + presentation.getItemText() + "()");
177                 }
178
179                 @Override
180                 public void handleInsert(InsertionContext context) {
181                   super.handleInsert(context);
182                   context.getDocument().insertString(context.getStartOffset(), PsiKeyword.NEW + " ");
183                 }
184               };
185             }
186           });
187           LookupElement[] result = new LookupElement[sortedLookups.size() + 1];
188           result[0] = nullLookupItem;
189           for (int i = 0; i < sortedLookups.size(); i++) {
190             LookupElement lookup = sortedLookups.get(i);
191             result[i + 1] = lookup;
192           }
193           return result;
194         } else {
195           if (PsiUtil.hasDefaultConstructor(aClass)) {
196             final PsiExpression newExpression = elementFactory
197               .createExpressionFromText(PsiKeyword.NEW + " " + type.getCanonicalText(false) + "()", variable);
198             return new LookupElement[]{nullLookupItem, new ExpressionLookupItem(newExpression)};
199           }
200         }
201       }
202     }
203     final String defaultValue = PsiTypesUtil.getDefaultValueOfType(type);
204     final PsiExpression expression = elementFactory.createExpressionFromText(defaultValue, variable);
205     return new LookupElement[] {new ExpressionLookupItem(expression)};
206   }
207
208   @Override
209   public boolean startInWriteAction() {
210     return true;
211   }
212 }